Ë
    tPêjk3  ã                   óÞ  — U d dl Z d dlZej                  j                  d e j                  j	                  d«      «       d dlZd dlZd dlmZmZ d dl	Z	d dl
mZ d dlZd dlmZ d dlmZmZ d dlmZm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m Z  dZ!ee"d<   ejF                  jH                  d ejJ                  jH                  dejL                  jH                  diZ'ee"d<    G d„ de«      Z(y)é    Nz../..)ÚFinalÚLiteral)ÚHTTPException)Úverbose_proxy_logger)ÚCustomGuardrailÚlog_guardrail_information)Úget_async_httpx_clientÚhttpxSpecialProvider)ÚUserAPIKeyAuth)Ú should_proceed_based_on_metadata)Ú
get_secret)ÚGuardrailEventHooksÚGuardrailItemÚLakeraCategoryThresholdsÚRoleÚdefault_rolesÚlakera_prompt_injectionÚGUARDRAIL_NAMEé   é   ÚINPUT_POSITIONING_MAPc                   ó  ‡ — e Zd Zedee   fd„«       Z	 	 	 	 dded   dedz  de	dz  de	dz  fˆ fd	„Z
d
eddfd„Zdededed   fd„Zededej"                  deded   dee	z  ez  dz  f
d„«       Zedededed   fd„«       Zˆ xZS )ÚlakeraAI_ModerationÚreturnc                 óB   — t         j                  t         j                  gS )N)r   Úpre_callÚduring_call)Úclss    úw/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.pyÚget_supported_event_hooksz-lakeraAI_Moderation.get_supported_event_hooks2   s    € ô  ×(Ñ(Ü×+Ñ+ð
ð 	
ó    NÚmoderation_check)r   Úin_parallelÚcategory_thresholdsÚapi_baseÚapi_keyc                 óX  •— |j                  dt        | j                  «       «      «       t        t        j
                  ¬«      | _        |xs# t        j                  j                  d«      xs d| _
        || _        || _        |xs t        d«      xs d| _        t        ‰| �@  di |¤Ž y )NÚsupported_event_hooks)Úllm_providerÚLAKERA_API_KEYÚ ÚLAKERA_API_BASEzhttps://api.lakera.ai© )Ú
setdefaultÚlistr    r	   r
   ÚGuardrailCallbackÚasync_handlerÚosÚenvironÚgetÚlakera_api_keyr"   r$   r   r%   ÚsuperÚ__init__)Úselfr"   r$   r%   r&   ÚkwargsÚ	__class__s         €r   r7   zlakeraAI_Moderation.__init__9   s�   ø€ ð 	×ÑÐ1´4¸×8VÑ8VÓ8XÓ3YÔZÜ3ÔAU×AgÑAgÔhˆÔØ%ÒO¬¯©¯©Ð8HÓ)IÒOÈRˆÔØ 0ˆÔØ#6ˆÔ Ø Ò\¤JÐ/@Ó$AÒ\ÐE\ˆŒÜ‰ÑÑ"˜6Ó"r!   Úresponsec                 óÒ  — |j                  dg «      }t        |«      dk  ry |d   j                  dd«      }|d   j                  dd «      }| j                  �||�yt        di |¤Ž}d|v r3d| j                  v r%|d   | j                  d   k\  rt	        dd|d	œ¬
