o
    ưi[)                    @   s  d dl Z d dlZd dlZd dlmZmZmZmZmZm	Z	 d dl
mZ d dlmZ d dlZd dlmZmZ d dlmZ d dlmZ d dlmZmZmZmZmZmZmZ d d	lmZ e d
d ej!" D Z#d dl$m%Z% d dl&m'Z' d dl(m)Z) d dl*m+Z+m,Z,m-Z-m.Z. e Z/erd dl0m1Z2 d dl3m4Z4 e2Z1neZ1eZ4dd Z5dZ6dede7fddZ8deee7e7f  dee7 fddZ9de:defddZ;dedee defd d!Z<G d"d# d#Z=d$ed%e1dee fd&d'Z>		(	d`ded)ee7 d*e?d+ee7 de:f
d,d-Z@G d.d/ d/ZA		dade:ded$ed%e1d0eee7ef  d1ee7 fd2d3ZBde:d$eddfd4d5ZCde:d$eddfd6d7ZDd0ee: d$edeeE fd8d9ZFd:ede?fd;d<ZGd=e:d0ee: d$ed>e?de?f
d?d@ZHdAee: dBee: de:dCe7ddf
dDdEZIdAee: dBee: de:dCe7ddf
dFdGZJde:dHe7d$efdIdJZKdKe7de?fdLdMZLdKe7dee7 fdNdOZM	dbde:dPdQdRedSeee7ef  deNeEe7 e:e7e7f f f
dTdUZO		dade:dCe7dPdQdVeee7  dWeee7ef  ddfdXdYZPde:dCe7d$eddfdZd[ZQde:de:fd\d]ZRde:defd^d_ZSdS )c    N)TYPE_CHECKINGAnyDictListOptionalUnion)Request)Headers)verbose_loggerverbose_proxy_logger)ServiceLoggingsafe_json_loads)AddTeamCallbackCommonProxyErrorsLitellmDataForBackendLLMCallLitellmUserRolesSpecialHeadersTeamCallbackMetadataUserAPIKeyAuth)_safe_get_request_headersc                 c   s    | ]}|j  V  qd S N)valuelower).0v r   [/home/app/Keep/.python/lib/python3.10/site-packages/litellm/proxy/litellm_pre_call_utils.py	<genexpr>   s    

r   )Router)ANTHROPIC_API_HEADERS)ServiceTypes)LlmProvidersProviderSpecificHeader!StandardLoggingUserAPIKeyMetadataSupportedCacheControls)ProxyConfigPolicyMatchContextc                 C   sD   i }|  d}|D ]}d|v r| d\}}|||< q	d||< q	|S )Nz, =T)split)Zcache_control
cache_dictZ
directives	directivekeyr   r   r   r   parse_cache_control3   s   


r.   )Zbatchesz/v1/messages	responsesfilesrequestreturnc                    s:   | j j d v sd v rdS t fddtD rdS dS )z
    Helper to return what the "metadata" field should be called in the request data

    For all /thread or /assistant endpoints we need to call this "litellm_metadata"

    For ALL other endpoints we call this "metadata"
    threadZ	assistantlitellm_metadatac                 3   s    | ]}| v V  qd S r   r   )r   Zroutepathr   r   r   V   s    z._get_metadata_variable_name.<locals>.<genexpr>metadata)urlr6   anyLITELLM_METADATA_ROUTESr1   r   r5   r   _get_metadata_variable_nameI   s   r<   headersc                 C   s.   | sdS dd |   D }|dp|dS )a  
    Extract chain id for call chaining from request headers.

    x-litellm-trace-id and x-litellm-session-id are interchangeable; when both
    are present, x-litellm-trace-id takes precedence. Header keys are matched
    case-insensitively so this works with raw header dicts from any transport.

    Used by MCP (and other paths that have raw_headers but no Request) to set
    litellm_trace_id/litellm_session_id for spend logs and logging consistency.
    Nc                 S   s$   i | ]\}}t |tr| |qS r   )
isinstancestrr   )r   kr   r   r   r   
<dictcomp>i   s   $ z-get_chain_id_from_headers.<locals>.<dictcomp>x-litellm-trace-idx-litellm-session-id)itemsget)r=   
normalizedr   r   r   get_chain_id_from_headers\   s   rG   datac              
   C   s   zt |drt|j}d|v r|d | d< W d S W d S W d S  ty'   Y d S  tyA } ztdt| W Y d }~d S d }~ww )Nquery_paramsapi-versionapi_versionz.error checking api version in query params: %s)hasattrdictrI   KeyError	Exceptionr
   	exceptionr?   )rH   r1   rI   er   r   r   &safe_add_api_version_from_query_paramso   s   

