
    <i"                        S SK Jr  S SKJrJrJrJr  S SKJrJ	r	J
r
  S SKJrJr  S SKJr  S SKJr  S SKJrJrJr  S SKJr  S S	KJrJr  S S
KJr  S SKJr   S SKJ r   \	" S5      r#\$\$\%S4   \&\%\4   4   r' " S S\\ 5      r(g! \! a    \"r  N0f = f)    )annotations)AsyncIteratorCallableIteratorSequence)AnyTypeVarcast)UUIDuuid4)BaseCallbackHandler)BaseMessage)ChatGenerationChatGenerationChunk	LLMResult)NS_SEP)
TAG_HIDDENTAG_NOSTREAM)StreamChunk)Command)_StreamingCallbackHandlerT.c                     \ rS rSrSrSr SS.       SS jjrSS.SS	 jjrSS
 jr      SS jr	SS jr
SSSS.               SS jjrSSSS.             SS jjrSS.         SS jjrSS.         S S jjrSSSS.               S!S jjrSS.         S"S jjrSS.         S S jjrSrg)#StreamMessagesHandler   zA callback handler that implements stream_mode=messages.

Collects messages from:
(1) chat model stream events; and
(2) node outputs.
TN)	parent_nsc               T    Xl         X l        0 U l        [        5       U l        X0l        g)aC  Configure the handler to stream messages from LLMs and nodes.

Args:
    stream: A callable that takes a StreamChunk and emits it.
    subgraphs: Whether to emit messages from subgraphs.
    parent_ns: The namespace where the handler was created.
        We keep track of this namespace to allow calls to subgraphs that
        were explicitly requested as a stream with `messages` mode
        configured.

Example:
    parent_ns is used to handle scenarios where the subgraph is explicitly
    streamed with `stream_mode="messages"`.

    ```python
    def parent_graph_node():
        # This node is in the parent graph.
        async for event in some_subgraph(..., stream_mode="messages"):
            do something with event # <-- these events will be emitted
        return ...

    parent_graph.invoke(subgraphs=False)
    ```
