g% pddlmZmZmZmZmZddlmZddlm Z ddl m Z m Z ddl mZGdde ZdZy ) )DictListLiteralOptionalUnion)verbose_logger) CustomLogger)DynamicGuardrailParamsGuardrailEventHooks)#StandardLoggingGuardrailInformationc eZdZ ddeedeeedeeffd Zdede eeee ee fffdZ d e eeee ee ffde fd Zd ede fd Zd edefdZde fdZde eeefd ededddfdZxZS)CustomGuardrailNguardrail_namesupported_event_hooks event_hookc ||_||_||_|r|r||vrtd|d|t |di|y)Nz Event hook z% is not in the supported event hooks )rrr ValueErrorsuper__init__)selfrrrkwargs __class__s Z/var/www/openai/venv/lib/python3.12/site-packages/litellm/integrations/custom_guardrail.pyrzCustomGuardrail.__init__ s\-%:"9C j0EE !*-RShRij "6"datareturncZ|jdxsi}|jdxsg}|S)zF Returns the guardrail(s) to be run from the metadata metadata guardrails)get)rrrrequested_guardrailss rget_guardrail_from_metadataz+CustomGuardrail.get_guardrail_from_metadatas1 88J'-2'||L9?R##rr"c|D]E}t|tr|j|vs"yt|ts5|j|k(sEyy)NTF) isinstancedictrstr)rr" _guardrails r%_guardrail_is_in_requested_guardrailsz5CustomGuardrail._guardrail_is_in_requested_guardrails(sI/J*d+&&*4J,&&*4 /r event_typec&|j|}tjd|j||j||jr!|j |s|j dk7ry|jr|j|j k7ryy)Nzxinside should_run_guardrail for guardrail=%s event_type= %s guardrail_supported_event_hooks= %s requested_guardrails= %s logging_onlyFT)r#rdebugrrr)value)rrr*r"s rshould_run_guardrailz$CustomGuardrail.should_run_guardrail5s#??E G     OO   OO>>?ST  N2 ??t*2B2BBr request_datac|j|}|D]c}t|ts|j|vs#t di||j}|j duricS|j dicSiS)a Returns `extra_body` to be added to the request body for the Guardrail API call Use this to pass dynamic params to the guardrail API call - eg. success_threshold, failure_threshold, etc. ``` [{"lakera_guard": {"extra_body": {"foo": "bar"}}}] ``` Will return: for guardrail=`lakera-guard`: { "foo": "bar" } Args: request_data: The original `request_data` passed to LiteLLM Proxy T extra_bodyr)r#r%r&rr _validate_premium_userr!)rr0r" guardrailguardrail_configs r)get_guardrail_dynamic_request_body_paramsz9CustomGuardrail.get_guardrail_dynamic_request_body_paramsLs$ $?? M.I)T*t/B/Bi/O;Q< 3 34< ..0<I(++L"==. rcvddlm}m}|dur-tjd|j j yy)z< Returns True if the user is a premium user r)CommonProxyErrors premium_userTz5Trying to use premium guardrail without premium user F)litellm.proxy.proxy_serverr8r9rwarningnot_premium_userr.)rr8r9s rr3z&CustomGuardrail._validate_premium_useros> O t #  " "GHYHjHjHpHpGqr rguardrail_json_responseguardrail_status)successfailurecRddlm}|dur0tjd|jd|j yt |tr t|}t|j|j ||}d|vr ||dd <yd |vr ||d d <ytjd y) z Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc. r)r9Tz`Guardrail Tracing is only available for premium users. Skipping guardrail logging for guardrail=z event_hook=N)rguardrail_modeguardrail_responser>r&standard_logging_guardrail_informationlitellm_metadatazFunable to log guardrail information. No metadata found in request_data) r:r9rr;rrr% Exceptionr'r )rr=r0r>r9slgs r:add_standard_logging_guardrail_information_to_request_datazJCustomGuardrail.add_standard_logging_guardrail_information_to_request_data|s < t #  " "rswtGtGsHHTUYUdUdTef   -y 9&)*A&B #1..??6-    %QTL $%M N < / + ,8   " "X r)NNN)__name__ __module__ __qualname__rr'rr rr&rrr r#boolr)r/r6r3rFrrH __classcell__)rs@rrr s)-EI48 #  # (-@(AB#01 #&$$ tCy$tC)?$?@AA B$ #DItD>T9T4U/V$VW   5HT.!d!t!F  !!&y#t';!<!!""67 !  !rrcddlddl}dd|jfd|jfd|jfd}|S)ad Decorator to add standard logging guardrail information to any function Add this decorator to ensure your guardrail response is logged to DataDog, OTEL, s3, GCS etc. Logs for: - pre_call - during_call - TODO: log post_call. This is more involved since the logs are sent to DD, s3 before the guardrail is even run rNc.|j||d|S)Nr?r=r0r>rH)rresponser0s rprocess_responsez3log_guardrail_information..process_responses& GG$,%& H rc.|j||d|)Nr@rPrQ)rer0s r process_errorz0log_guardrail_information..process_errors& GG$%%& H rcK|d}|jdxs|jdxsi} |i|d{}|||S7#t$r}|||cYd}~Sd}~wwxYwwNrrr0r!rF argsrrr0rRrUfuncrVrSs r async_wrapperz0log_guardrail_information..async_wrappers} $Q JJv  B&**^"< B  8!42622H#D(LA A3 8 q,7 7 8sD.A. AA  A A. A A+ A& A+!A.&A++A.c|d}|jdxs|jdxsi} |i|}|||S#t$r}|||cYd}~Sd}~wwxYwrXrYrZs r sync_wrapperz/log_guardrail_information..sync_wrappersr $Q JJv  B&**^"< B  8T,V,H#D(LA A 8 q,7 7 8sA A AAAcFjr|i|S|i|S)N)iscoroutinefunction)r[rr]asyncior\r_s rwrapperz*log_guardrail_information..wrappers1  & &t , $1&1 1T,V,,r)rb functoolswraps)r\rdrcr]rbrVrSr_s` @@@@@rlog_guardrail_informationrfss__T 8 8__T 8 8__T-- NrN)typingrrrrrlitellm._loggingr"litellm.integrations.custom_loggerr litellm.types.guardrailsr r litellm.types.utilsr rrfrrrrls+77+;PCTlTn<r