rR   team_callback_settings_objc                 C   sN  |d u rt  }| jdkr"|jd u rg |_| j|jvr!|j| j nb| jdkr=|jd u r/g |_| j|jvr<|j| j nG| jrE| jdkr|jd u rMg |_|jd u rUg |_|jd u r]g |_| j|jvrj|j| j | j|jvrw|j| j | j|jvr|j| j | j D ]\}}|jd u ri |_t	t
jj||dp||j|< q|S )NsuccessZfailureZsuccess_and_failure)default_value)r   Zcallback_typesuccess_callbackZcallback_nameappendfailure_callback	callbackscallback_varsrD   r?   litellmutilsZ
get_secret)rH   rS   varr   r   r   r   (convert_key_logging_metadata_to_callback}   sH   








r^   c                   @   s4   e Zd ZdZedefddZedefddZdS )KeyAndTeamLoggingSettingszO
    Helper class to get the dynamic logging settings for the key and team
    user_api_key_dictc                 C   "   | j d urd| j v r| j d S d S Nloggingr7   r`   r   r   r    get_key_dynamic_logging_settings      


z:KeyAndTeamLoggingSettings.get_key_dynamic_logging_settingsc                 C   ra   rb   )team_metadatare   r   r   r   !get_team_dynamic_logging_settings   rg   z;KeyAndTeamLoggingSettings.get_team_dynamic_logging_settingsN)__name__
__module____qualname____doc__staticmethodr   rf   ri   r   r   r   r   r_      s    r_   r`   proxy_configc                 C   s   d }t | }t | }|d ur"|D ]}ttdi ||d}q|S |d ur8|D ]}ttdi ||d}q(|S | jd ur]d| jv r]	 | j}|dd pMi }tdi |}t	d| |S | j
d urjtj| j
|d}|S )N)rH   rS   callback_settingsz$Team callback settings activated: %s)team_idro   r   )r_   rf   ri   r^   r   rh   rE   r   r   debugrq   LiteLLMProxyRequestSetup$add_team_based_callbacks_from_config)r`   ro   callback_settings_objZkey_dynamic_logging_settingsZteam_dynamic_logging_settingsitemrh   rp   r   r   r   _get_dynamic_logging_metadata   sF   ,#

	
rw   Flitellm_key_header_name!forward_llm_provider_auth_headersauthenticated_with_headerc           
      C   s   ddl m} i }|dur| nd}|  D ]`\}}| }	|	dkr5||r5|du s0| dkr4|||< q|	dkrJ|rI|du sE| dkrI|||< q|rf|	tv rf|rW|	|krWq|	dkr\q|	dkraq|||< q|	tvrv|du sr|	|krv|||< q|S )a  
    Removes litellm api key from headers
    
    Args:
        headers: Request headers
        litellm_key_header_name: Custom header name for LiteLLM API key
        forward_llm_provider_auth_headers: Whether to forward provider auth headers
        authenticated_with_header: Which header was used for LiteLLM authentication
            (e.g., "x-litellm-api-key", "authorization", "x-api-key")
    
    Returns:
        Cleaned headers dict
    r   is_anthropic_oauth_keyNauthorization	x-api-keyx-litellm-api-key)#litellm.llms.anthropic.common_utilsr|   r   rD   _SPECIAL_HEADERS_CACHE)
r=   rx   ry   rz   r|   clean_headersZlitellm_key_lowerheaderr   Zheader_lowerr   r   r   r     s>   
r   c                   @   s  e Zd Zededee fddZededee fddZededee	 fddZ
ededee fd	d
Zedeeef fddZedededee fddZedee dededefddZe	d;dedee dee fddZe	d;dedee dee fddZedededefddZededededefddZedee fddZedd dededeeeef  defd!d"Zededed#edefd$d%Zededefd&d'Zededed#edefd(d)Zeded*ed#edefd+d,Z ed-ee ded#efd.d/Z!ed0ee" d1ee" de"fd2d3Z#ed4ed5e$dee% fd6d7Z&ed8ee' dededee(e  fd9d:Z)dS )<rs   r=   r2   c                 C       |  dd}|durt|S dS )a  
        Workaround for client request from Vercel's AI SDK.

        Allow's user to set a timeout in the request headers.

        Example:

        ```js
        const openaiProvider = createOpenAI({
            baseURL: liteLLM.baseURL,
            apiKey: liteLLM.apiKey,
            compatibility: "compatible",
            headers: {
                "x-litellm-timeout": "90"
            },
        });
        ```
        zx-litellm-timeoutNrE   float)r=   Ztimeout_headerr   r   r   _get_timeout_from_request<  s   z2LiteLLMProxyRequestSetup._get_timeout_from_requestc                 C   r   )zD
        Get the `stream_timeout` from the request headers.
        zx-litellm-stream-timeoutNr   )r=   Zstream_timeout_headerr   r   r    _get_stream_timeout_from_requestU     z9LiteLLMProxyRequestSetup._get_stream_timeout_from_requestc                 C   r   )zE
        Workaround for client request from Vercel's AI SDK.
        zx-litellm-num-retriesN)rE   int)r=   Znum_retries_headerr   r   r   _get_num_retries_from_request_  r   z6LiteLLMProxyRequestSetup._get_num_retries_from_requestc                 C   s,   ddl m} | dd}|dur||S dS )zI
        Get the `spend_logs_metadata` from the request headers.
        r   r   zx-litellm-spend-logs-metadataN)*litellm.litellm_core_utils.safe_json_loadsr   rE   )r=   r   Zspend_logs_metadata_headerr   r   r   -_get_spend_logs_metadata_from_request_headersi  s
   zFLiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headersc                 C   sV   i }|   D ]"\}}| dr| ds|||< q| dr(|||< q|S )z
        Get the headers that should be forwarded to the LLM Provider.

        Looks for any `x-` headers and sends them to the LLM Provider.

        [07/09/2025] - Support 'anthropic-beta' header as well.
        zx-zx-stainlesszanthropic-beta)rD   r   
startswith)r=   Zforwarded_headersr   r   r   r   r   _get_forwardable_headersu  s   
z1LiteLLMProxyRequestSetup._get_forwardable_headersr-   c                 C   s.   |   D ]\}}| | kr|  S qdS )zL
        Get a case-insensitive header from the headers dictionary.
        N)rD   r   )r=   r-   r   r   r   r   r   _get_case_insensitive_header  s
   z5LiteLLMProxyRequestSetup._get_case_insensitive_headergeneral_settingsr`   c                 C   sN   | d u r|S |  d}|s|S t|}|s|S t||}|r%||_|S |S )NZuser_header_mappings)rE   rs   %get_internal_user_header_from_mappingr   user_id)r   r`   r=   user_header_mappingheader_nameheader_valuer   r   r   #add_internal_user_from_user_mapping  s"   
z<LiteLLMProxyRequestSetup.add_internal_user_from_user_mappingNc                 C   sv   |du rdS | d}|du s|dkrdS t|ts#tdt| t| |}|dur9td| d| d |S )zg
        Get the user from the specified header if `general_settings.user_header_name` is set.
        NZuser_header_name z.Expected user_header_name to be a str but got zfound user "z" in header "")	rE   r>   r?   	TypeErrortypers   r   r
   info)r=   r   r   userr   r   r   get_user_from_headers  s   

z.LiteLLMProxyRequestSetup.get_user_from_headersc                 C   sV   |dur| ddurdS |  D ]\}}| dkr(td| d |  S qdS )z9
        Get the OpenAI Org ID from the headers.
        NZforward_openai_org_idTzopenai-organizationzfound openai org id: z, sending to llm)rE   rD   r   r
   r   )r=   r   r   r   r   r   r   get_openai_org_id_from_headers  s   z7LiteLLMProxyRequestSetup.get_openai_org_id_from_headersc                 C   sL   t | }tjdu r$t j|d}| D ]\}}|dur#||d|< q|S )z
        Add headers to the LLM call

        - Checks request headers for forwardable headers
        - Checks if user information should be added to the headers
        Tre   Nzx-litellm-{})rs   r   r[   Z#add_user_information_to_llm_headers'get_sanitized_user_information_from_keyrD   format)r=   r`   Zreturned_headersZ litellm_logging_metadata_headersr@   r   r   r   r   add_headers_to_llm_call  s   

