
    3fi'<                        d Z ddlmZ ddlZddlmZ ddlmZmZm	Z	m
Z
 ddlmZmZmZ ddlmZmZmZmZ ddlmZmZ dd	lmZmZ erdd
lmZ e	d   ZdZ G d de      Z G d deee   ef         Z ddgZ!y)z>Agent middleware that integrates OpenAI's moderation endpoint.    )annotationsN)Sequence)TYPE_CHECKINGAnyLiteralcast)AgentMiddleware
AgentStatehook_config)	AIMessageBaseMessageHumanMessageToolMessage)AsyncOpenAIOpenAI)
ModerationModerationModel)Runtime)inputoutputtoolzQI'm sorry, but I can't comply with that request. It was flagged for {categories}.c                  8     e Zd ZdZ	 	 	 	 	 	 	 	 	 	 d fdZ xZS )OpenAIModerationErrorzKRaised when OpenAI flags content and `exit_behavior` is set to ``"error"``.c               N    t         |   |       || _        || _        || _        y)a  Initialize the error with violation details.

        Args:
            content: The content that was flagged.
            stage: The stage where the violation occurred.
            result: The moderation result from OpenAI.
            message: The error message.
        N)super__init__contentstageresult)selfr   r   r   message	__class__s        k/var/www/auto_recruiter/arenv/lib/python3.12/site-packages/langchain_openai/middleware/openai_moderation.pyr   zOpenAIModerationError.__init__   s'      	!
    )
r   strr   ViolationStager   r   r!   r%   returnNone)__name__
__module____qualname____doc__r   __classcell__r"   s   @r#   r   r      sB    U  	
   
 r$   r   c            	          e Zd ZdZddddddddd	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d fdZ edg	      	 	 	 	 	 	 dd
       Z edg	      	 	 	 	 	 	 dd       Z edg	      	 	 	 	 	 	 dd       Z edg	      	 	 	 	 	 	 dd       Z		 	 	 	 d dZ
	 	 	 	 d dZ	 	 	 	 d dZ	 	 	 	 d dZ	 	 	 	 d dZ	 	 	 	 d dZ	 	 	 	 d dZ	 	 	 	 d dZ	 	 	 	 	 	 	 	 	 	 	 	 d!dZd"dZd"dZd#dZd$dZd%dZ	 	 	 	 	 	 d&dZd'dZ xZS )(OpenAIModerationMiddlewarez:Moderate agent traffic using OpenAI's moderation endpoint.zomni-moderation-latestTFendN)modelcheck_inputcheck_outputcheck_tool_resultsexit_behaviorviolation_messageclientasync_clientc                   t         	|           || _        || _        || _        || _        || _        || _        || _        || _	        y)a  Create the middleware instance.

        Args:
            model: OpenAI moderation model to use.
            check_input: Whether to check user input messages.
            check_output: Whether to check model output messages.
            check_tool_results: Whether to check tool result messages.
            exit_behavior: How to handle violations
                (`'error'`, `'end'`, or `'replace'`).
            violation_message: Custom template for violation messages.
            client: Optional pre-configured OpenAI client to reuse.
                If not provided, a new client will be created.
            async_client: Optional pre-configured AsyncOpenAI client to reuse.
                If not provided, a new async client will be created.
        N)
r   r   r2   r3   r4   r5   r6   r7   _client_async_client)
r    r2   r3   r4   r5   r6   r7   r8   r9   r"   s
            r#   r   z#OpenAIModerationMiddleware.__init__4   sN    6 	