«      ‚d|v r3d| j                  v r%|d   | j                  d   k\  rt	        dd|d	œ¬
«      ‚y |du rt	        dd|d	œ¬
«      ‚y )NÚresultsr   ÚflaggedFÚcategory_scoresÚ	jailbreaki�  zViolated jailbreak threshold)ÚerrorÚlakera_ai_response)Ústatus_codeÚdetailÚprompt_injectionz#Violated prompt_injection thresholdTzViolated content safety policyr-   )r4   Úlenr$   r   r   )r8   r;   Ú_resultsr>   r?   Útyped_cat_scoress         r   Ú_check_response_flaggedz+lakeraAI_Moderation._check_response_flaggedJ   s2  € Ø"Ÿ,™, y°"Ó5ˆÜˆx‹=˜AÒØà! !™Ÿ™¨°EÓ:ˆØ.6°q©k¯o©oÐ>OÐQUÓ.Vˆà×#Ñ#Ð/ØÐ*Ü*BÑ*UÀ_Ñ*UÐ ØÐ"2Ñ2°{Àd×F^ÑF^Ñ7^à'¨Ñ4¸×8PÑ8PÐQ\Ñ8]Ò]Ü+Ø(+à)GØ6>ñ$ôð ð &Ð)9Ñ9Ð>PÐTX×TlÑTlÑ>lØ'Ð(:Ñ;¸t×?WÑ?WÐXjÑ?kÒkÜ+Ø(+à)NØ6>ñ$ôð ð  	ð ˜‰_ÜØà=Ø*2ñôð ð 	r!   ÚdataÚuser_api_key_dictÚ	call_type)Ú
completionÚtext_completionÚ
embeddingsÚimage_generationÚ
moderationÚaudio_transcriptionÚpass_through_endpointÚrerankÚ	responsesÚmcp_callÚanthropic_messagesc              ƒ   ó   K  — t        |t        ¬«      ƒ d {  –—† du ry d}d}d|v �rEt        |d   t        «      �r1t        j
                  j                  d«      }|�|j                  }nd }|€t        }g }|�U|D ]P  }	t        |	t        «      r|j                  |	j                  «       Œ/t        |	t        «      sŒ@|j                  |	«       ŒR t        D �	ci c]  }	|	d “Œ }
}	d }g }|d   D ]\  }|j                  d«      }	|	|v sŒd|v r	g |¢|d   ¢}|	t        j                  j                  k(  r|}ŒF|	|j                  d«      d	œ|
|	<   Œ^ |�£t        j                  s“|j                  d«      }g }|D ]  }d
|v sŒ|j                  |d
   d   «       Œ  t!        |«      dkD  r|ddj#                  |«      z   z  }t        j                  j                  |d	œ|