z0LiteLLMProxyRequestSetup.add_headers_to_llm_callrH   c                 C   sz   ddl m} ddlm} | d}|dur;tjdur;tjjdur;|||tjj|j|j	dr;t
||}|i kr;|| d< | S )z<
        Add headers to the LLM call by model group
        r   )_check_model_access_helper)
llm_routermodelN)r   r   modelsteam_model_aliasesrq   r=   )Zlitellm.proxy.auth.auth_checksr   litellm.proxy.proxy_serverr   rE   r[   Zmodel_group_settings!forward_client_headers_to_llm_apir   rq   rs   r   )rH   r=   r`   r   r   Z
data_model_headersr   r   r   &add_headers_to_llm_call_by_model_group  s,   


z?LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_groupc                 C   sz   | sd S t | tr| n| g}|D ]*}t |tsq|d}|d}|d u s(|s)qt| ttj kr:|  S qd S )NZlitellm_user_roler   )r>   listrM   rE   r?   r   r   ZINTERNAL_USER)r   rD   rv   Zroler   r   r   r   r     s"   


z>LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping)r   c           	      C   s   t  }|r|ddu rt| |}|i kr||d< t| |}|dur(||d< t| }|dur5||d< t| }|durB||d< t| }|durO||d< |S )	zc
        - Adds user from headers
        - Adds forwardable headers
        - Adds org id
        r   Tr=   NZorganizationtimeoutstream_timeoutnum_retries)r   rE   rs   r   r   r   r   r   )	r=   r`   r   rH   r   Z_organizationr   r   r   r   r   r   %add_litellm_data_for_backend_llm_call'  s4   

z>LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call_metadata_variable_namec                 C   s   ddl m} | }t| }|dur||d< | d}| dp$| d}|r3||d< td	|  |rM||d
< ||d< ||d< ||d< td|  t|| tr[|| 	| |S )z
        Add litellm metadata from request headers

        Relevant issue: https://github.com/BerriAI/litellm/issues/14008
        r   )!LitellmMetadataFromRequestHeadersNspend_logs_metadatazx-litellm-agent-idrB   rC   agent_idz Extracted agent_id from header: Ztrace_idZ