&("4*!2)r$   )can_jump_toc                    | j                   s| j                  syt        |j                  dg             }|sy| j	                  |      S )a  Moderate user input and tool results before the model is called.

        Args:
            state: Current agent state containing messages.
            runtime: Agent runtime context.

        Returns:
            Updated state with moderated messages, or `None` if no changes.
        Nmessages)r3   r5   listget_moderate_inputsr    stateruntimer?   s       r#   before_modelz'OpenAIModerationMiddleware.before_modelZ   sD     (?(?		*b12$$X..r$   c                z    | j                   syt        |j                  dg             }|sy| j                  |      S )a  Moderate model output after the model is called.

        Args:
            state: Current agent state containing messages.
            runtime: Agent runtime context.

        Returns:
            Updated state with moderated messages, or `None` if no changes.
        Nr?   )r4   r@   rA   _moderate_outputrC   s       r#   after_modelz&OpenAIModerationMiddleware.after_modelp   s<       		*b12$$X..r$   c                   K   | j                   s| j                  syt        |j                  dg             }|sy| j	                  |       d{   S 7 w)zAsync version of before_model.

        Args:
            state: Current agent state containing messages.
            runtime: Agent runtime context.

        Returns:
            Updated state with moderated messages, or `None` if no changes.
        Nr?   )r3   r5   r@   rA   _amoderate_inputsrC   s       r#   abefore_modelz(OpenAIModerationMiddleware.abefore_model   sN      (?(?		*b12++H5555s   AAAAc                   K   | j                   syt        |j                  dg             }|sy| j                  |       d{   S 7 w)zAsync version of after_model.

        Args:
            state: Current agent state containing messages.
            runtime: Agent runtime context.

        Returns:
            Updated state with moderated messages, or `None` if no changes.
        Nr?   )r4   r@   rA   _amoderate_outputrC   s       r#   aafter_modelz'OpenAIModerationMiddleware.aafter_model   sF        		*b12++H5555s   A A	AA	c                   t        |      }d}| j                  r*| j                  |      }|rd|v r|S t        d|d         }d}| j                  r*| j                  |      }|rd|v r|S t        d|d         }d}|rd|iS y NFjump_tolist[BaseMessage]r?   T)r@   r5   _moderate_tool_messagesr   r3   _moderate_user_messager    r?   workingmodifiedactions        r#   rB   z+OpenAIModerationMiddleware._moderate_inputs   s     x.""11':F&!M2F:4FG009F&!M2F:4FG((r$   c                2  K   t        |      }d}| j                  r2| j                  |       d {   }|rd|v r|S t        d|d         }d}| j                  r2| j                  |       d {   }|rd|v r|S t        d|d         }d}|rd|iS y 7 b7 &wrQ   )r@   r5   _amoderate_tool_messagesr   r3   _amoderate_user_messagerV   s        r#   rK   z,OpenAIModerationMiddleware._amoderate_inputs   s      x.""88AAF&!M2F:4FG77@@F&!M2F:4FG((% B As!   .BB=B.B/%BBc                    | j                  |t              }|y ||   }| j                  |      }|sy | j                  |      }|j                  sy | j                  ||d||      S Nr   indexr   r   r   )_find_last_indexr   _extract_text	_moderateflagged_apply_violationr    r?   last_ai_idx
ai_messagetextr   s         r#   rH   z+OpenAIModerationMiddleware._moderate_output   sz     ++Hi@k*
!!*-%~~$$Kxf % 
 	
r$   c                   K   | j                  |t              }|y ||   }| j                  |      }|sy | j                  |       d {   }|j                  sy | j                  ||d||      S 7 'wr^   )ra   r   rb   
_amoderaterd   re   rf   s         r#   rN   z,OpenAIModerationMiddleware._amoderate_output   s      ++Hi@k*
!!*-t,,~~$$Kxf % 
 	
	 -   AA3	A1
(A3c                   | j                  |t              }|y t        |      }d}t        |dz   t	        |            D ]|  }||   }t        |t              s| j                  |      }|s-| j                  |      }|j                  sK| j                  ||d||      }	|	sdd|	v r|	c S t        d|	d         }d}~ |rd|iS y 	NF   r   r_   rR   rS   r?   T)ra   r   r@   rangelen
isinstancer   rb   rc   rd   re   r   
r    r?   rg   rW   rX   idxmsgri   r   rY   s
             r#   rT   z2OpenAIModerationMiddleware._moderate_tool_messages  s     ++Hi@x.q#g,7 	 C#,Cc;/%%c*D^^D)F>>**s&$v + F &!M2F:4FG)	 , ((r$   c                  K   | j                  |t              }|y t        |      }d}t        |dz   t	        |            D ]  }||   }t        |t              s| j                  |      }|s-| j                  |       d {   }|j                  sS| j                  ||d||      }	|	sld|	v r|	c S t        d|	d         }d} |rd|iS y 7 Lwrn   )ra   r   r@   rp   rq   rr   r   rb   rk   rd   re   r   rs   s
             r#   r[   z3OpenAIModerationMiddleware._amoderate_tool_messages5  s      ++Hi@x.q#g,7 	 C#,Cc;/%%c*D??400F>>**s&$v + F &!M2F:4FG)	 , ((! 1s   B CC)C-#Cc                    | j                  |t              }|y ||   }| j                  |      }|sy | j                  |      }|j                  sy | j                  ||d||      S Nr   r_   )ra   r   rb   rc   rd   re   r    r?   rt   r!   ri   r   s         r#   rU   z1OpenAIModerationMiddleware._moderate_user_messageZ  sx     ##Hl;;3-!!'*%~~$$CwV % 
 	
r$   c                   K   | j                  |t              }|y ||   }| j                  |      }|sy | j                  |       d {   }|j                  sy | j                  ||d||      S 7 'wrx   )ra   r   rb   rk   rd   re   ry   s         r#   r\   z2OpenAIModerationMiddleware._amoderate_user_messagen  s      ##Hl;;3-!!'*t,,~~$$CwV % 
 	
	 -rl   c                  | j                  ||      }| j                  dk(  rt        ||||      | j                  dk(  rdt        |      gdS |y t	        |      }||   }t        t        |j                  d|i            ||<   d|iS )	Nerror)r   r   r   r!   r1   )r   )rR   r?   r   )updater?   )_format_violation_messager6   r   r   r@   r   r   