t        j                  j                  <   t%        |
j'                  «       d„ ¬«      D ��cg c]	  \  }}|�|‘Œ }}}t!        |«      dk(  rt)        j*                  d«       y d|i}t-        j.                  |fi | j1                  |¬«      ¤Ž}nšd|v rBt        |d   t        «      r/|d   }t-        j.                  d|i| j1                  |¬«      ¥«      }nTd|v rPt        |d   t        «      r=dj#                  |d   «      }t-        j.                  d|i| j1                  |¬«      ¥«      }t)        j*                  d|«       	 	 | j2                  j5                  | j6                  › d�|d| j8                  z   ddœ¬«      ƒ d {  –—† }t)        j*                  d|jB                  «       |jD                  dk(  r"	 | jG                  |j-                  «       ¬«       y y 7 �Œ�c c}	w c c}}w 7 Œd# t:        j<                  $ r$}t?        |j@                  jB                  «      ‚d }~ww xY w­w)N)rJ   Úguardrail_nameFr+   ÚmessagesrE   ÚroleÚ
tool_callsÚcontent)r[   r]   ÚfunctionÚ	argumentsr   z Function Input: ú c                 ó   — t         | d      S )Nr   )r   )Úxs    r   ú<lambda>z,lakeraAI_Moderation._check.<locals>.<lambda>Ê   s   € ÔLaÐbcÐdeÑbfÑLg€ r!   )Úkeyz>Skipping lakera prompt injection, no roles with messages foundÚinput)Úrequest_dataú
zLakera AI Request Args %sz/v1/prompt_injectionzBearer zapplication/json)ÚAuthorizationzContent-Type)ÚurlrJ   ÚheaderszLakera AI response: %séÈ   )r;   )$r   r   Ú
isinstancer/   ÚlitellmÚguardrail_name_config_mapr4   Úenabled_rolesr   r   ÚappendÚvalueÚstrr   ÚSYSTEMÚadd_function_to_promptrF   ÚjoinÚsortedÚitemsr   ÚdebugÚjsonÚdumpsÚ)get_guardrail_dynamic_request_body_paramsr1   Úpostr%   r5   ÚhttpxÚHTTPStatusErrorÚ	Exceptionr;   ÚtextrC   rI   )r8   rJ   rK   rL   r€   Ú
_json_dataÚprompt_injection_objro   Ústringified_rolesr[   Úlakera_input_dictÚsystem_messageÚtool_call_messagesÚmessager]   Úfunction_inputÚ	tool_callÚkÚvÚlakera_inputÚ_datar;   Úes                          r   Ú_checkzlakeraAI_Moderation._checks   sI  è ø€ ô& 3ØÜ-ô÷ ð ñ	ð ØˆØˆ
Ø˜Ò¤*¨T°*Ñ-=¼tÕ"DÜ9@×9ZÑ9Z×9^Ñ9^Ð_qÓ9rÐ Ø#Ð/Ø 4× BÑ B‘à $�àÐ$Ü -�à24ÐØÐ(Ø)ò 7�DÜ! $¬Ô-Ø)×0Ñ0°·±Õ<Ü# D¬#Õ.Ø)×0Ñ0°Õ6ð	7ô
 F[Ö-[¸T¨d°D©jÐ-[ÐÐ-[Ø!ˆNØ')ÐØ 
Ñ+ò �Ø—{‘{ 6Ó*�ØÐ,Ò,Ø# wÑ.ð.Ø/ð.à$ \Ñ2ð.Ð*ð œtŸ{™{×0Ñ0Ò0Ø)0˜Ø ð !%Ø#*§;¡;¨yÓ#9ñ/Ð% dÒ+ðð, Ð)Ü×5Ò5Ø,×0Ñ0°Ó;�GØ,.�NØ%7ò V˜	Ø%¨Ò2Ø*×1Ñ1°)¸JÑ2GÈÑ2TÕUðVô ˜>Ó*¨QÒ.ØÐ#6¸¿¹À.Ó9QÑ#QÑQ˜ä $§¡× 1Ñ 1Ø#*ñ<Ð%¤d§k¡k×&7Ñ&7Ñ8ô #Ð#4×#:Ñ#:Ó#<ÑBgÔh÷#á�A�qØ�=ò ð#ˆLñ #ô
 �<Ó  AÒ%Ü$×*Ñ*Ð+kÔlØØ# \Ð2ˆEÜŸ™Øñà×@Ñ@ÈdÐ@ÓSñ‰Jð ˜‰_¤¨D°©M¼3Ô!?Ø˜‘=ˆDÜŸ™à˜Tðà×DÑDÐRVÐDÓWðó‰Jð ˜‰_¤¨D°©M¼4Ô!@Ø—9‘9˜T '™]Ó+ˆDÜŸ™à˜Tðà×DÑDÐRVÐDÓWðóˆJô 	×"Ñ"Ð#>À
ÔKð
	ð
	-Ø$(×$6Ñ$6×$;Ñ$;Ø—}‘}�oÐ%9Ð:Øà%.°×1DÑ1DÑ%DØ$6ñð %<ó %÷ ˆHô 	×"Ñ"Ð#;¸X¿]¹]ÔKØ×Ñ 3Ò&ðð2 ×(Ñ(°(·-±-³/Ð(ÕBð7 'ðwûò4 .\ùóP#ð\ùô ×$Ñ$ò 	-Ü˜AŸJ™JŸO™OÓ,Ð,ûð	-üsz   ‚O>—N4˜B!O>Â:O>Ã
N7Ã#O>ÄA4O>Å8BO>ÈN<ÈDO>Ì>O ÍOÍO Í"AO>Î7O>ÏO ÏO;ÏO6Ï6O;Ï;O>Úcache)
rM   rN   rO   rP   rQ   rR   rS   rT   rV   rW   c              ƒ   óÊ   K  — ddl m} | j                  €| j                  dk(  r!y | j	                  ||j
                  ¬«      dury | j                  |||¬«      ƒ d {  –—† S 7 Œ­w)Nr   ©r   r#   ©rJ   Ú
event_typeT©rJ   rK   rL   )Úlitellm.types.guardrailsr   Ú
event_hookr"   Úshould_run_guardrailr   r�   )r8   rK   r�   rJ   rL   r   s         r   Úasync_pre_call_hookz'lakeraAI_Moderation.async_pre_call_hook  si   è ø€ õ& 	Aà�?‰?Ð"Ø×$Ñ$¨Ò5Øð ×(Ñ(¨dÐ?R×?[Ñ?[Ð(Ó\ÐdhÑhØà—[‘[ dÐ>OÐ[d�[Óe×eÐeÐeús   ‚AA#ÁA!ÁA#)rM   rO   rP   rQ   rR   rU   rV   rW   c              ƒ   óÎ   K  — | j                   €| j                  dk(  r)y ddlm} |j                  }| j                  ||¬«      dury | j                  |||¬«      ƒ d {  –—† S 7 Œ­w)Nr   r   r’   r“   Tr•   )r—   r"   r–   r   r   r˜   r�   )r8   rJ   rK   rL   r   r”   s         r   Úasync_moderation_hookz)lakeraAI_Moderation.async_moderation_hook>  sm   è ø€ ð  �?‰?Ð"Ø×$Ñ$¨
Ò2Øõ Eà5H×5TÑ5TˆJØ×(Ñ(¨d¸zÐ(ÓJÐRVÑVØà—[‘[ dÐ>OÐ[d�[Óe×eÐeÐeús   ‚AA%ÁA#ÁA%)r#   NNN)Ú__name__Ú
__module__Ú__qualname__Úclassmethodr/   r   r    r   r   rr   r7   ÚdictrI   r   r�   r   rm   Ú	DualCacher   r™   r›   Ú__classcell__)r:   s   @r   r   r   1   s`  ø„ Øð
¨$Ð/BÑ*Cò 
ó ð
ð @MØ?CØ#Ø"ñ#à!Ð";Ñ<ð#ð 6¸Ñ<ð#ð ˜‘*ð	#ð
 �t‘õ#ð"'°ð '¸ó 'ðRiCàðiCð *ðiCð ð
"ñ
ó	iCðV ðfà)ðfð × Ñ ðfð ð	fð
 ð	"ñ
ðfð" 
�S‰˜4Ñ	 $Ñ	&ò#fó ðfð> ðfàðfð *ðfð ð"ñ	
ò	fó ôfr!   r   ))r2   ÚsysÚpathÚinsertÚabspathry   Útypingr   r   r}   Úfastapir   rm   Úlitellm._loggingr   Ú%litellm.integrations.custom_guardrailr   r   Ú&litellm.llms.custom_httpx.http_handlerr	   r
   Úlitellm.proxy._typesr   Ú*litellm.proxy.guardrails.guardrail_helpersr   Úlitellm.secret_managers.mainr   r–   r   r   r   r   r   r   Ú__annotations__rs   rq   ÚUSERÚ	ASSISTANTr   r   r-   r!   r   ú<module>r²      s¸   ðô 
Û 
à ‡�‡���2—7‘7—?‘? 7Ó+Ô ,Û Û 
ß !ã Ý !ã Ý 1÷÷õ 0Ý WÝ 3÷õ ð 2€�Ó 1ð 	‡K�K×Ñ�qØ‡I�I‡O�O�QØ‡N�N×Ñ˜!ð Ð �uó ôhf˜/õ hfr!   