session_idZlitellm_session_idZlitellm_trace_idz6Extracted chain_id from header (trace-id/session-id): )
litellm.proxy._typesr   rs   r   rE   r   rr   r>   rM   update)r=   rH   r   r   Zmetadata_from_headersr   Zagent_id_from_headerZchain_idr   r   r   )add_litellm_metadata_from_request_headersT  s8   
zBLiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headersc                 C   sR   t | j| j| j| j| j| j| j| j| j	| j
| j| j| jr!| j nd | jd}|S )N)Zuser_api_key_hashZuser_api_key_aliasuser_api_key_spenduser_api_key_max_budgetZuser_api_key_team_idZuser_api_key_project_idZuser_api_key_user_idZuser_api_key_org_idZuser_api_key_team_aliasZuser_api_key_end_user_idZuser_api_key_user_emailuser_api_key_request_routeZuser_api_key_budget_reset_atuser_api_key_auth_metadata)r$   api_key	key_aliasspend
max_budgetrq   Z
project_idr   Zorg_id
team_aliasend_user_idZ
user_emailrequest_routeZbudget_reset_at	isoformatr7   )r`   user_api_key_logged_metadatar   r   r   r     s&   z@LiteLLMProxyRequestSetup.get_sanitized_user_information_from_keyc                 C   sz   t j|d}| | | |j| | d< t|dd}| | d}|p$|}|| | d< t|dd| | d< || | d< | S )zK
        Adds the `UserAPIKeyAuth` object to the request metadata.
        re   Zuser_api_keyr   NZend_user_max_budgetZuser_api_end_user_max_budgetZuser_api_key_auth)rs   r   r   r   getattrrE   )rH   r`   r   r   Z_key_agent_idZ_existing_agent_idZ_resolved_agent_idr   r   r   )add_user_api_key_auth_to_request_metadata  s$   
zBLiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadatamanagement_endpoint_metadatac                 C   sz   || vr| S ddl m}m} i }| D ]\}}||| vr"|||< q| | ddu r2i | | d< | | d | | S )z
        Adds the `UserAPIKeyAuth` metadata to the request metadata.

        ignore any sensitive fields like logging, api_key, etc.
        r   ))LiteLLM_ManagementEndpoint_MetadataFields1LiteLLM_ManagementEndpoint_MetadataFields_Premiumr   N)r   r   r   rD   rE   r   )rH   r   r   r   r   Zadded_metadatar@   r   r   r   r   4add_management_endpoint_metadata_to_request_metadata  s$   zMLiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadatakey_metadatac                 C   sd  | d u r|S d| v r*i |d< t | d tr*| d  D ]\}}|tv r)||d |< qd| v rF| d d urFtj|| d| d d|| d< d| v rYt | d trY| d || d< d| v rt | d trd|| v rt || d tr| d  D ]\}}||| d vr||| d |< qyn| d || d< d| v rt | d tr| d |d< tj|| |d}|S )Ncachetagsrequest_tagstags_to_adddisable_global_guardrailsr   Zdisable_fallbacksrH   r   r   )	r>   rM   rD   r%   rs   _merge_tagsrE   boolr   )r   rH   r   r@   r   r-   r   r   r   r   add_key_level_controls  s`   




z/LiteLLMProxyRequestSetup.add_key_level_controlsr   r   c                 C   sJ   g }| rt | tr||  |r#t |tr#|D ]}||vr"|| q|S )a5  
        Helper function to merge two lists of tags, ensuring no duplicates.

        Args:
            request_tags (Optional[list]): List of tags from the original request
            tags_to_add (Optional[list]): List of tags to add

        Returns:
            list: Combined list of unique tags
        )r>   r   extendrW   )r   r   Z
final_tagstagr   r   r   r     s   

z$LiteLLMProxyRequestSetup._merge_tagsrq   ro   c                 C   sx   |j | d}t|trt|dkrdS i |d|}|dd |dd |dd t|dd|dd|dS )	z:
        Add team-based callbacks from the config
        )rq   r   NrZ   rq   rV   rX   )rV   rX   rZ   )Zload_team_configr>   rM   lenrE   popr   )rq   ro   Zteam_configZcallback_vars_dictr   r   r   rt   6  s   

z=LiteLLMProxyRequestSetup.add_team_based_callbacks_from_configr   c                 C   sp   d }d|v r't |d tr|d d}dd |D }nt |d tr'|d }d|v r6t |d tr6|d }|S )Nzx-litellm-tags,c                 S   s   g | ]}|  qS r   )strip)r   r   r   r   r   
<listcomp>Y      zHLiteLLMProxyRequestSetup.add_request_tag_to_metadata.<locals>.<listcomp>r   )r>   r?   r*   r   )r   r=   rH   r   Z_tagsr   r   r   add_request_tag_to_metadataM  s   z4LiteLLMProxyRequestSetup.add_request_tag_to_metadatar   )*rj   rk   rl   rn   rM   r   r   r   r   r   r   r   r   r	   r   r?   r   r   r   r   r   r   r   r   r   r   r   r   r   r$   r   r   r   r   r   r   r&   r   rt   r   r   r   r   r   r   r   rs   ;  s   		
	,3 : 
rs   r   versionc           %   
      s(  ddl m}m} ddlm}m}	 |t|}
d}|r!|dd}|s)tt	dd}d}d|j
v r3d}n
d|j
v r;d}nd	}t|j
|durI|d
nd||d}td|  td|
  |rpd	|v rp|d	 | d< td t }t|j|j|t| |d| d< t| | t|}| |ddu ri | |< | tj|||d tj|| |d || v rt| | tr|| | d< tj| ||d} t|||}t||}|dur|jdu r||_d| vr|| d< |	|
d| d< z	|j}t|}W n t y   i }Y nw |d}|dur|| d< t!| |d |dd}|r,t"|}|d| d< td|  d | v rk| d  durkt| d  tr`t#| d  | d < t| d  ts`t$d!| d    t%| d  | | d"< d#| v r| d# durt| d# trt#| d# }t|tst$d$| d#   n|| d#< t| d# tr| d# & D ]\}}|| | vr|| | |< qtj'| ||d%} || | d&< |dur|d'd| | d'< |j(}tj)|| |d(} |j*pi }d)|v r|d) durtj+| | d)|d) d*| | d)< d+|v rt|d+ t,r|d+ | | d+< d,|v r`t|d, tr`d,| | v rXt| | d, trX|d, & D ]\}}|| | d, vrU|| | d, |< q@n|d, | | d,< |j-pei }d)|v r|d) durtj+| | d)|d) d*| | d)< tj.| ||d-} |j/| | d.< |j0| | d/< |j1| | d0< |j2| | d1< |j3| | d2< |j4| | d3< |j5| | d4< |j6| | d5< |j7| | d6< |j(| | d7< |j*| | d8< t|d9d| | d:< t|d;d| | d<< || | d< t|j| | d=< |j8| | d>< t9| |d? |j:dur|j:| d@< t }dA}	 |dur?|dCdBu r?|dur?t;|dr?dD|j
v r?|j
dD }n|dur[t;|dEr[t;|j<dFr[|j<dur[|j<j=}|| | dG< dA}|duryt;|drydH|j
v ry|j
dH }|| | dI< tj>||| dJ}|dur|| | d)< t?||dK} | dur| j@| dL< | jA| dM< | jBdur| jB& D ]	\}!}"|"| |!< q|j(rdN|j(v r|j(dN }#|#rt|#tCr|#| dN< tD| ||dOI dH  tE| |dP tF| |dP tdQ|  tG| |||dR t }$tHItJjKtLjM|$| dS||$|j8dT | S )Ua  
    Adds LiteLLM-specific data to the request.

    Args:
        data (dict): The data dictionary to be modified.
        request (Request): The incoming request.
        user_api_key_dict (UserAPIKeyAuth): The user API key dictionary.
        general_settings (Optional[Dict[str, Any]], optional): General settings. Defaults to None.
        version (Optional[str], optional): Version. Defaults to None.

    Returns:
        dict: The modified data dictionary.

    r   )r   premium_user)RedactedDictSecretFieldsFry   Nr   r}   r~   rx   )rx   ry   rz   zRequest Headers: zRaw Headers: r   zUSetting client-provided x-api-key as api_key parameter (will override deployment key))r8   methodr=   bodyarrival_timeZproxy_server_request)r=   r`   r   )r=   rH   r   r=   )rH   r=   r`   r   )Zraw_headersZsecret_fieldsrJ   rK   )rH   r=   zCache-Controlzs-maxagettlzreceiving data: %sr7   z9Failed to parse 'metadata' as JSON dict. Received value: Zrequester_metadatar4   zAFailed to parse 'litellm_metadata' as JSON dict. Received value: )rH   r`   r   Zlitellm_api_versionZglobal_max_parallel_requests)r   rH   r   r   r   r   r   r   Zuser_api_key_team_max_budgetZuser_api_key_team_spendr   r   r   Zuser_api_key_model_max_budgetZ&user_api_key_end_user_model_max_budgetZuser_api_key_user_spendZuser_api_key_user_max_budgetZuser_api_key_metadataZuser_api_key_team_metadataZobject_permission_idZ!user_api_key_object_permission_idZteam_object_permission_idZ&user_api_key_team_object_permission_idZendpointZlitellm_parent_otel_spanr;   allowed_model_regionr   TZuse_x_forwarded_forzx-forwarded-forclienthostrequester_ip_addressz
user-agent
user_agent)r   r=   rH   )r`   ro   rV   rX   Zlitellm_disabled_callbacks)rH   r   r`   )rH   r`   z5[PROXY] returned data from litellm_pre_call_utils: %s)request_bodyr   r`   r   add_litellm_data_to_request)ZservicedurationZ	call_type
start_timeend_timeparent_otel_span)Nr   r   r   Z*litellm.types.proxy.litellm_pre_call_utilsr   r   r   rE   r   r[   r=   r   r   rr   timer?   r8   r   copyrR   r<   r   rs   r   r   r>   rM   r   r   r   r   rI   rN   (add_provider_specific_headers_to_requestr.   r   warningdeepcopyrD   r   r7   r   rh   r   r   project_metadatar   Zteam_max_budgetZ
team_spendr   r   r   Zmodel_max_budgetZend_user_model_max_budgetZ
user_spendZuser_max_budgetr   _add_otel_traceparent_to_datar   rL   r   r   r   rw   rV   rX   rZ   r   move_guardrails_to_metadata"_update_model_if_team_alias_exists!_update_model_if_key_alias_exists_enforced_params_checkasynciocreate_taskservice_logger_objZasync_service_success_hookr!   ZPROXY_PRE_CALL)%rH   r1   r`   ro   r   r   r   r   r   r   Z_raw_headersZforward_llm_authrz   r   r   r   r   rI   Z
query_dictZdynamic_api_versionZcache_control_headerr+   Zparsed_litellm_metadatar-   r   r   rh   r   r   r   r   r   ru   r@   r   Zdisabled_callbacksr   r   r   r   r   c  sJ  



	

