model_copy)	r    r?   r`   r   r   r   violation_textnew_messagesoriginals	            r#   re   z+OpenAIModerationMiddleware._apply_violation  s     77H('&	  &$9^3T2UVV=H~&",,Y4O,P
U L))r$   c                    | j                   | j                         | _         | j                   j                  j                  | j                  |      }|j
                  d   S N)r2   r   r   )r;   _build_clientmoderationscreater2   resultsr    ri   responses      r#   rc   z$OpenAIModerationMiddleware._moderate  sP    <<--/DL<<++2242P""r$   c                   K   | j                   | j                         | _         | j                   j                  j                  | j                  |       d {   }|j
                  d   S 7 wr   )r<   _build_async_clientr   r   r2   r   r   s      r#   rk   z%OpenAIModerationMiddleware._amoderate  sj     %!%!9!9!;D++77>>**D ? 
 
 ""
s   AA.A,A.c                8    t               | _        | j                  S N)r   r;   r    s    r#   r   z(OpenAIModerationMiddleware._build_client  s    x||r$   c                8    t               | _        | j                  S r   )r   r<   r   s    r#   r   z.OpenAIModerationMiddleware._build_async_client  s    (]!!!r$   c                   |j                   j                         }|j                         D cg c]  \  }}|r|j                  dd       }}}|rdj	                  |      nd}| j
                  xs t        }t        j                  |j                  j                         d      }		 |j                  ||	|      }
|
S c c}}w # t        $ r |}
Y |
S w xY w)N_ z, zOpenAI's safety policiesT)	sort_keys)
categoriescategory_scoresoriginal_content)r   
model_dumpitemsreplacejoinr7   DEFAULT_VIOLATION_TEMPLATEjsondumpsr   formatKeyError)r    r   r   categories_dictnamerd   r   category_labeltemplatescores_jsonr!   s              r#   r~   z4OpenAIModerationMiddleware._format_violation_message  s     ++668 "1!6!6!8
g LLc"

 
 &0DIIj!5O 	 ))G-Gjj!7!7!B!B!DPTU	oo) +!( & G %
   	G	s   B>(C CCc                f    t        t        |      dz
  dd      D ]  }t        ||   |      s|c S  y )Nro   )rp   rq   rr   )r    r?   message_typert   s       r#   ra   z+OpenAIModerationMiddleware._find_last_index  s=     X*B3 	C(3-6
	 r$   c                    |j                   y t        |dd       }|t        |j                         S t        |      }|r|S d S )Nri   )r   getattrr%   )r    r!   text_accessorri   s       r#   rb   z(OpenAIModerationMiddleware._extract_text  sI    ??"6 w''=!t%%r$   )r2   r   r3   boolr4   r   r5   r   r6   z"Literal['error', 'end', 'replace']r7   
str | Noner8   zOpenAI | Noner9   zAsyncOpenAI | Noner'   r(   )rD   zAgentState[Any]rE   zRuntime[Any]r'   dict[str, Any] | None)r?   Sequence[BaseMessage]r'   r   )r?   r   r`   
int | Noner   r&   r   r%   r   r   r'   r   )ri   r%   r'   r   )r'   r   )r'   r   )r   r%   r   r   r'   r%   )r?   r   r   ztype[BaseMessage]r'   r   )r!   r   r'   r   )r)   r*   r+   r,   r   r   rF   rI   rL   rO   rB   rK   rH   rN   rT   r[   rU   r\   re   rc   rk   r   r   r~   ra   rb   r-   r.   s   @r#   r0   r0   1   sq   D
 ": !#(<A(, $+/$* $* 	$*
 $* !$* :$* &$* $* )$* 
$*L eW%/$//;/	/ &/* eW%/$//;/	/ &/* eW%6$6/;6	6 &6* eW%6$6/;6	6 &6*-	6-	6
-
	
(
-
	
(#-#	#J#-#	#J
-
	
(
-
	
(*'* 	*
 * * * 
*@##".-=N	&r$   r0   )"r,   
__future__r   r   collections.abcr   typingr   r   r   r   !langchain.agents.middleware.typesr	   r
   r   langchain_core.messagesr   r   r   r   openair   r   openai.typesr   r   langgraph.runtimer   r&   r   RuntimeErrorr   r0   __all__ r$   r#   <module>r      s}    D "  $ 4 4 V V U U & 4)23 X 
L 2m&C#1E!F m&b  r$   