N)stream	subgraphsmetadatasetseenr   )selfr   r   r   s       j/home/dmtnaga/Documents/work/airagagent/rag_env/lib/python3.13/site-packages/langgraph/pregel/_messages.py__init__StreamMessagesHandler.__init__(   s#    > "*,$'E	"    Fdedupec                  U(       a  UR                   U R                  ;   a  g UR                   c  [        [        5       5      Ul         U R                  R	                  UR                   5        U R                  US   SX!S   445        g )Nr   messages   )idr"   strr   addr   )r#   metamessager)   s       r$   _emitStreamMessagesHandler._emitM   s`    gjjDII-zz! \
IIMM'**%KKa*wQ.@ABr'   c                   [        U[        5      (       a  U R                  XSS9  g [        U[        5      (       a1  U H*  n[        U[        5      (       d  M  U R                  XSS9  M,     g [        U[        5      (       a  UR                  5        Hq  n[        U[        5      (       a  U R                  XSS9  M*  [        U[        5      (       d  MA  U H*  n[        U[        5      (       d  M  U R                  XSS9  M,     Ms     g [        US5      (       a  [        UR                  5      (       a  [        U5       H}  n [        X%5      n[        U[        5      (       a  U R                  XSS9  M6  [        U[        5      (       a2  U H*  n[        U[        5      (       d  M  U R                  XSS9  M,     M}  M     g g g ! [         a     M  f = f)NTr(   __dir__)
isinstancer   r2   r   dictvalueshasattrcallabler5   dirgetattrAttributeError)r#   r0   responsevalueitemkeys         r$   _find_and_emit_messages-StreamMessagesHandler._find_and_emit_messagesV   ss   h,,JJtdJ3(++!e[11JJt4J8 " $''!*e[11JJt4J8x00 %%dK88 JJt$J? !&	 + Xy))hx7G7G.H.H8}	#H2E!%55

4t
<#E844$)D)$<< $

4d
 C %* 5 % /I) & s   0G4-G%G
GGc                    U$ N r#   run_idoutputs      r$   tap_output_aiter&StreamMessagesHandler.tap_output_aiterr   s	     r'   c                    U$ rE   rF   rG   s      r$   tap_output_iter%StreamMessagesHandler.tap_output_iterw   s    r'   )parent_run_idtagsr    c                  U(       a  U(       a
  [         U;  a  [        [        [        US   5      R	                  [
        5      5      S S nU R                  (       d  [        U5      S:  a  XR                  :w  a  g U(       a5  U V	s/ s H  oR                  S5      (       a  M  U	PM     sn	=n
(       a  XS'   X4U R                  U'   g g g s  sn	f )Nlanggraph_checkpoint_nsr   zseq:steprP   )r   tupler
   r.   splitr   r   lenr   
startswithr    )r#   
serializedr+   rH   rO   rP   r    kwargsnstfiltered_tagss              r$   on_chat_model_start)StreamMessagesHandler.on_chat_model_startz   s     Tl$&>tC*C!DEKKFSTB >>c"gkbNN6J04$U1LL<TQ$UU=U'4V$%'NDMM&! '?8 %Vs   C
#C
)chunkrO   rP   c                   [        U[        5      (       d  g U R                  R                  U5      =n(       a  U R	                  XrR
                  5        g g rE   )r6   r   r    getr2   r1   )r#   tokenr_   rH   rO   rP   rY   r0   s           r$   on_llm_new_token&StreamMessagesHandler.on_llm_new_token   sD     %!455==$$V,,4,JJt]]+ -r'   )rO   c               L   U R                   R                  U5      =n(       af  UR                  (       aU  UR                  S   (       aA  UR                  S   S   n[        U[        5      (       a  U R                  XVR                  SS9  U R                   R                  US 5        g )Nr   Tr(   )r    ra   generationsr6   r   r2   r1   pop)r#   r>   rH   rO   rY   r0   gens          r$   
on_llm_end StreamMessagesHandler.on_llm_end   s~     ==$$V,,4,##(<(<Q(?**1-a0c>22JJt[[J>&$'r'   c               <    U R                   R                  US 5        g rE   r    rg   r#   errorrH   rO   rY   s        r$   on_llm_error"StreamMessagesHandler.on_llm_error        	&$'r'   c               F   U(       Ga  UR                  S5      UR                  S5      :X  Gas  U(       a  [        U;  Ga`  [        [        [        US   5      R                  [        5      5      S S nU R                  (       d  [        U5      S:  a  g X4U R                  U'   [        U[        5      (       a  UR                  5        H  u  p[        U
[        5      (       a6  U
R                  b'  U R                  R!                  U
R                  5        MN  MP  [        U
["        5      (       d  Mg  [        U
[        5      (       a  M~  U
 HN  n[        U[        5      (       d  M  UR                  c  M)  U R                  R!                  UR                  5        MP     M     g g g g g )Nnamelanggraph_noderR   rS   r   )ra   r   rT   r
   r.   rU   r   r   rV   r    r6   r7   itemsr   r-   r"   r/   r   )r#   rX   inputsrH   rO   rP   r    rY   rZ   rA   r?   r@   s               r$   on_chain_start$StreamMessagesHandler.on_chain_start   s0    

6"hll3C&DDZt3tC*C!DEKKFSTB >>c"gk%'NDMM&!&$''"(,,.JC!%55 88/ IIMM%((3 0#E844Zs=S=S$)D)$<<#'77#6$(IIMM$''$: %* #1 ( 4 E r'   c                  U R                   R                  US 5      =n(       a  [        U[        5      (       a  U R	                  XQR
                  5        g [        U[        5      (       ad  [        S U 5       5      (       aM  U HF  n[        U[        5      (       a  U R	                  XVR
                  5        M5  U R	                  XV5        MH     g U R	                  XQ5        g g )Nc              3  B   #    U  H  n[        U[        5      v   M     g 7frE   )r6   r   ).0r?   s     r$   	<genexpr>5StreamMessagesHandler.on_chain_end.<locals>.<genexpr>   s      88@u
5'**s   )r    rg   r6   r   rB   updater   any)r#   r>   rH   rO   rY   r0   r?   s          r$   on_chain_end"StreamMessagesHandler.on_chain_end   s     ==$$VT2242(G,,,,T??CHh//C 88@8 5 5 &E!%1144T<<H44TA	 & ,,T< 3r'   c               <    U R                   R                  US 5        g rE   rl   rm   s        r$   on_chain_error$StreamMessagesHandler.on_chain_error   rq   r'   )r    r   r"   r   r   )r   zCallable[[StreamChunk], None]r   boolr   ztuple[str, ...] | NonereturnNone)r0   Metar1   r   r)   r   r   r   )r0   r   r>   r   r   r   )rH   r   rI   AsyncIterator[T]r   r   )rH   r   rI   Iterator[T]r   r   )rX   dict[str, Any]r+   zlist[list[BaseMessage]]rH   r   rO   UUID | NonerP   list[str] | Noner    dict[str, Any] | NonerY   r   r   r   )rb   r.   r_   zChatGenerationChunk | NonerH   r   rO   r   rP   r   rY   r   r   r   )
r>   r   rH   r   rO   r   rY   r   r   r   )
rn   BaseExceptionrH   r   rO   r   rY   r   r   r   )rX   r   rv   r   rH   r   rO   r   rP   r   r    r   rY   r   r   r   )
r>   r   rH   r   rO   r   rY   r   r   r   )__name__
__module____qualname____firstlineno____doc__
run_inliner%   r2   rB   rJ   rM   r]   rc   ri   ro   rw   r   r   __static_attributes__rF   r'   r$   r   r      st    JX -1##-## ##
 *## 
##J IN C8$4	
 &*!%*.3"3 *3
 3 #3 3 (3 3 
34 -1%)!%,, *	,
 , #, , , 
,( &*(( 	(
 #( ( 
(( &*(( 	(
 #( ( 
(  &*!%*.;"; ;
 ; #; ; (; ; 
;L &*== 	=
 #= = 
=< &*(( 	(
 #( ( 
( (r'   r   N))
__future__r   collections.abcr   r   r   r   typingr   r	   r
   uuidr   r   langchain_core.callbacksr   langchain_core.messagesr   langchain_core.outputsr   r   r   langgraph._internal._constantsr   langgraph.constantsr   r   langgraph.pregel.protocolr   langgraph.typesr   !langchain_core.tracers._streamingr   ImportErrorobjectr   rT   r.   r7   r   r   rF   r'   r$   <module>r      s    " G G 
  8 / Q Q 1 8 1 #'K CLU38_d38n,-](/1J ](  ' &'s   B BB