r   c                 C   s0   |  d}|r|jr||jv r|j| | d< dS )aS  
    Update the model if the team alias exists

    If a alias map has been set on a team, then we want to make the request with the model the team alias is pointing to

    eg.
        - user calls `gpt-4o`
        - team.model_alias_map = {
            "gpt-4o": "gpt-4o-team-1"
        }
        - requested_model = "gpt-4o-team-1"
    r   N)rE   r   rH   r`   _modelr   r   r   r     s   

r   c                 C   s<   |  d}|r|jrt|jtr||jv r|j| | d< dS )ap  
    Update the model if the key alias exists

    If an alias map has been set on a key, then we want to make the request with the model the key alias is pointing to

    eg.
        - user calls `modelAlias`
        - key.aliases = {
            "modelAlias": "xai/grok-4-fast-non-reasoning"
        }
        - requested_model = "xai/grok-4-fast-non-reasoning"
    r   N)rE   aliasesr>   rM   r   r   r   r   r     s   


r   c                 C   s   d }| d ur*|  d}d| v r*t|du r*| d }d|v r*|d u r#g }||d  |j dd d urA|d u r9g }||jd  |S )Nenforced_paramsservice_account_settingsT)rE   !check_if_token_is_service_accountr   r7   )r   r`   r  r  r   r   r   _get_enforced_params7  s   
r  valid_tokenc                 C   s   | j r
d| j v r
dS dS )zr
    Checks if the token is a service account

    Returns:
        bool: True if token is a service account

    Zservice_account_idTFrd   )r  r   r   r   r  M  s   
r  r   r   c                 C   s   t ||d}|du rdS |r|durtd| dtjj |D ]M}|d}t|dkr>|d | vr=td	|d  d
q t|dkrm|d | vrTtd	|d  d
|d | |d  vrmtd|d  d|d  dq dS )z]
    If enforced params are set, check if the request body contains the enforced params.
    )r   r`   NTz;Enforced Params is an Enterprise feature. Enforced Params: z. .   r   zBadRequest please pass param=z* in request body. This is a required param   zBadRequest please pass param=[z][z+] in request body. This is a required param)r  
ValueErrorr   Znot_premium_userr   r*   r   )r   r   r`   r   r  Zenforced_paramZ_enforced_paramsr   r   r   r   [  s8   	
r   r   rh   metadata_variable_namec                 C   s   ddl m} t }| r(d| v r(t| d tr(t| d dkr(|  || d  |rGd|v rGt|d trGt|d dkrG|  ||d  |rSt||| d< dS dS )a  
    Helper add guardrails from key or team metadata to request data

    Key guardrails are set first, then team guardrails are appended (without duplicates).

    Args:
        key_metadata: The key metadata dictionary to check for guardrails
        team_metadata: The team metadata dictionary to check for guardrails
        data: The request data to update
        metadata_variable_name: The name of the metadata field in data

    r   _premium_user_check
guardrailsN)litellm.proxy.utilsr  setr>   r   r   r   )r   rh   rH   r  r  Zcombined_guardrailsr   r   r   )_add_guardrails_from_key_or_team_metadata  s"   r  c                 C   s
  ddl m} ddlm} ddlm} ddlm} ddlm	} t
 }	| r@d| v r@t| d tr@t| d dkr@|  |	| d  |r_d|v r_t|d tr_t|d dkr_|  |	|d  |	scdS |d	|	  | }
|
 sy|d
 dS ||dd}|
 }t
 }|	D ]+}|
