uPj!UdZddlmZddlZddlZddlmZddlmZddl Z dZ de d<d Z de d <d Z de d <d Zde d<dZde d<ej dZde d<ej dZde d<ee j)fZde d<d#dZd$dZd%dZd&dZ d'dZd(dZ d)dZ d* d+dZd,d Z d-d!Z d.d"Z y)/at Resolve AWS Bedrock ``requestMetadata`` from LiteLLM proxy identity and caller metadata. Bedrock attaches request metadata to CloudTrail records and to the dimension AWS Cost Explorer groups on, so everything here is opt-in: nothing is forwarded unless the operator sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the proxy). Two properties are load-bearing for that billing record and are asserted by the tests: proxy identity is resolved first so it can never be evicted by caller-supplied pairs, and the whole ``user_api_key_`` prefix is reserved so a caller cannot write a proxy-authoritative looking key. Values that break Bedrock's constraints are dropped rather than sanitised or rejected, because an operator flipping this setting on must not turn a working request into a 400 and a silently rewritten attribution key is worse than an absent one. ) annotationsN)Mapping)FinalzX-Amzn-Bedrock-Request-MetadatarBEDROCK_REQUEST_METADATA_HEADER"BEDROCK_REQUEST_METADATA_MAX_PAIRS user_api_key_(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIXspend_logs_metadata%BEDROCK_REQUEST_METADATA_CLIENT_FIELD)metadatalitellm_metadatazFinal[tuple[str, ...]]_METADATA_PARAM_NAMESz!^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$ _KEY_PATTERNz!^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$_VALUE_PATTERNzFinal[frozenset[str]]_OWNED_HEADER_NAMEScbtj|duxrtj|duSN)rmatchr)keyvalues j/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/request_metadata.py_is_forwardabler$s.   c "$ . Z>3G3G3NVZ3ZZcdt|tsytd|jDS)Nc3nK|]-\}}t|tst|ts(||f/ywr) isinstancestr.0rrs r z_text_pairs..+s.r*#u*SRUBV[efkmp[q#urs 55 5)rrtupleitems)sources r _text_pairsr&(s& fg & r r rrrctj}t|ttfsyt d|Dt fdt DS)z The operator allow-list, deduplicated so a field repeated in config cannot consume a second reserved slot and shrink the client budget for nothing. First occurrence wins, which keeps the operator's declared precedence intact. rc32K|]}t|ywr)r)r!fields rr"z"_allowed_fields..7s=#e*=sc38K|]\}}|d|vs|ywrr)r!indexr)fieldss rr"z"_allowed_fields..8s$\<5%VTZUZ^@[\s)litellmbedrock_request_metadata_fieldsrlistr# enumerate) configuredr,s @r_allowed_fieldsr2.sG !( G GJ j4- 0=*==F \9V+<\ \\rc:ytfdtDS)zZ``metadata`` on /v1/chat/completions, ``litellm_metadata`` on the LITELLM_METADATA_ROUTES.rc3dK|]'}j|}t|tr|)ywr)getrr)r!namer%litellm_paramss rr"z$_metadata_sources..?s3 %))$/F fg & s-0)r#r)r7s`r_metadata_sourcesr8;s& ) rc:tfd|DdtS)Nc3K|]7}|jtr t|}|t||r||f9ywr) startswithr _first_textr)r!r)rsourcess rr"z"_identity_pairs..KsG    D E!'51E  !> s=A)r#r)r=allowed_fieldss` r_identity_pairsr?Gs, #  *)  ++rc2tfd|DdS)Nc3fK|](}t|jxts%*ywr)rr5r)r!r%r)rs rr"z_first_text..Us)\6 FJJuDU;U5WZ0[\s&11)next)r=r)rs `@rr<r<Ts \W\^b ccrct|vrtd|Dnd}tdgt||DtfdtDd|S)Nc3fK|])}t|jtD]}|+ywr)r&r5r )r!r%pairs rr"z _client_pairs.._s+ tv[LqAr5s tTd td ts/1rc3lK|],\}}|jtst||r||f.ywr)r;r rr s rr"z _client_pairs..cs8 C~~FGO\_afLg e s24c3\K|]#\}}|dtdd|Dvr|%yw)rc3&K|] \}}| ywrr)r!earlier_s rr"z*_client_pairs...ksMJGQMsN)r#)r!r+rE candidatess rr"z _client_pairs..hs; E4 7%M*Ve:LMM M s),)r r#r&r0)r=r>caller_metadatabudgetspend_logs_pairsrKs @r _client_pairsrOXs{ 1N B  t tt  LK8L;KLJ $Z0 v  rc t}|syt|}t||}t|||tt |z }g||Dcic]\}}|| }}}|xsdScc}}w)a% Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when the feature is off or nothing survives Bedrock's constraints. The result is a plain dict because it is written straight onto the Converse body, which Bedrock types as ``dict[str, str]``. ``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. It has already been validated (and rejected with a 400) by the Converse transformation, so it is only filtered here for the reserved identity prefix and the remaining slot budget. N)r=r>rLrM)r2r8r?rOrlen) r7rLr>r=identityclientrrresolveds r resolve_bedrock_request_metadatarUos,-N &~6G%g~>H!%'1CMA F 5Ih4H4HIjc5sEzIHI  tJs A$c(ttS)a Whether the proxy OWNS the request-metadata field and header name for this request. Ownership follows the operator's opt-in alone, never whether anything resolved, because a caller can suppress the resolver by omitting the allow-listed fields or by sending values that all fail Bedrock's rules. Owned-but-empty has to mean "absent on the wire" rather than "fall back to whatever the caller supplied", or the reserved-prefix guarantee is bypassable by anyone who can make the resolver produce nothing. )boolr2rrr!bedrock_request_metadata_is_ownedrXs ! ""rcts tdfSt|}|tdfSttt j |dfffS)a The signed ``X-Amzn-Bedrock-Request-Metadata`` header for the Invoke paths, which have no body field for request metadata. Returns the header names the proxy OWNS and, separately, the pairs to send. Ownership is reported whenever forwarding is enabled, including when nothing resolves, because a caller can suppress the resolver (omit the allow-listed fields, or send values that all fail Bedrock's rules) and an owned-but-empty result must still evict the caller's header rather than fall back to it. r),:) separators)rX frozensetrUrrjsondumps)r7rTs r bedrock_request_metadata_headersr`sQ - .{B6~FH"B&& "A4::hcmCn!o q qqrc|s|ss|Std|Dgfd|jDfd|D|Dcic]\}}|| c}}Scc}}w)ab Merge the ``X-Amzn-*`` headers the Invoke paths derive from params. ``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry proxy-authenticated identity into an AWS billing record that the caller must not be able to write. Names are compared case-insensitively so a caller cannot leave a second spelling in the dict and let the transport pick the winner. c3<K|]}|jywrlower)r!r6s rr"z/merge_bedrock_invoke_headers..s%Gtdjjl%Gsc3RK|]\}}|jvs||f ywrrc)r!nvproxy_owned_namess rr"z/merge_bedrock_invoke_headers..s' WAAGGIEV4Vq!f W' 'c3RK|]\}}|jvs||f ywrrc)r!rfrgexisting_namess rr"z/merge_bedrock_invoke_headers..s& QA.1Pq!f Qri)r]r$)headers caller_owned proxy_ownedrhr6rrks ` @rmerge_bedrock_invoke_headersrosv  4E%%Gw%GGN W W Q Q    D% e   s A)rrrrreturnrW)r%objectrptuple[tuple[str, str], ...])rptuple[str, ...])r7Mapping[str, object] | Nonerp tuple[Mapping[str, object], ...])r=rur>rsrprr)r=rur)rrpz str | None) r=rur>rsrLrqrMintrprrr)r7rtrLrqrpzdict[str, str] | None)rprW)r7rtrpz2tuple[frozenset[str], tuple[tuple[str, str], ...]]) rldict[str, str]rmrrrnrrrhzfrozenset[str]rprw)!__doc__ __future__rr^recollections.abcrtypingrr-r__annotations__rr r rcompilerrr]rdrrr&r2r8r?r<rOrUrXr`rorrrrs # #)JJ,."E.2A(%A/D%uD0P-P bjj!EF eF" #GHH-68W8]8]8_7a-b*b[s ]  + - +# +! +d -#   ! 2#/8 #r/r7r* --&   r