|r|j|||d}||j |d| d|j  q|d| d q|sdS ||vri ||< || dg }t|tsg }t
|}|| t||| d< d|| vrg || d< || d t|	 |dt|  dS )a#  
    Helper to resolve guardrails from policies attached to key/team metadata.

    This function:
    1. Gets policy names from key and team metadata
    2. Resolves guardrails from those policies (including inheritance)
    3. Adds resolved guardrails to request metadata

    Args:
        key_metadata: The key metadata dictionary to check for policies
        team_metadata: The team metadata dictionary to check for policies
        data: The request data to update
        metadata_variable_name: The name of the metadata field in data
    r   r   get_policy_registryPolicyResolverr  r'   policiesNz<Policy engine: resolving guardrails from key/team policies: zGPolicy engine not initialized, skipping policy resolution from metadatar   )r   )policy_namer  contextz0Policy engine: resolved guardrails from policy 'z': zPolicy engine: policy 'z' not found in registryr  Zapplied_policieszLPolicy engine: added guardrails from key/team policies to request metadata: )litellm._loggingr   +litellm.proxy.policy_engine.policy_registryr  +litellm.proxy.policy_engine.policy_resolverr  r  r  !litellm.types.proxy.policy_enginer(   r  r>   r   r   r   rr   is_initializedrE   get_all_policiesZ
has_policyZresolve_policy_guardrailsr  r   r   )r   rh   rH   r  r   r  r  r  r(   policy_namesregistryr  Zall_policiesresolved_guardrailsr  Zresolved_policyexisting_guardrailscombinedr   r   r   )_add_guardrails_from_policies_in_metadata  s|   



r&  r   c                    sf  |j }|j}|od|v pd|v }|od|v pd|v }d| v p&d| v p&d| v }|s@|s@|s@ddlm} |  s@| dd dS t|j |j| |d t|j |j| |d t| ||dI dH  d| v r| d}	d| | v rt	| | d t
r| | d |	 n|	| | d< d| v r| d}
d| | v rt	| | d tr| | d |
 dS |
| | d< dS dS )	aa  
    Helper to add guardrails from request to metadata

    - If guardrails set on API Key metadata then sets guardrails on request metadata
    - If guardrails not set on API key, then checks request metadata
    - Adds guardrails from policies attached to key/team metadata
    - Adds guardrails from policy engine based on team/key/model context
    r  r  Zguardrail_configr   r  N)r   rh   rH   r  )rH   r  r`   )r7   rh   r  r  r  r   r  r&  !add_guardrails_from_policy_enginer>   r   r   rM   r   )rH   r   r`   r   rh   Zhas_key_configZhas_team_configZhas_request_configr  Zrequest_body_guardrailsZrequest_body_guardrail_configr   r   r   r      sj   





r   sc                 C   s    ddl m} t| to| |S )zPReturn True if string is a policy version ID (starts with policy_<uuid> prefix).r   POLICY_VERSION_ID_PREFIX)r  r*  r>   r?   r   r(  r*  r   r   r   _is_policy_version_id{  s   r,  c                 C   s0   ddl m} t| sdS | t|d  pdS )zNExtract raw UUID from policy_<uuid> string, or None if not a valid version ID.r   r)  N)r  r*  r,  r   r   r+  r   r   r   _extract_policy_id  s   r-  r  r(   request_body_policiespolicies_overridec                    s  ddl m} ddlm}m} ddlm} ddlm} | }	|		|}
dd |
D }dd	 |
D  |
d
|  t|}|rPt|trP|| |
d|  |sVg i fS |jt|||d}|
d|  |D ]}|| |d qj fdd	|D }|| |d | fS )z
    Match policies via attachments and request body, track them in metadata.

    Returns:
        Tuple of (applied_policy_names, policy_reasons)
    r   r  )add_policy_sources_to_metadata%add_policy_to_applied_policies_header)get_attachment_registry)PolicyMatcherc                 S   s   g | ]}|d  qS )r  r   r   mr   r   r   r     r   z-_match_and_track_policies.<locals>.<listcomp>c                 S   s   i | ]	}|d  |d qS )r  Zmatched_viar   r4  r   r   r   rA     s    z-_match_and_track_policies.<locals>.<dictcomp>z1Policy engine: matched policies via attachments: z9Policy engine: added dynamic policies from request body: )r!  r  r  z6Policy engine: applied policies (conditions matched): )request_datar  c                    s   i | ]}| v r| | qS r   r   )r   nameZpolicy_reasonsr   r   rA     s
    )r6  Zpolicy_sources)r  r   Z)litellm.proxy.common_utils.callback_utilsr0  r1  Z/litellm.proxy.policy_engine.attachment_registryr2  Z*litellm.proxy.policy_engine.policy_matcherr3  Z"get_attached_policies_with_reasonsrr   r  r>   r   r   Z%get_policies_with_matching_conditions)rH   r  r.  r/  r   r0  r1  r2  r3  Zattachment_registryZmatches_with_reasonsZmatching_policy_namesZall_policy_namesapplied_policy_namesr  Zapplied_reasonsr   r8  r   _match_and_track_policies  sL   

r:  r!  r  c                 C   s
  ddl m} ddlm} |j|||d}|d|  |j|||d}|| vr,i | |< t }	|rO||}	|| | d< |	| | d< |dt	| d	|	  |sU|sUd
S | | 
dg }
t|
tsdg }
t|
}|| ||	8 }t|| | d< |dt|  d
S )z<Apply resolved guardrails and pipelines to request metadata.r   r  r  )r  r  r!  z$Policy engine: resolved guardrails: Z_guardrail_pipelinesZ_pipeline_managed_guardrailszPolicy engine: resolved z" pipeline(s), managed guardrails: Nr  z5Policy engine: added guardrails to request metadata: )r  r   r  r  Zresolve_guardrails_for_contextrr   Zresolve_pipelines_for_contextr  Zget_pipeline_managed_guardrailsr   rE   r>   r   r   )rH   r  r  r!  r  r   r  r#  Z	pipelinesZpipeline_managed_guardrailsr$  r%  r   r   r   &_apply_resolved_guardrails_to_metadata  sX   

r;  c              
      s  ddl m} ddlm} ddlm} ddlm} | dd}| }|	d|
  d	t|   |
 s>|	d
 dS || pCd}	|j}
|j}|t|
trR|
ndt|trZ|nd| d|	d}|	d|j d|j d|j d|j  g }g }|rt|tr|D ]}t|tsqt|rt|}|r|| q|| qt| }g }|D ].}|j|d}|dur|\}}|||< || |	d| d|  q|	d| d q|| }t| |||r|ndd\}}t| |||r|nd|r|ndd dS )ah  
    Add guardrails from the policy engine based on request context.

    This function:
    1. Extracts "policies" from request body (if present) for dynamic policy application
    2. Supports policy_<uuid> in policies to execute a specific version (e.g. published)
    3. Gets matching policies based on team_alias, key_alias, and model (via attachments)
    4. Combines dynamic policies with attachment-based policies
    5. Resolves guardrails from all policies (including inheritance)
    6. Adds guardrails to request metadata
    7. Tracks applied policies in metadata for response headers
    8. Removes "policies" from request body so it's not forwarded to LLM provider

    Args:
        data: The request data to update
        metadata_variable_name: The name of the metadata field in data
        user_api_key_dict: The user's API key authentication info
    r   r  )get_tags_from_request_bodyr  r'   r  Nz$Policy engine: registry initialized=z, policy_count=z7Policy engine not initialized, skipping policy matchingr   )r   r   r   r   z8Policy engine: matching policies for context team_alias=z, key_alias=z, model=z, tags=)	policy_idz+Policy engine: loaded version by ID policy_z -> zPolicy engine: policy version z not found in cache, skipping)r/  )r!  r  )r  r   -litellm.proxy.common_utils.http_parsing_utilsr<  r  r  r  r(   r   rr   r  r   r   r   r   r>   r?   rE   r   r   r   r,  r-  rW   rM   Zget_policy_by_id_for_requestr:  r;  )rH   r  r`   r   r<  r  r(   Zrequest_body_policies_rawr"  Zall_tagsZ_team_aliasZ
_key_aliasr  Zrequest_body_namesZrequest_body_version_idsrv   r=  Zmerged_policiesZfetched_policy_namesresultZpnamepolicyr.  r9  _r   r   r   r'    s   







	


r'  c                 C   s   ddl m} i }d}tD ]}||v r|| }|||< d}q| D ]\}}| dkr7||r7|||< d} nq!|du rRttjj dtj	j dtj
j |d| d< d S )	Nr   r{   FTr}   r   )Zcustom_llm_providerextra_headersZprovider_specific_header)r   r|   r    rD   r   r#   r"   Z	ANTHROPICr   ZBEDROCKZ	VERTEX_AI)rH   r=   r|   Zanthropic_headersZadded_headerr   r   r   r   r   r   r     s*   
r   c                 C   s   ddl m} | d u rd S |d u rd S tjdu r8|jr:d|jv r<d| vr'i | d< | d }d|vr>|jd |d< d S d S d S d S d S )Nr   )open_telemetry_loggerTZtraceparentrB  )r   rC  r[   Z#forward_traceparent_to_llm_providerr=   )rH   r1   rC  Z_exra_headersr   r   r   r     s"   

r   )NFN)NNr   )Tr   r   r   typingr   r   r   r   r   r   Zfastapir   Zstarlette.datastructuresr	   r[   r  r
   r   Zlitellm._service_loggerr   r   r   r   r   r   r   r   r   r   r   r>  r   	frozenset_member_map_valuesr   Zlitellm.routerr   Zlitellm.types.llms.anthropicr    Zlitellm.types.servicesr!   Zlitellm.types.utilsr"   r#   r$   r%   r   r   r&   Z_ProxyConfigr  r(   r.   r:   r?   r<   rG   rM   rR   r^   r_   rw   r   r   rs   r   r   r   r   r  r  r   r  r&  r   r,  r-  tupler:  r;  r'  r   r   r   r   r   r   <module>   sz    $	"
.
@
:    1
   "



'
.
p
[
K

C
o
"