
    A'h?                      d dl mZ d dlZd dlZd dlZd dlZd dlZd dlmZm	Z	 d dl
mZmZmZmZ d dlmZ d dlmZmZmZmZ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 m!Z! d dl"m#Z#m$Z$m%Z% d dl&m'Z' d dl(m)Z)m*Z* d dl+m,Z, d dl-m.Z. d dl/m0Z0 d dl1m2Z2 d dl3m4Z4m5Z5m6Z6m7Z7 d dl8m9Z9m:Z:m;Z;m<Z<m=Z=m>Z>m?Z?m@Z@mAZAmBZBmCZCmDZDmEZEmFZFmGZGmHZHmIZImJZJmKZKmLZLmMZMmNZNmOZOmPZPmQZQmRZR d dlSmTZTmUZUmVZVmWZW d dlXmYZY d dlZm[Z[m\Z\m]Z]m^Z^m_Z_ d dl`maZa d dlbmcZcmdZd d dlemfZf d dlgmhZh d dlimjZjmkZk d dllmmZmmnZnmoZo d dlpmqZqmrZr d dlsmtZt d dlumvZv d d lwmxZx d d!lymzZz d d"l{m|Z| d d#l}m~Z~ d d$lmZmZ d d%lmZmZ d d&lmZ d d'lmZmZmZmZmZmZmZmZmZ d d(lmZmZmZmZmZmZ d d)lmZ d d*lmZmZ d d+lmZmZ 	 d d,lmZ eee ge!f   ef   Z G d- d.      Z G d/ d0ev      Zd2d1Zy# e$ r dZY 0w xY w)3    )annotationsN)defaultdictdeque)AsyncIteratorIteratorMappingSequence)partial)AnyCallableUnioncastget_type_hintsoverload)UUIDuuid5)	get_debug)RunnableSequence)InputOutput)RunnableConfig%get_async_callback_manager_for_configget_callback_manager_for_config)Graph)ConfigurableFieldSpecget_unique_config_specs)	BaseModel)Self)	BaseCache)BaseChannel)BaseCheckpointSaver
CheckpointCheckpointTuplecopy_checkpoint)CACHE_NS_WRITESCONFCONFIG_KEY_CACHECONFIG_KEY_CHECKPOINT_DURINGCONFIG_KEY_CHECKPOINT_IDCONFIG_KEY_CHECKPOINT_NSCONFIG_KEY_CHECKPOINTERCONFIG_KEY_NODE_FINISHEDCONFIG_KEY_READCONFIG_KEY_RESUMINGCONFIG_KEY_RUNNER_SUBMITCONFIG_KEY_SENDCONFIG_KEY_STORECONFIG_KEY_STREAMCONFIG_KEY_STREAM_WRITERCONFIG_KEY_TASK_IDCONFIG_KEY_THREAD_IDENDERRORINPUT	INTERRUPTNS_ENDNS_SEPNULL_TASK_IDPUSH	SCHEDULED)	ErrorCodeGraphRecursionErrorInvalidUpdateErrorcreate_error_message)ManagedValueSpec)PregelTaskWritesapply_writes
local_readlocal_writeprepare_next_tasks)
identifier)create_checkpointempty_checkpoint)tasks_w_writes)
draw_graph)	map_inputread_channels)AsyncPregelLoopStreamProtocolSyncPregelLoop)AsyncChannelsManagerChannelsManager)StreamMessagesHandler)PregelProtocol)
PregelNode)RetryPolicy)PregelRunner)get_new_channel_versions)validate_graphvalidate_keys)ChannelWriteChannelWriteEntry)	BaseStore)	AllCachePolicyCheckpointer	InterruptLoopProtocolStateSnapshotStateUpdateStreamChunk
StreamMode)ensure_configmerge_configspatch_checkpoint_mappatch_configpatch_configurablerecast_checkpoint_ns)get_enhanced_type_hints)create_modelis_supported_by_pydantic)
AsyncQueue	SyncQueue)_StreamingCallbackHandlerc                      e Zd Zeeddd	 	 	 	 	 	 	 dd              Zeeddd	 	 	 	 	 	 	 dd              Zeddd	 	 	 	 	 	 	 d	d       Ze	 	 	 	 	 	 d
d       Zy)ChannelN)keytagsc                    y N clschannelsrw   rx   s       `/home/kushmeetdev/Regenta/Chatbot/venv/lib/python3.12/site-packages/langgraph/pregel/__init__.pysubscribe_tozChannel.subscribe_to            c                    y rz   r{   r|   s       r   r   zChannel.subscribe_to   r   r   c          
     b   t        |t              s|t        d      t        t	        t
        t        t           t        t        t        f   f   t        |t              r|||in#t        |t              r|gn|D ci c]  }|| c}      t        |t              r	|g|      S ||      S c c}w )ziRuns process.invoke() each time channels are updated,
        with a dict of the channel values as input.z9Can't specify a key when subscribing to multiple channels)r~   triggersrx   )
isinstancestr
ValueErrorrW   r   r   listr   )r}   r~   rw   rx   chans        r   r   zChannel.subscribe_to   s     (C(S_K  d3ic!223 "(C0S_ (O &h4 "
5=>TdDj> $.h#<hZ
 	
 CK
 	
 ?s   ;
B,c                    t        |D cg c]  }t        |       c}|j                         D cg c]*  \  }}t        |      rt        ||      nt        ||      , c}}z         S c c}w c c}}w )zEWrites to channels the result of the lambda, or None to skip writing.)mapper)value)r]   r^   itemscallable)r}   r~   kwargsckvs         r   write_tozChannel.write_to   su     +34aq!4 #LLN Aq  { &a2*1A67

 
	
4s
   A*/A/)r~   r   rw   
str | Nonerx   list[str] | NonereturnrW   )r~   Sequence[str]rw   Nonerx   r   r   rW   )r~   str | Sequence[str]rw   r   rx   r   r   rW   )r~   r   r   
WriteValuer   r]   )__name__
__module____qualname__r   classmethodr   r   r{   r   r   rv   rv      s   
 !% 	
  
   
 !% 	
  
   
 !%
%
 	

 
 

 
< 

 
 
	
 
r   rv   c                  v    e Zd ZU dZded<   ded<   dZded<   	 d	Zd
ed<   	 ded<   dZded<   	 ded<   ded<   ded<   dZded<   	 d
ed<   	 dZ	ded<   	 dZ
ded<   	 dZded<   	 dZded <   	 dZd!ed"<   	 dZd#ed$<   dZd%ed&<   dZd'ed(<   d)Zd*ed+<   d,ed-<   d.dd	ddddddddddddddd)d/	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d]d0Z	 d^d	d1	 	 	 	 	 d_d2Z	 d^d	d1	 	 	 	 	 d_d3Zd`d4Zd^dad5Zd^dbd6Zdcd7Zeddd8       Zdd9de fd:Zdd9	 	 	 dfd;Zedgd<       Zd^dh fd=Z	 d^	 	 	 did>Zedgd?       Z 	 d^	 	 	 dh fd@Z!	 d^	 	 	 didAZ"edjdB       Z#edkdC       Z$dd	dD	 	 	 	 	 dldEZ%dd	dD	 	 	 	 	 dmdFZ&dndGZ'	 	 do	 	 	 	 	 	 	 	 	 dpdHZ(	 	 do	 	 	 	 	 	 	 	 	 dpdIZ)d	dJ	 	 	 	 	 dqdKZ*d	dJ	 	 	 	 	 dqdLZ+ddddM	 	 	 	 	 	 	 	 	 drdNZ,ddddM	 	 	 	 	 	 	 	 	 dsdOZ-	 	 	 	 	 	 dtdPZ.	 	 	 	 	 	 dtdQZ/	 d^	 	 	 	 	 	 	 dudRZ0	 d^	 	 	 	 	 	 	 dvdSZ1	 	 	 	 	 	 	 	 	 	 	 	 	 	 dwdTZ2	 d^ddddddd	dU	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 dxdVZ3	 d^ddddddd	dU	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 dydWZ4	 d^dddddddX	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 dzdYZ5	 d^dddddddX	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 dzdZZ6d^d{d[Z7d^d{d\Z8 xZ9S )|Pregeluh   Pregel manages the runtime behavior for LangGraph applications.

    ## Overview

    Pregel combines [**actors**](https://en.wikipedia.org/wiki/Actor_model)
    and **channels** into a single application.
    **Actors** read data from channels and write data to channels.
    Pregel organizes the execution of the application into multiple steps,
    following the **Pregel Algorithm**/**Bulk Synchronous Parallel** model.

    Each step consists of three phases:

    - **Plan**: Determine which **actors** to execute in this step. For example,
        in the first step, select the **actors** that subscribe to the special
        **input** channels; in subsequent steps,
        select the **actors** that subscribe to channels updated in the previous step.
    - **Execution**: Execute all selected **actors** in parallel,
        until all complete, or one fails, or a timeout is reached. During this
        phase, channel updates are invisible to actors until the next step.
    - **Update**: Update the channels with the values written by the **actors**
        in this step.

    Repeat until no **actors** are selected for execution, or a maximum number of
    steps is reached.

    ## Actors

    An **actor** is a `PregelNode`.
    It subscribes to channels, reads data from them, and writes data to them.
    It can be thought of as an **actor** in the Pregel algorithm.
    `PregelNodes` implement LangChain's
    Runnable interface.

    ## Channels

    Channels are used to communicate between actors (`PregelNodes`).
    Each channel has a value type, an update type, and an update function – which
    takes a sequence of updates and
    modifies the stored value. Channels can be used to send data from one chain to
    another, or to send data from a chain to itself in a future step. LangGraph
    provides a number of built-in channels:

    ### Basic channels: LastValue and Topic

    - `LastValue`: The default channel, stores the last value sent to the channel,
       useful for input and output values, or for sending data from one step to the next
    - `Topic`: A configurable PubSub Topic, useful for sending multiple values
       between *actors*, or for accumulating output. Can be configured to deduplicate
       values, and/or to accumulate values over the course of multiple steps.

    ### Advanced channels: Context and BinaryOperatorAggregate

    - `Context`: exposes the value of a context manager, managing its lifecycle.
      Useful for accessing external resources that require setup and/or teardown. eg.
      `client = Context(httpx.Client)`
    - `BinaryOperatorAggregate`: stores a persistent value, updated by applying
       a binary operator to the current value and each update
       sent to the channel, useful for computing aggregates over multiple steps. eg.
      `total = BinaryOperatorAggregate(int, operator.add)`

    ## Examples

    Most users will interact with Pregel via a
    [StateGraph (Graph API)][langgraph.graph.StateGraph] or via an
    [entrypoint (Functional API)][langgraph.func.entrypoint].

    However, for **advanced** use cases, Pregel can be used directly. If you're
    not sure whether you need to use Pregel directly, then the answer is probably no
    – you should use the Graph API or Functional API instead. These are higher-level
    interfaces that will compile down to Pregel under the hood.

    Here are some examples to give you a sense of how it works:

    Example: Single node application
        ```python
        from langgraph.channels import EphemeralValue
        from langgraph.pregel import Pregel, Channel, ChannelWriteEntry

        node1 = (
            Channel.subscribe_to("a")
            | (lambda x: x + x)
            | Channel.write_to("b")
        )

        app = Pregel(
            nodes={"node1": node1},
            channels={
                "a": EphemeralValue(str),
                "b": EphemeralValue(str),
            },
            input_channels=["a"],
            output_channels=["b"],
        )

        app.invoke({"a": "foo"})
        ```

        ```con
        {'b': 'foofoo'}
        ```

    Example: Using multiple nodes and multiple output channels
        ```python
        from langgraph.channels import LastValue, EphemeralValue
        from langgraph.pregel import Pregel, Channel, ChannelWriteEntry

        node1 = (
            Channel.subscribe_to("a")
            | (lambda x: x + x)
            | Channel.write_to("b")
        )

        node2 = (
            Channel.subscribe_to("b")
            | (lambda x: x + x)
            | Channel.write_to("c")
        )


        app = Pregel(
            nodes={"node1": node1, "node2": node2},
            channels={
                "a": EphemeralValue(str),
                "b": LastValue(str),
                "c": EphemeralValue(str),
            },
            input_channels=["a"],
            output_channels=["b", "c"],
        )

        app.invoke({"a": "foo"})
        ```

        ```con
        {'b': 'foofoo', 'c': 'foofoofoofoo'}
        ```

    Example: Using a Topic channel
        ```python
        from langgraph.channels import LastValue, EphemeralValue, Topic
        from langgraph.pregel import Pregel, Channel, ChannelWriteEntry

        node1 = (
            Channel.subscribe_to("a")
            | (lambda x: x + x)
            | {
                "b": Channel.write_to("b"),
                "c": Channel.write_to("c")
            }
        )

        node2 = (
            Channel.subscribe_to("b")
            | (lambda x: x + x)
            | {
                "c": Channel.write_to("c"),
            }
        )


        app = Pregel(
            nodes={"node1": node1, "node2": node2},
            channels={
                "a": EphemeralValue(str),
                "b": EphemeralValue(str),
                "c": Topic(str, accumulate=True),
            },
            input_channels=["a"],
            output_channels=["c"],
        )

        app.invoke({"a": "foo"})
        ```

        ```pycon
        {'c': ['foofoo', 'foofoofoofoo']}
        ```

    Example: Using a BinaryOperatorAggregate channel
        ```python
        from langgraph.channels import EphemeralValue, BinaryOperatorAggregate
        from langgraph.pregel import Pregel, Channel


        node1 = (
            Channel.subscribe_to("a")
            | (lambda x: x + x)
            | {
                "b": Channel.write_to("b"),
                "c": Channel.write_to("c")
            }
        )

        node2 = (
            Channel.subscribe_to("b")
            | (lambda x: x + x)
            | {
                "c": Channel.write_to("c"),
            }
        )


        def reducer(current, update):
            if current:
                return current + " | " + "update"
            else:
                return update

        app = Pregel(
            nodes={"node1": node1, "node2": node2},
            channels={
                "a": EphemeralValue(str),
                "b": EphemeralValue(str),
                "c": BinaryOperatorAggregate(str, operator=reducer),
            },
            input_channels=["a"],
            output_channels=["c"]
        )

        app.invoke({"a": "foo"})
        ```

        ```con
        {'c': 'foofoo | foofoofoofoo'}
        ```

    Example: Introducing a cycle
        This example demonstrates how to introduce a cycle in the graph, by having
        a chain write to a channel it subscribes to. Execution will continue
        until a None value is written to the channel.

        ```python
        from langgraph.channels import EphemeralValue
        from langgraph.pregel import Pregel, Channel, ChannelWrite, ChannelWriteEntry

        example_node = (
            Channel.subscribe_to("value")
            | (lambda x: x + x if len(x) < 10 else None)
            | ChannelWrite(writes=[ChannelWriteEntry(channel="value", skip_none=True)])
        )

        app = Pregel(
            nodes={"example_node": example_node},
            channels={
                "value": EphemeralValue(str),
            },
            input_channels=["value"],
            output_channels=["value"]
        )

        app.invoke({"value": "a"})
        ```

        ```con
        {'value': 'aaaaaaaaaaaaaaaa'}
        ```
    dict[str, PregelNode]nodesz)dict[str, BaseChannel | ManagedValueSpec]r~   valuesrh   stream_modeFboolstream_eagerr   output_channelsNstr | Sequence[str] | Nonestream_channelsAll | Sequence[str]interrupt_after_nodesinterrupt_before_nodesinput_channelsfloat | Nonestep_timeoutdebugrb   checkpointerBaseStore | NonestoreBaseCache | Nonecacher{   zSequence[RetryPolicy]retry_policyCachePolicy | Nonecache_policytype[Any] | Noneconfig_typetype[BaseModel] | Noneinput_modelRunnableConfig | Noneconfig	LangGraphr   nameMapping[str, Sequence[str]]trigger_to_nodesT)auto_validater   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   c                  || _         |xs i | _        || _        || _        || _        || _        || _        |	| _        |
| _        || _	        ||n	t               | _        || _        || _        || _        t        |t               r|fn|| _        || _        || _        || _        || _        |xs i | _        || _        |r| j1                          y y rz   )r   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   rX   r   r   r   r   r   r   r   validate)selfr   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   s                          r   __init__zPregel.__init__  s    4 
 B&(..%:"&<#,(#/UY[
(

),D\O, 	 )&& 0 6B	MMO r   xrayc                  |rO| j                         D ci c]3  \  }}||j                  |t        |t              s|dk  r|n|dz
        5 }}}ni }t	        t        | j                  |      | j                  | j                  | j                  | j                  | j                  | j                  | j                  |	      S c c}}w ):Return a drawable representation of the computation graph.r      r   r   specsr   r   r   r   r   	subgraphs)get_subgraphs	get_graphr   r   rM   rj   r   r   r~   r   r   r   r   r   )r   r   r   r   r   r   s         r   r   zPregel.get_graph?  s    
  !..0
 Aq	 1;;!+D$!7419$QR(   I  I$++v.**--.."&"<"<#'#>#>!22**

 
	
s   8Cc                 K   r%| j                         2 ci c3 d{   \  }}||i }t        t        | j                        | j                  | j                  | j                  | j                  | j                  | j                  | j                  |	      S 7 {6 nc c}}w }}}t        |t        j                  fd|j	                         D          d{  7        D ci c]  \  }}||
 nc c}}w }}}ӭw)r   Nc              3  x   K   | ]1  }|j                  t        t              sd k  rndz
         3 yw)r   r   r   N)
aget_graphr   r   ).0pr   r   s     r   	<genexpr>z$Pregel.aget_graph.<locals>.<genexpr>j  sH       !" LL &#-dD#9TQY &*%)AX	 ) s   7:r   )aget_subgraphszipasynciogatherr   rM   rj   r   r   r~   r   r   r   r   r   )r   r   r   r   r   
subpregelsr   s    ``    r   r   zPregel.aget_graph[  s	     '+':':'<5 5#q!1( I$++v.**--.."&"<"<#'#>#>!22**

 
	
/5 5 5J 5
  !.. &0%6%6%8
 
 
Aq 1 I sN   DB BBBB A0DBB =DC
D)C76Dc                V    t        |       | j                         j                         dS )z0Mime bundle used by Jupyter to display the graph)z
text/plainz	image/png)reprr   draw_mermaid_png)r   r   s     r   _repr_mimebundle_zPregel._repr_mimebundle_  s(     t*)::<
 	
r   c                N    i | j                   |xs i } | j                  di |S )Nr{   )__dict__	__class__)r   updateattrss      r   copyzPregel.copy  s-    34==3V\r3t~~&&&r   c                n    | j                  dt        | j                  |t        t        |            i      S )z:Create a copy of the Pregel object with an updated config.r   )r   rj   r   r   r   )r   r   r   s      r   with_configzPregel.with_config  s/    yy}T[[&$~v:VWX
 	
r   c                Z   t        | j                  | j                  j                         D ci c]  \  }}t	        |t
              s|| c}}| j                  | j                  | j                  | j                  | j                         t        | j                        | _        | S c c}}w rz   )r[   r   r~   r   r   r    r   r   r   r   r   _trigger_to_nodesr   )r   r   r   s      r   r   zPregel.validate  s    JJ"mm113Rdaz![7QQTR    &&''	
 !2$** = Ss   B'B'c                
   t        | j                  j                         D cg c]  }|j                  D ]  }|  c}}t	        | j
                  t              r| j
                  j                  ng z   | j                  ;t        | j                        D cg c]  \  }}}}t        ||||       c}}}}ng z         D cg c]&  }|j                  t        t        t        t        fvr|( c}S c c}}w c c}}}}w c c}w )N)id
annotationdefaultdescription)r   r   r   config_specsr   r   r!   r   ro   r   r   r-   r0   r+   r.   )r   nodespecr   typr   r   s          r   r   zPregel.config_specs  s    0"&**"3"3"5T$$BSBST$TT "$"3"35HI %%22	$ ''3	 @W ,,@
 
 <D#w .#'*$+(3	
 '!
0 ww'#	3 !
 !	
 U
!
s   C2
C8+D includec               |   | j                   t        | j                         st        |   |      S |xs g }d| j                   d fit	        t
              j                         D ci c]"  \  }}||D cg c]
  }|dk7  s	| c}v r||d f$ c}}}}t        | j                  d      |      S c c}w c c}}}w )Nr   configurableConfigfield_definitions)	r   rq   superconfig_schemar   r   r   rp   get_name)r   r   
field_name
field_typeifieldsr   s         r   r   zPregel.config_schema  s    
 #+CDDTDT+U7((99-RT--t4
 /=^.L.R.R.T *J
W!L^8K!!LL Z..
 DMM(3vNN "Ms   )B7
6
B2B2B7
2B7
c               ~    | j                  |      }t        |d      r|j                         S |j                         S )Nr   model_json_schema)r   hasattrr  schema)r   r   r  s      r   get_config_jsonschemazPregel.get_config_jsonschema  s=     ##G#46./++--==?"r   c                    t        | j                  t              r6| j                  | j                     }t        |t              r|j
                  S y y rz   )r   r   r   r~   r    
UpdateTyper   channels     r   	InputTypezPregel.InputType  sF    d))3/mmD$7$78G';/))) 0 0r   c           
        | j                   | j                   S t        | j                  |      }t        | j                  t
              rt        |   |      S t        | j                  d      | j                  xs | j                  j                         D ci c]2  }| j                  |   x}rt        |t              r||j                  d f4 c}      S c c}w )Nr   r   )r   rj   r   r   r   r   r   get_input_schemarp   r   r~   keysr    r	  r   r   r   r   r   s       r   r  zPregel.get_input_schema  s    '###t{{F3d))3/7+F33g& "00HDMM4F4F4H#!]]1---:a3M d++# #s   7C
c                |    | j                  |      }t        |d      r|j                         S |j                         S Nr  )r  r  r  r  r   r   r  s      r   get_input_jsonschemazPregel.get_input_jsonschema  s:     &&v.6./++--==?"r   c                    t        | j                  t              r6| j                  | j                     }t        |t              r|j
                  S y y rz   )r   r   r   r~   r    	ValueTyper
  s     r   
OutputTypezPregel.OutputType  sF    d**C0mmD$8$89G';/((( 0 1r   c           
     P   t        | j                  |      }t        | j                  t              rt
        |   |      S t        | j                  d      | j                  D ci c]2  }| j                  |   x}rt        |t              r||j                  d f4 c}      S c c}w )Nr   r   )rj   r   r   r   r   r   get_output_schemarp   r   r~   r    r  r  s       r   r  zPregel.get_output_schema  s     t{{F3d**C07,V44h' "11#!]]1---:a3M T**# #s   $7B#
c                |    | j                  |      }t        |d      r|j                         S |j                         S r  )r  r  r  r  r  s      r   get_output_jsonschemazPregel.get_output_jsonschema  s:     ''/6./++--==?"r   c                D    | j                   }t        |t              r|gS |S rz   )stream_channels_asisr   r   )r   r   s     r   stream_channels_listzPregel.stream_channels_list$  s)    33!+OS!A_	
GV	
r   c                    | j                   xs8 | j                  D cg c]"  }t        | j                  |   t              s!|$ c}S c c}w rz   )r   r~   r   r    )r   r   s     r   r  zPregel.stream_channels_asis+  sC    ## 
}}(

4==3C[(QA(
 	
 (
s   "A A	namespacerecursec             #  ~  K   | j                   j                         D ]  \  }||j                        s|j                  r|j                  d   nd}|s:|k(  r|f  y||f |sRt	        |t
              sc||t              dz   d }fd|j                  ||      D        E d{     y7 w)aq  Get the subgraphs of the graph.

        Args:
            namespace: The namespace to filter the subgraphs by.
            recurse: Whether to recurse into the subgraphs.
                If False, only the immediate subgraphs will be returned.

        Returns:
            Iterator[tuple[str, PregelProtocol]]: An iterator of the (namespace, subgraph) pairs.
        Nr   r   c              3  @   K   | ]  \  }} t          | |f  y wrz   )r;   )r   nsr   s      r   r   z'Pregel.get_subgraphs.<locals>.<genexpr>Q  s/        Aq !6&!-q1 s   r   )r   r   
startswithr   r   r   lenr   )r   r!  r"  r   graphr   s        @r   r   zPregel.get_subgraphs1  s      ****, 	JD$$ ++D1 *.DNN1%TE 9$+%$+%z%8 ,$-c$i!mo$>	 $)$7$7&/ %8 %   '	&s$   AB=B=0B=2B=3B;4B=c              P   K   | j                  ||      D ]  \  }}||f  yw)av  Get the subgraphs of the graph.

        Args:
            namespace: The namespace to filter the subgraphs by.
            recurse: Whether to recurse into the subgraphs.
                If False, only the immediate subgraphs will be returned.

        Returns:
            AsyncIterator[tuple[str, PregelProtocol]]: An iterator of the (namespace, subgraph) pairs.
        r   N)r   )r   r!  r"  r   r   s        r   r   zPregel.aget_subgraphsX  s5      ,,y',R 	JD$*	s   $&c                     y)z1Migrate a saved checkpoint to new channel layout.Nr{   )r   
checkpoints     r   _migrate_checkpointzPregel._migrate_checkpointh  s    r   c                   |st        i d|d d d dd      S | j                  |j                         t        | j                  |j                  t        |j                  |j                  j                  dd      dz   |j                  j                  dd      dz         d	      5 \  }}t        |j                  |j                  xs g | j                  |||j                  |j                  j                  dd      dz   d| j                  t        | j                  t              r| j                  nd d 
      }t!        | j#                               }|j                  t$           j                  t&        d      }	i }
|j)                         D ]  }|j*                  |vr|j*                   t,         |j.                   }|	r|	 t0         | }|s4t$        d|j                  t$           d   t&        |ii}||
|j.                  <   vt$        t2        |d|j                  t$           d   t&        |ii}||j*                     j5                  |d      |
|j.                  <    |j                  xs g D cg c]  }|d   t6        k(  s|dd   c}x}r4t9        |j                  |t;        dt<        |g       gd | j>                         |r|j                  r|j                  D ]?  \  }}}|t@        tB        tD        fv r||vr ||   jF                  jI                  ||f       A |j)                         D cg c]  }|jF                  s| c}x}r#t9        |j                  ||d | j>                         tK        |j)                         |j                  |
| jL                        }t        tO        || jL                        tQ        d |j)                         D              tS        |j                  |j                        |j                  |j                  d   tS        |jT                  |j                        |tQ        |D cg c]  }|jV                  D ]  }|  c}}            cd d d        S c c}w c c}w c c}}w # 1 sw Y   y xY w)Nr{   r   nextr   metadata
created_atparent_configtasks
interruptsstepr      r   r6  stopTskip_contextfor_executionr   r   manager 	thread_idr   r   c              3  N   K   | ]  }|j                   r|j                    y wrz   writesr   r   ts     r   r   z1Pregel._prepare_state_snapshot.<locals>.<genexpr>       JaffJ   %%ts),re   r-  r,  rT   r~   rd   r   r1  getrH   pending_writesr   r   r   r   r!   dictr   r&   r*   r   r   r:   r   r;   r+   	get_stater<   rE   rD   r8   r   r7   r9   r>   rE  appendrL   r  rO   tuplerk   r3  r5  )r   r   savedr"  apply_pending_writesr~   managed
next_tasksr   	parent_nstask_statestasktask_nswnull_writestidr   r   rG  r4  tasks_with_writesr  s                         r   _prepare_state_snapshotzPregel._prepare_state_snapshotl  sU     "	 	 	  !1!12MM||^^''3a7^^''3a7
 	
 b	 !h+  $$*

""62.2"jj "$"3"35HI %%J$ T//12IT*../GLIEGK"))+ 99I-!YYKxy9!*F8G9=G 'd);K)H4gF ,2K( 3W'd);K)H4gF ,5TYY+?+I+I$ ,J ,K(3<  %339rQqT\=Q!" {  $$%b%bAB)) $(<(<!&!5!5 :ICAUIy99 *, sO**111a&9: )3(9(9(;H1qxxQHH5H (((E4AVAV !/!!#$$))	! !h(A(ABJj&7&7&9JJ$U\\5>>B  &$U%8%8%..I!%6PTP1qPqPQ	sb	 b	z" I& QCb	 b	sE   $GQ2*Q"<Q"B'Q2*Q'<Q' C3Q23Q,Q2"Q22Q;c                	  K   |st        i d|d d d dd      S | j                  |j                         t        | j                  |j                  t        |j                  |j                  j                  dd      dz   |j                  j                  dd      dz         d	      4 d {   \  }}t        |j                  |j                  xs g | j                  |||j                  |j                  j                  dd      dz   d| j                  t        | j                  t              r| j                  nd d 
      }| j!                         2 	ci c3 d {   \  }}	||	7 7 6 nc c}	}w }
}}	|j                  t"           j                  t$        d      }i }|j'                         D ]  }|j(                  |
vr|j(                   t*         |j,                   }|r| t.         | }|s4t"        d|j                  t"           d   t$        |ii}|||j,                  <   vt"        t0        |d|j                  t"           d   t$        |ii}|
|j(                     j3                  |d       d {  7  ||j,                  <    |j                  xs g D cg c]  }|d   t4        k(  s|dd   nc c}w c}x}r4t7        |j                  |t9        dt:        |g       gd | j<                         |r|j                  r|j                  D ]?  \  }}}|t>        t@        tB        fv r||vr ||   jD                  jG                  ||f       A |j'                         D cg c]  }|jD                  s| nc c}w c}x}r#t7        |j                  ||d | j<                         tI        |j'                         |j                  || jJ                        }t        tM        || jJ                        tO        d |j'                         D              tQ        |j                  |j                        |j                  |j                  d   tQ        |jR                  |j                        |tO        |D cg c]  }|jT                  D ]  }|  nc c}}w c}}            cd d d       d {  7   S # 1 d {  7  sw Y   y xY ww)Nr{   r/  r6  r7  r   r8  r9  Tr;  r=  r@  rA  rB  r   c              3  N   K   | ]  }|j                   r|j                    y wrz   rD  rF  s     r   r   z2Pregel._aprepare_state_snapshot.<locals>.<genexpr>[  rH  rI  rJ  )+re   r-  r,  rS   r~   rd   r   r1  rK  rH   rL  r   r   r   r   r!   r   r&   r*   r   r   r:   r   r;   r+   
aget_stater<   rE   rD   r8   r   r7   r9   r>   rE  rO  rL   r  rO   rP  rk   r3  r5  )r   r   rQ  r"  rR  r~   rS  rT  r%  gr   rU  rV  rW  rX  rY  rZ  r[  r   r   rG  r4  r\  r  s                           r   _aprepare_state_snapshotzPregel._aprepare_state_snapshot  s      "	 	 	  !1!12'MM||^^''3a7^^''3a7
 	
 f	 f	
 ,  $$*

""62.2"jj "$"3"35HI %%J$ 150C0C0EFF1AAf	@ GFFIFT*../GLIEGK"))+ 99I-!YYKxy9!*F8G9=G 'd);K)H4gF ,2K( 3W'd);K)H4gF 2;4991E1P1P$ 2Q 2 , ,K(3<  %339rQqT\=Q!"  {  $$%b%bAB)) $(<(<!&!5!5 :ICAUIy99 *, sO**111a&9: )3(9(9(;H1qxxQHHH5H (((E4AVAV !/!!#$$))	! !h(A(ABJj&7&7&9JJ$U\\5>>B  &$U%8%8%..I!%6PTP1qPqPPQ	{f	 f	 f	 f	 f	s   B(S*E&+S.B&R,E,E*
E(E*
E,&S(E*
*E,+DR,1I42&R,J2*J21B-R,M50M54C9R,-RR,S%R(&S,R>2R53R>:SrB  c               L   t        |      t           j                  t        | j                        }|st        d      |t           j                  t        d      x}rht        |t           vrYt        |      }| j                  |d      D ])  \  }}|j                  t        |t        |i      |      c S  t        d| d      | j                  rt        | j                  |      n|}| j                  du r?t        t        |t           t                 }t        |t        t        t        |      ii      }|t           t           }	t!        |	t              st        |	      |t           t        <   |j#                  |      }
| j%                  ||
|r|ndt&        |t           v	      S )
#Get the current state of the graph.No checkpointer setr@  Tr   rB  	Subgraph 
 not foundNr"  rR  )ri   r&   rK  r+   r   r   r*   rn   r   rN  rm   r   rj   r   r   r5   r   	get_tupler]  r)   r   r   r   r   checkpoint_nsrecast_pregelnsrA  rQ  s              r   rN  zPregel.get_stated  s    4A3H3N3R3R#T%6%64
 233 $D\--.FKKMK%VD\9)-8F!//&$/O A	6''&v0G/VW' (  A !9VHJ!?@@7;{{t{{F3$c6$<(@ABB" 8:Nr:RSTF 4L!56	)S)14YF4L-.&&v.++$-L4!9!M	 , 
 	
r   c                 K   t        |      t           j                  t        | j                        }|st        d      |t           j                  t        d      x}rft        |t           vrWt        |      }| j                  |d      2 3 d{   \  }}|j                  t        |t        |i      |       d{   c S | j                  rt        | j                  |      n|}| j                  du r?t        t        |t           t                 }t        |t        t        t        |      ii      }|t           t           }	t!        |	t              st        |	      |t           t        <   |j#                  |       d{   }
| j%                  ||
|r|ndt&        |t           v	       d{   S 7 &7 6 t        d| d      7 G7 w)
rd  re  r@  Tr   NrB  rf  rg  rh  )ri   r&   rK  r+   r   r   r*   rn   r   r`  rm   r   rj   r   r   r5   r   
aget_tuplerb  r)   rj  s              r   r`  zPregel.aget_state  s     4A3H3N3R3R#T%6%64
 233 $D\--.FKKMK%VD\9)-8F#'#6#6QU#6#V A Aia#..&v0G/VW' /    8<{{t{{F3$c6$<(@ABB" 8:Nr:RSTF 4L!56	)S)14YF4L-."--f5522$-L4!9!M	 3 
 
 	
'A $W !9VHJ!?@@ 6
s[   B
GF;F6F;*G>F9?CGG-G1G2G6F;9G;GGfilterbeforelimitc             #    K   t        |      }t        |      t           j                  t        | j                        }|st        d      |t           j                  t        d      x}rrt        |t           vrct        |      }| j                  |d      D ]3  \  }}	|	j                  t        |t        |i      |||      E d{     y t        d| d      t        | j                  |t        t        |t        t        |t           t                 ii      }t        |j                  ||||	            D ]   }
| j!                  |
j                  |
       " y7 w)
z*Get the history of the state of the graph.re  r@  Tr   rr  Nrf  rg  rt  ru  rs  )ri   r&   rK  r+   r   r   r*   rn   r   get_state_historyrm   rj   r   r5   r   r   r]  )r   r   rs  rt  ru  r   rk  rl  rm  rn  checkpoint_tuples              r   rx  zPregel.get_state_history  s     v&3@3H3N3R3R#T%6%64
 233 $D\--.FKKMK%VD\9)-8F!//&$/O 	A	6!33&v0G/VW!!	 4    	A !9VHJ!?@@KK,m(#fTl;O.P*Q	
 !%fV5P!
 	 .. '')9 	+s   CE EBE c                K   t        |      }t        |      t           j                  t        | j                        }|st        d      |t           j                  t        d      x}rmt        |t           vr^t        |      }| j                  |d      2 3 d{   \  }}	|	j                  t        |t        |i      |||      2 3 d{   }
|
 t        | j                  |t        t        |t        t        |t           t                 ii      }|j                  ||||	      2 cg c3 d{   }|7 7 q6  y6 t        d| d      7 6 nc c}w c}D ]*  }| j!                  |j                  |       d{  7   , yw)
z9Asynchronously get the history of the state of the graph.re  r@  Tr   Nrr  rf  rg  rw  )ri   r&   rK  r+   r   r   r*   rn   r   aget_state_historyrm   rj   r   r5   r   alistrb  )r   r   rs  rt  ru  r   rk  rl  rm  rn  stater   ry  s                r   r{  zPregel.aget_state_history  s     v&3@3H3N3R3R#T%6%64
 233 $D\--.FKKMK%VD\9)-8F#'#6#6QU#6#V 
A 
Aia#)#<#<&v0G/VW!!	 $= $    %  K
 KK,m(#fTl;O.P*Q	
 (--vU6 . !
 !
 1
A  $  $W !9VHJ!?@@!
 !
 !
 	 55 '')9   	s   BFED<E)FE D>E AF.E0E4E
5E8E<E>E  FEE+FFFc                   
 t        |      t           j                  t         j                        

st        d      t        |      dk(  rt        d      t        d |D              rt        d      |t           j                  t        d      x}rgt        |t           vrXt        |      } j                  |d      D ](  \  }}|j                  t        |t        
i      |      c S  t        d	| d
      	 	 	 	 	 	 d
 fd}t        |t        t        |t           t                 i      }|D ]  }	 |||	      } |S )aF  Apply updates to the graph state in bulk. Requires a checkpointer to be set.

        Args:
            config: The config to apply the updates to.
            supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state.
                        Each update is a tuple of the form `(values, as_node)`.

        Raises:
            ValueError: If no checkpointer is set or no updates are provided.
            InvalidUpdateError: If an invalid update is provided.

        Returns:
            RunnableConfig: The updated config.
        re  r   No supersteps providedc              3  8   K   | ]  }t        |      d k(    ywr   Nr(  r   us     r   r   z+Pregel.bulk_update_state.<locals>.<genexpr>?       /qs1v{/   No updates providedr@  Tr   rf  rg  c                :   t        'j                  |       }&j                  |      }|'j                  |j                         |rt        |j                        n	t               }|r|j                  d   j                         ni }|r|j                  j                  dd      nd}t        |t        |t           j                  t        d      i      }|d   }|r-t        ||j                  t                 }i |j                  |}t        'j                  |t        ||dz   |dz               5 \  }	}
|d	   \  }}||t         k(  rt#        |      dkD  rt%        d
      |at'        ||j(                  xs g 'j*                  |	|
|j                  |j                  j                  dd      dz   d'j,                  t/        'j0                  t2              r'j0                  nd d       }|j(                  xs g D cg c]  }|d	   t4        k(  r|dd   c}x}r4t7        |j                  |	t9        dt:        |g       gd 'j<                         |j(                  xs g D ]?  \  }}}|t>        t@        tB        fv r||vr ||   jD                  jG                  ||f       A t7        ||	|jI                         d 'j<                         &jK                  |tM        |d |      i |d|dz   i |r|j                  j                  di       ni di       }tO        ||r|j                  nd       cd d d        S ||t#        |      dkD  rt%        d      tM        |d |      }&jK                  ||i |d|dz   i |r|j                  j                  di       ni di       }tO        ||r|j                  nd       cd d d        S |t:        k(  r@t#        |      dkD  rt%        d      tQ        tS        'jT                  |            x}rt7        ||	t9        dt:        |g       g&jV                  'j<                         |r |j                  j                  d      |dz   nd}&jK                  |tM        ||	|      i |d|tY        |      dt[        ||d               }&j]                  ||t_        ta        tc        |d         t:                           tO        ||r|j                  nd       cd d d        S t%        d'jT                         ||dk(  rt#        |      dkD  rt%        d      tM        |d |      }&jK                  |r|jd                  xs |j                  n||i |d|dz   |r|j                  j                  di       ni di       }tO        ||r|j                  nd       cd d d        S tf        |t           vr||j(                  rxt'        ||j(                  'j*                  |	|
|j                  |j                  j                  dd      dz   d'j,                  t/        'j0                  t2              r'j0                  nd d       }|j(                  xs g D cg c]  }|d	   t4        k(  r|dd   c}x}r4t7        |j                  |	t9        dt:        |g       gd 'j<                         |j(                  D ]?  \  }}}|t>        t@        tB        fv r||vr ||   jD                  jG                  ||f       A |jI                         D cg c]  }|jD                  s| c}x}rt7        ||	|d 'j<                         d d d        g }t#        |      dk(  r6|d	   \  }}|1t#        'j*                        dk(  rti        'j*                        d	   }n|btk        d |d   jI                         D              s?t/        'jT                  t^              r~'jT                  'j*                  v rf'jT                  }nY|Wtm        'fd|d   jo                         D              }|r0t#        |      dk(  r	|d	   d   }n|d   d	   |d   d	   k7  r|d   d   }|t%        d      |'j*                  vrt%        d | d!      |jG                  ||f       nG|D ]B  \  }}|t%        d"      |'j*                  vrt%        d | d!      |jG                  ||f       D g }g }|D ]'  \  }}'j*                  |   jp                  }|st%        d | d#      tQ               }t9        d||t@        g      }t_        ta        tc        |d         t@                    } |jG                  |       |jG                  |        t#        |      dkD  rts        | n|d	   }!|!ju                  |tw        |'jx                  d$z   tz        t}        t~        |j                  'j*                  j                               t        t}        t        	
|      i%             * t        ||      D ]D  \  } }|jD                  D cg c]  }|d	   t        k7  s| }"}|s/|"s2&j]                  ||"|        F t7        |	|&jV                  'j<                        \  }#}$|#rJ d&       tM        ||	|dz         }&jK                  ||i |d|dz   |D ci c]  \  }}||
 c}}|r|j                  j                  di       ni dt[        ||d               }t        ||      D ]B  \  } }|jD                  D cg c]  }|d	   t        k(  s| c}x}%s0&j]                  ||%|        D tO        ||r|j                        S d       S c c}w c c}w c c}w # 1 sw Y   8xY wc c}w c c}}w c c}w )'Nchannel_versionsr6  r7  r@  r1  r   r8  r9  r   1Cannot apply multiple updates when clearing stateTr=  r{   r   parentssourcer6  rE  r  4Cannot create empty checkpoint with multiple updates4Cannot apply multiple updates when updating as inputinputr  r6  rE  r   Received no input writes for __copy__,Cannot copy checkpoint with multiple updatesforkr  r6  r  c              3  J   K   | ]  }|j                         D ]  }|   y wrz   )r   )r   vvr   s      r   r   zFPregel.bulk_update_state.<locals>.perform_superstep.<locals>.<genexpr>Z  s3      1YY[1  11s   !#versions_seenc              3  r   K   | ].  \  }}|j                   v r|j                         D ]  }||f 
 0 y wrz   r   r   r   r%  seenr   r   s       r   r   zFPregel.bulk_update_state.<locals>.perform_superstep.<locals>.<genexpr>e  G      /#At

?!%	/  A//   47!Ambiguous update, specify as_nodeNode  does not exist2as_node is required when applying multiple updates has no writersUpdateStaterun_namer   -Can't write to SharedValues from update_state)Fri   r   ri  r-  r,  r$   rK   r   r1  rK  rm   r*   r&   rT   r~   rd   r6   r(  rA   rH   rL  r   r   r   r   r!   r<   rE   rD   r8   r   r7   r9   r>   rE  rO  r   putrJ   rk   r   rN   r   get_next_versionrM  rZ   
put_writesr   r   r   r3  r)   rP  anysortedr   flat_writersr   invokerl   r   r0   r
   rG   extendr  r-   rF   r   r=   (input_configupdatesr   rQ  r,  checkpoint_previous_versionsr6  checkpoint_configcheckpoint_metadatar~   rS  r   as_noderT  rY  rZ  r[  r   r   next_confignext_checkpointinput_writes	next_steprG  r4  valid_updateslast_seen_by_node	run_tasksrun_task_idswritersrE  rW  task_idrunchannel_writes	mv_writesrm  push_writesr   r   s(                                         r   perform_superstepz3Pregel.bulk_update_state.<locals>.perform_superstepQ  s    #4;;=F **62E (()9)9:5: 0 01@P@R  @E  !3499;" ) 6;5>>%%fb1D 2,fTl.>.>0"/! #)"4$6vu||D?Q$R!&O&O;N&O# FqI f %(G")!* >gn7|a'0O  (%7&!006B JJ$#!LL!NN..vr:Q>*."&**)$*;*;=PQ *.):):!%$(&
$ &+%9%9%?R+ ! t|3 abE+ ; 
 ) % 0 0 (!1"e["!M N $ $ 5 5 */)=)=)C BICA UIy$AA ("*4 (&sO22991a&AB %&$&--/  11 #/"2"2))*dDA1&.$(1H&($ (-~~'9'9)R'H!# #K 0#uU^^$Uf f\ >go7|a'0R  '8
D$&OO"."2"2)'1&.$(1H&($ (-~~'9'9)R'H!# #K 0#uU^^$If fR e#7|a'0R  (-Yt7J7JF-S'TT|T$&$-b%rJK(99 11  %););F)C)O !1H!# "
 '3&6&6--j(IN"5*1(1*.|*<	 5 < *+= >'  %//'(d:d+;&<e DE  4'5d if fp 1;D<O<O;PQ 
 >g&;7|a'0J  '8
D$&OO"."2"2  ++;u||.'1&,$(1H$ (-~~'9'9)R'H!# #K 0#uU^^$if fr -F4L@),, "4",,

 **626:&*"jj  *$*;*;=PQ !--!% $"J( "'!5!5!;'Q4</ !"' { 
 %!,,$-b%bIJ  11 &+%9%9 >	Q	9 ==$j0$"3..55q!f=> -7,=,=,? Lq188 LLuL$&%t?T?TIfN FHM7|q ")!*?s4::!';#DJJ/2G_S 1(9@@B1 . #4#6#6< //4::="&"5"5_(. /'1/'B'H'H'J/ )% )01Q6&7&:1&=G.r2159J29Nq9QQ&7&;A&>G?,-PQQ$**,,uWI_-MNN$$gv%67'. <OFG0P  djj005	1QRR!(('6):;< 13I&(L#0  **W-::,uWI_-MNN16'GVi[IeDD)9$:IFG  &##G,47L14D&0'RS*

 !%]!: ,W + & $

 1.
 ,W * ( ' $	.& D "%\9!= X-1[[!IAaDDL!!I!I^ ++,=~wW	X (--%%LIq !Q"QQ=*:xJJ&**!)& 1HFST?7FwTDIu~~11)R@r )0*=O2PK "%\9!= O.2kk"JQqTT\1"JJ;J ++KgNO
 (uU^^WWRVWWw	+`'& !MGf fv	 "J( U #Ksy   2Cn >m1Dn 4Bn ?D(n 1B7n 2B3n %m6=Bn m;(m;,n 1nn=n"n4n1n  n
r  r   r  zSequence[StateUpdate]r   r   )ri   r&   rK  r+   r   r   r(  r  r*   rn   r   bulk_update_staterm   r5   r   )r   r   
superstepsrk  rl  rm  rn  r  current_config	superstepr   s   `         @r   r  zPregel.bulk_update_state"  su   ( 4A3H3N3R3R#T%6%64
 233z?a566/J//233 $D\--.FKKMK%VD\9)-8F!//&$/O A	6//&v0G/VW A !9VHJ!?@@w	X(w	X3Hw	Xw	Xr ,)3vd|<P/Q+RS
 $ 	JI.~yIN	Jr   c                   
K   t        |      t           j                  t         j                        

st        d      t        |      dk(  rt        d      t        d |D              rt        d      |t           j                  t        d      x}ret        |t           vrVt        |      } j                  |d      2 3 d	{   \  }}|j                  t        |t        
i      |       d	{   c S 	 	 	 	 	 	 d
 fd}t        |t        t        |t           t                 i      }|D ]  }	 |||	       d	{   } |S 7 7 X6 t        d
| d      7 w)aU  Asynchronously apply updates to the graph state in bulk. Requires a checkpointer to be set.

        Args:
            config: The config to apply the updates to.
            supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state.
                        Each update is a tuple of the form `(values, as_node)`.

        Raises:
            ValueError: If no checkpointer is set or no updates are provided.
            InvalidUpdateError: If an invalid update is provided.

        Returns:
            RunnableConfig: The updated config.
        re  r   r  c              3  8   K   | ]  }t        |      d k(    ywr  r  r  s     r   r   z,Pregel.abulk_update_state.<locals>.<genexpr>  r  r  r  r@  Tr   Nrf  rg  c                p  K   t        'j                  |       }&j                  |       d {   }|'j                  |j                         |rt        |j                        n	t               }|r|j                  d   j                         ni }|r|j                  j                  dd      nd}t        |t        |t           j                  t        d      i      }|d   }|r-t        ||j                  t                 }i |j                  |}t        'j                  |t        ||dz   |dz               4 d {   \  }	}
|d	   \  }}||t         k(  rt#        |      dkD  rt%        d
      |at'        ||j(                  xs g 'j*                  |	|
|j                  |j                  j                  dd      dz   d'j,                  t/        'j0                  t2              r'j0                  nd d       }|j(                  xs g D cg c]  }|d	   t4        k(  r|dd   c}x}r4t7        |j                  |	t9        dt:        |g       gd 'j<                         |j(                  xs g D ]?  \  }}}|t>        t@        tB        fv r||vr ||   jD                  jG                  ||f       A t7        ||	|jI                         d 'j<                         &jK                  |tM        |d |      i |d|dz   i |r|j                  j                  di       ni di        d {   }tO        ||r|j                  nd       cd d d       d {    S ||t#        |      dkD  rt%        d      tM        |d |      }&jK                  ||i |d|dz   i |r|j                  j                  di       ni di        d {   }tO        ||r|j                  nd       cd d d       d {    S |t:        k(  rYt#        |      dkD  rt%        d      tQ        tS        'jT                  |            x}rt7        ||	t9        dt:        |g       g&jV                  'j<                         |r |j                  j                  d      |dz   nd}&jK                  |tM        ||	|      i |d|tY        |      dt[        ||d                d {   }&j]                  ||t_        ta        tc        |d         t:                           d {    tO        ||r|j                  nd       cd d d       d {    S t%        d'jT                         ||dk(  rt#        |      dkD  rt%        d      tM        |d |      }&jK                  |r|jd                  xs |j                  n||i |d|dz   |r|j                  j                  di       ni di        d {   }tO        ||r|j                  nd       cd d d       d {    S tf        |t           vr||j(                  rxt'        ||j(                  'j*                  |	|
|j                  |j                  j                  dd      dz   d'j,                  t/        'j0                  t2              r'j0                  nd d       }|j(                  xs g D cg c]  }|d	   t4        k(  r|dd   c}x}r4t7        |j                  |	t9        dt:        |g       gd 'j<                         |j(                  D ]?  \  }}}|t>        t@        tB        fv r||vr ||   jD                  jG                  ||f       A |jI                         D cg c]  }|jD                  s| c}x}rt7        ||	|d 'j<                         d d d       d {    g }t#        |      dk(  r|d	   \  }}|1t#        'j*                        dk(  rti        'j*                        d	   }n|A|s?t/        'jT                  t^              r~'jT                  'j*                  v rf'jT                  }nY|Wtk        'fd|d   jm                         D              }|r0t#        |      dk(  r	|d	   d   }n|d   d	   |d   d	   k7  r|d   d   }|t%        d      |'j*                  vrt%        d| d       |jG                  ||f       nG|D ]B  \  }}|t%        d!      |'j*                  vrt%        d| d       |jG                  ||f       D g }g }|D ]/  \  }}'j*                  |   jn                  }|st%        d| d"      tQ               }t9        d||t@        g      }t_        ta        tc        |d         t@                    } |jG                  |       |jG                  |        t#        |      dkD  rtq        | n|d	   }!|!js                  |tu        |'jv                  d#z   tx        t{        t|        |j~                  'j*                  j                               t        t{        t        	
|      i$             d {    2 t        ||      D ]L  \  } }|jD                  D cg c]  }|d	   t        k7  s| }"}|s/|"s2&j]                  ||"|        d {    N t7        |	|&jV                  'j<                        \  }#}$|#rJ d%       tM        ||	|dz         }&jK                  ||i |d|dz   |D ci c]  \  }}||
 c}}|r|j                  j                  di       ni dt[        ||d                d {   }t        ||      D ]J  \  } }|jD                  D cg c]  }|d	   t        k(  s| c}x}%s0&j]                  ||%|        d {    L tO        ||r|j                        S d       S 7 7 
c c}w 7 7 7 #7 7 7 7 7 7 c c}w c c}w 7 S# 1 d {  7  sw Y   dxY w7 c c}w 7 uc c}}w 7 c c}w 7 w)&Nr  r6  r7  r@  r1  r   r8  r9  r   r  Tr=  r{   r   r  r  r  r  r  r  r   r  r  r  r  r  c              3  r   K   | ].  \  }}|j                   v r|j                         D ]  }||f 
 0 y wrz   r  r  s       r   r   zHPregel.abulk_update_state.<locals>.aperform_superstep.<locals>.<genexpr>  r  r  r  r  r  r  r  r  r  r  r  r  )Eri   r   rq  r-  r,  r$   rK   r   r1  rK  rm   r*   r&   rS   r~   rd   r6   r(  rA   rH   rL  r   r   r   r   r!   r<   rE   rD   r8   r   r7   r9   r>   rE  rO  r   aputrJ   rk   r   rN   r   r  rM  rZ   aput_writesr   r   r   r3  r)   rP  r  r   r  r   ainvokerl   r   r0   r
   rG   r  r  r-   rF   r   r=   r  s(                                         r   aperform_superstepz5Pregel.abulk_update_state.<locals>.aperform_superstep   s8     #4;;=F&11&99E (()9)9:5: 0 01@P@R  @E  !3499;" ) 6;5>>%%fb1D 2,fTl.>.>0"/! #)"4$6vu||D?Q$R!&O&O;N&O#+FqI f f")!*>gn7|a'0O  (%7&!006B JJ$#!LL!NN..vr:Q>*."&**)$*;*;=PQ *.):):!%$(&
$ &+%9%9%?R+ ! t|3 abE+ ; 
 ) % 0 0 (!1"e["!M N $ $ 5 5 */)=)=)C BICA UIy$AA ("*4 (&sO22991a&AB %&$&--/  11 )5(9(9))*dDA1&.$(1H&($ (-~~'9'9)R'H!# ) #K 0#uU^^$Wf f f^ >go7|a'0R  '8
D$&OO(4(9(9)'1&.$(1H&($ (-~~'9'9)R'H!# ) #K 0#uU^^$Kf f fT e#7|a'0R  (-Yt7J7JF-S'TT|T$&$-b%rJK(99 11  %););F)C)O !1H!# "
 -9,=,=--j(IN"5*1(1*.|*<	 5 < *+= >- '  +66'(d:d+;&<e DE    4'5d kf f fr 1;D<O<O;PQ 
 >g&;7|a'0J  '8
D$&OO(4(9(9  ++;u||.'1&,$(1H$ (-~~'9'9)R'H!# ) #K 0#uU^^$kf f ft -F4L@),, "4",,

 **626:&*"jj  *$*;*;=PQ !--!% $"J( "'!5!5!;'Q4</ !"' { 
 %!,,$-b%bIJ  11 &+%9%9 >	Q	9 ==$j0$"3..55q!f=> -7,=,=,? Lq188 LLuL$&%t?T?TIf fN FHM7|q ")!*?s4::!';#DJJ/2G_U"4#6#6< //4::="&"5"5_(. /'1/'B'H'H'J/ )% )01Q6&7&:1&=G.r2159J29Nq9QQ&7&;A&>G?,-PQQ$**,,uWI_-MNN$$gv%67'. <OFG0P  djj005	1QRR!(('6):;< 13I&(L#0  **W-::,uWI_-MNN16'GVi[IeDD)9$:IFG  &##G,47L14D&0'RS*kk !%]!: ,W + & $

 1.
 ,W * ( ' $	.&   D "%\9!= -1[[!IAaDDL!!I!I^&22)>7  	 (--%%LIq !Q"QQ=*:xJJ , 1 1!)& 1HFST?7FwTDIu~~11)R@r )0*=O2P! K "%\9!= V.2kk"JQqTT\1"JJ;J&22;WUUUV (uU^^WWRVWWa :.fB+8#{fn#of~' _fL#Mfb'$ !MGf f f f~0 "J* U  #KUs  +p6oDp6 op6Cpo(C>p&o'pp6o!p6A+po$pp6+o',p61Cpo*=po-p*p66o07p6<B!po3p;p6o6p6B3p o9Bp1o>o>p%p60p1J"p6p'p6;ppp6p6p6/p$0Ap6p';p6p-$p6<p/p/p6p6.p4/%p6p6p!p6$p'p6*p-p0p63p6p69
pp6ppp	p6p6'p6/p6r  )ri   r&   rK  r+   r   r   r(  r  r*   rn   r   abulk_update_staterm   r5   r   )r   r   r  rk  rl  rm  rn  r  r  r  r   s   `         @r   r  zPregel.abulk_update_state  s    ( 4A3H3N3R3R#T%6%64
 233z?a566/J//233 $D\--.FKKMK%VD\9)-8F#'#6#6QU#6#V A Aia#66&v0G/VW  u	X(u	X3Hu	Xu	Xn ,)3vd|<P/Q+RS
 $ 	QI#5ni#PPN	QIA $W !9VHJ!?@@z QsI   C E"EE	E)E"5E6AE"E E"EE"E"c                >    | j                  |t        ||      gg      S )zUpdate the state of the graph with the given values, as if they came from
        node `as_node`. If `as_node` is not provided, it will be set to the last node
        that updated the state, if not ambiguous.
        )r  rf   r   r   r   r  s       r   update_statezPregel.update_state~  s$     %%fFG0L/M.NOOr   c                Z   K   | j                  |t        ||      gg       d{   S 7 w)zAsynchronously update the state of the graph with the given values, as if they came from
        node `as_node`. If `as_node` is not provided, it will be set to the last node
        that updated the state, if not ambiguous.
        N)r  rf   r  s       r   aupdate_statezPregel.aupdate_state  s.      ,,V{677S6T5UVVVVs   "+)+c                  |d   dk  rt        d      ||n| j                  }|| j                  }nt        || j                         |xs | j
                  }|xs | j                  }||n| j                  }t        |t              s|g}t        |j                  t        i       v rdg}| j                  du rd }nRt        |j                  t        i       v r|t           t           }n%| j                  du rt        d      | j                  }|rE|j                  t              s0t        d|j                   D cg c]  }|j"                   c}       t$        |j                  t        i       v r|t           t$           }	n| j&                  }	t(        |j                  t        i       v r|t           t(           }
n| j*                  }
|t-        |      |||||	|
fS c c}w )	Nrecursion_limitr   z"recursion_limit must be at least 1r   FTz1checkpointer=True cannot be used for root graphs.zHCheckpointer requires one or more of the following 'configurable' keys: )r   r   r  r\   r~   r   r   r   r   r   r4   rK  r&   r   r+   RuntimeErrorr   r   r1   r   r'   r   set)r   r   r   output_keysinterrupt_beforeinterrupt_afterr   r   r&  r   r   s              r   	_defaultszPregel._defaults  s   & #$q(ABB*

33K+t}}5+Jt/J/J)GT-G-G%0%<k$BRBR+t,&-KD"!55#*K%7;L$

4(<<!$<(?@L$&RSS,,L

4 0Zjv  kD  kD  \Eef\]\`\`  \E  [F  G  vzz$33&,Tl3C&DEJJEvzz$33&,Tl3C&DEJJE	
 		
 \Es   9G
)r   r  r  r  checkpoint_duringr   r   c          	   #  	  	K   t               d(	fd}
t        | j                  |      }t        |      }|j	                  d||j                  d| j                               |j                  d            }	 | j                  |||||      \  }}}}}}}}| j                  du r7t        t        |t           t                 }t        |      |t           t        <   d|v r.|j                  j                  t!        j"                               d	|v rfd
|t           t$        <   |||t           t&        <   t)        |fi d| j*                  dt-        j"                  |      d|d|d|d|d| j.                  d| j0                  d|d| j2                  d|d|d|d|d||n|t           j                  t&        d      d| j4                  d| j6                  d| j8                  d| j:                  5 t=        |t           j                  t>        tA        jB                  jD                              tA        jB                  jF                        |t           j                  tH                    }	r$jJ                  j                  t           tL        <   | jN                  s
	sd|v sd	|v r:djP                  jS                  jT                  jV                         d)fd}nd}jY                  | jZ                         rʉj]                         D ]*  }j_                  |j`                  |jb                  d!       , |jY                  jd                  jg                         D cg c]  }|jb                  r| c}| jh                  |jj                  "      D ]  } |
       E d{     jY                  | jZ                         rddd        |
       E d{    jl                  d#k(  r-to        d$|d%    d&tp        jr                  '      }tu        |      |jw                  jx                         yc c}w 7 # 1 sw Y   wxY w7 o# tz        $ r}|j}                  |        d}~ww xY ww)*a  Stream graph steps for a single input.

        Args:
            input: The input to the graph.
            config: The configuration to use for the run.
            stream_mode: The mode to stream output, defaults to self.stream_mode.
                Options are:

                - `"values"`: Emit all values in the state after each step, including interrupts.
                    When used with functional API, values are emitted once at the end of the workflow.
                - `"updates"`: Emit only the node or task names and updates returned by the nodes or tasks after each step.
                    If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately.
                - `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`.
                - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks.
                - `"debug"`: Emit debug events with as much information as possible for each step.
            output_keys: The keys to stream, defaults to all non-context channels.
            interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph.
            interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph.
            checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved.
            debug: Whether to print debug information during execution, defaults to False.
            subgraphs: Whether to stream subgraphs, defaults to False.

        Yields:
            The output of each step in the graph. The output shape depends on the stream_mode.

        Example: Using stream_mode="values":
            ```python
            import operator
            from typing_extensions import Annotated, TypedDict
            from langgraph.graph import StateGraph, START

            class State(TypedDict):
                alist: Annotated[list, operator.add]
                another_list: Annotated[list, operator.add]

            builder = StateGraph(State)
            builder.add_node("a", lambda _state: {"another_list": ["hi"]})
            builder.add_node("b", lambda _state: {"alist": ["there"]})
            builder.add_edge("a", "b")
            builder.add_edge(START, "a")
            graph = builder.compile()

            for event in graph.stream({"alist": ['Ex for stream_mode="values"']}, stream_mode="values"):
                print(event)

            # {'alist': ['Ex for stream_mode="values"'], 'another_list': []}
            # {'alist': ['Ex for stream_mode="values"'], 'another_list': ['hi']}
            # {'alist': ['Ex for stream_mode="values"', 'there'], 'another_list': ['hi']}
            ```

        Example: Using stream_mode="updates":
            ```python
            for event in graph.stream({"alist": ['Ex for stream_mode="updates"']}, stream_mode="updates"):
                print(event)

            # {'a': {'another_list': ['hi']}}
            # {'b': {'alist': ['there']}}
            ```

        Example: Using stream_mode="debug":
            ```python
            for event in graph.stream({"alist": ['Ex for stream_mode="debug"']}, stream_mode="debug"):
                print(event)

            # {'type': 'task', 'timestamp': '2024-06-23T...+00:00', 'step': 1, 'payload': {'id': '...', 'name': 'a', 'input': {'alist': ['Ex for stream_mode="debug"'], 'another_list': []}, 'triggers': ['start:a']}}
            # {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 1, 'payload': {'id': '...', 'name': 'a', 'result': [('another_list', ['hi'])]}}
            # {'type': 'task', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'input': {'alist': ['Ex for stream_mode="debug"'], 'another_list': ['hi']}, 'triggers': ['a']}}
            # {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}}
            ```

        Example: Using stream_mode="custom":
            ```python
            from langgraph.types import StreamWriter

            def node_a(state: State, writer: StreamWriter):
                writer({"custom_data": "foo"})
                return {"alist": ["hi"]}

            builder = StateGraph(State)
            builder.add_node("a", node_a)
            builder.add_edge(START, "a")
            graph = builder.compile()

            for event in graph.stream({"alist": ['Ex for stream_mode="custom"']}, stream_mode="custom"):
                print(event)

            # {'custom_data': 'foo'}
            ```

        Example: Using stream_mode="messages":
            ```python
            from typing_extensions import Annotated, TypedDict
            from langgraph.graph import StateGraph, START
            from langchain_openai import ChatOpenAI

            llm = ChatOpenAI(model="gpt-4o-mini")

            class State(TypedDict):
                question: str
                answer: str

            def node_a(state: State):
                response = llm.invoke(state["question"])
                return {"answer": response.content}

            builder = StateGraph(State)
            builder.add_node("a", node_a)
            builder.add_edge(START, "a")
            graph = builder.compile()

            for event in graph.stream({"question": "What is the capital of France?"}, stream_mode="messages"):
                print(event)

            # (AIMessageChunk(content='The', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], 'langgraph_path': ('__pregel_pull', 'a'), 'langgraph_checkpoint_ns': '...', 'checkpoint_ns': '...', 'ls_provider': 'openai', 'ls_model_name': 'gpt-4o-mini', 'ls_model_type': 'chat', 'ls_temperature': 0.7})
            # (AIMessageChunk(content=' capital', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], ...})
            # (AIMessageChunk(content=' of', additional_kwargs={}, response_metadata={}, id='...'), {...})
            # (AIMessageChunk(content=' France', additional_kwargs={}, response_metadata={}, id='...'), {...})
            # (AIMessageChunk(content=' is', additional_kwargs={}, response_metadata={}, id='...'), {...})
            # (AIMessageChunk(content=' Paris', additional_kwargs={}, response_metadata={}, id='...'), {...})
            ```
        c               3     K   	 	 j                  d      \  } }}rt        t              r| ||f n$t        t              r||f nr| |f n| V# t        j                  $ r Y y w xY ww)NF)block)rK  queueEmptyr   r   ro  modepayloadstreamr   r   s      r   outputzPregel.stream.<locals>.output`	  s     (.


(?%Bg K!>tW--T2/)w-'!M  {{ s'   A4A ?A4A1.A40A11A4Nr  run_idr   r  r   r  r  r  r   Tmessagescustomc                ,    j                  dd| f      S Nr{   r  )r  )r   r  s    r   <lambda>zPregel.stream.<locals>.<lambda>	  s    6::1%D r   r   r  r   r   r   r   r   r   r  stream_keysr  r  r?  r   r  r   migrate_checkpointr   r   )submitr  node_finishedc                 f    j                         r j                  j                        S S rz   )doner  wait)loopr  waiters   r   
get_waiterz!Pregel.stream.<locals>.get_waiter	  s,    !>V[[]%)[[%=F#)M#)Mr   
input_keyscachedtimeoutr  schedule_taskout_of_stepsRecursion limit of r  r reached without hitting a stop condition. You can increase the limit by setting the `recursion_limit` config key.message
error_coder   r   )r   zconcurrent.futures.Future[None])?rs   ri   r   r   on_chain_startrK  r   r  r   r   r   r&   r*   rn   inheritable_handlersrO  rU   r  r3   r(   rR   r   rQ   r   r~   r  r   r-  r   r   rY   r/   weakref
WeakMethodr  r  r,   r  r2   r   stackcallback_countreleasetickr   match_cached_writesoutput_writesr   rE  r4  r   r   accept_pushstatusrB   r?   GRAPH_RECURSION_LIMITr@   on_chain_endr  BaseExceptionon_chain_error)r   r  r   r   r  r  r  r  r   r   r  callback_managerrun_managerstream_modesinterrupt_before_interrupt_after_r   r   r   ro  runnerr  rW  rG  rm  msger  r  r  s      `     `                 @@@r   r  zPregel.stream  s    N 	" t{{F3:6B&55J8::h'	 6 
A	 ''!1 /  	!    D(#vd|,DEF9Mb9Qt56\)0077)&**5 <':t56 !,=Nt9: ,, &fjj,? 	
   * jj mm ( !55 "3 !1 $ " %0 #4D\%%&BDI%& "&!6!6'( $(#;#;)* "..+, "..- K,. %!$<++0'2D2DT[[2Q  '11$//B"(,"2"23K"L ;?;;DKK%&78 %% !\1</ @DF JJ''(=(=>* * "&J ii4+>+>i? $ 8 8 : N**477DKK*MN#[[$(JJ$5$5$7HqqxxH $ 1 1#-&*&6&6	 )  , $*8++, ii4+>+>i?CK,Z x{{n,*-f5F.G-H IM M  )>> *#..$$T[[1+ I ,WK, K,Z    	&&q)	su   A6R>FQ$ E6Q:Q
Q
,Q<Q="Q Q$ 2Q"3AQ$ RQQQ$ $	R-Q??RRc          	    F
  	 K   t                t        j                         t        t        t
        gdf   t        j                   j                              }
d* 	fd}t        | j                  |      }t        |      }|j                  d||j                  d| j                               |j                  d             d{   }t        t!        d |j"                  D        d      nd}	 | j%                  |||||      \  }}}}}}}}| j&                  d	u r7t        t(        |t*           t,                 }t/        |      |t*           t,        <   d
|v r$|j0                  j3                  t5        |
             d|v r fd|t*           t6        <   |||t*           t8        <   t;        |fi d| j<                  dt?         j                  |      d|d|d|d|d| j@                  d| jB                  d|d| jD                  d|d|d|d|d||n|t*           j                  t8        d	      d| jF                  d| jH                  d| jJ                  d| jL                  4 d{   }tO        |t*           j                  tP        tS        jT                  |jV                              tS        jT                  |jX                        ||t*           j                  tZ                     }	r$t?        |
|      |j                  t*           t\        <   | j^                  s
	sd
|v sd|v rd+ fd!}nd}|ja                  | jb                  "      r|je                          d{   D ]*  }|jg                  |jh                  |jj                  d	#       , |jm                  |jn                  jq                         D cg c]  }|jj                  r| c}| jr                  ||jt                  $      2 3 d{   } |       D ]  }| 	 ddd      d{     |       D ]  }| 	 jv                  d%k(  r-ty        d&|d'    d(tz        j|                  )      }t        |      |j                  |j                         d{    y7 7 %7 0c c}w 7 6 |ja                  | jb                  "      rl7 # 1 d{  7  sw Y   xY w7 L# t        $ r3}t        j                  |j                  |             d{  7    d}~ww xY ww),a  Asynchronously stream graph steps for a single input.

        Args:
            input: The input to the graph.
            config: The configuration to use for the run.
            stream_mode: The mode to stream output, defaults to self.stream_mode.
                Options are:

                - `"values"`: Emit all values in the state after each step, including interrupts.
                    When used with functional API, values are emitted once at the end of the workflow.
                - `"updates"`: Emit only the node or task names and updates returned by the nodes or tasks after each step.
                    If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately.
                - `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`.
                - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks.
                - `"debug"`: Emit debug events with as much information as possible for each step.
            output_keys: The keys to stream, defaults to all non-context channels.
            interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph.
            interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph.
            checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved.
            debug: Whether to print debug information during execution, defaults to False.
            subgraphs: Whether to stream subgraphs, defaults to False.

        Yields:
            The output of each step in the graph. The output shape depends on the stream_mode.

        Example: Using stream_mode="values":
            ```python
            import operator
            from typing_extensions import Annotated, TypedDict
            from langgraph.graph import StateGraph, START

            class State(TypedDict):
                alist: Annotated[list, operator.add]
                another_list: Annotated[list, operator.add]

            builder = StateGraph(State)
            builder.add_node("a", lambda _state: {"another_list": ["hi"]})
            builder.add_node("b", lambda _state: {"alist": ["there"]})
            builder.add_edge("a", "b")
            builder.add_edge(START, "a")
            graph = builder.compile()

            async for event in graph.astream({"alist": ['Ex for stream_mode="values"']}, stream_mode="values"):
                print(event)

            # {'alist': ['Ex for stream_mode="values"'], 'another_list': []}
            # {'alist': ['Ex for stream_mode="values"'], 'another_list': ['hi']}
            # {'alist': ['Ex for stream_mode="values"', 'there'], 'another_list': ['hi']}
            ```

        Example: Using stream_mode="updates":
            ```python
            async for event in graph.astream({"alist": ['Ex for stream_mode="updates"']}, stream_mode="updates"):
                print(event)

            # {'a': {'another_list': ['hi']}}
            # {'b': {'alist': ['there']}}
            ```

        Example: Using stream_mode="debug":
            ```python
            async for event in graph.astream({"alist": ['Ex for stream_mode="debug"']}, stream_mode="debug"):
                print(event)

            # {'type': 'task', 'timestamp': '2024-06-23T...+00:00', 'step': 1, 'payload': {'id': '...', 'name': 'a', 'input': {'alist': ['Ex for stream_mode="debug"'], 'another_list': []}, 'triggers': ['start:a']}}
            # {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 1, 'payload': {'id': '...', 'name': 'a', 'result': [('another_list', ['hi'])]}}
            # {'type': 'task', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'input': {'alist': ['Ex for stream_mode="debug"'], 'another_list': ['hi']}, 'triggers': ['a']}}
            # {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}}
            ```

        Example: Using stream_mode="custom":
            ```python
            from langgraph.types import StreamWriter

            async def node_a(state: State, writer: StreamWriter):
                writer({"custom_data": "foo"})
                return {"alist": ["hi"]}

            builder = StateGraph(State)
            builder.add_node("a", node_a)
            builder.add_edge(START, "a")
            graph = builder.compile()

            async for event in graph.astream({"alist": ['Ex for stream_mode="custom"']}, stream_mode="custom"):
                print(event)

            # {'custom_data': 'foo'}
            ```

        Example: Using stream_mode="messages":
            ```python
            from typing_extensions import Annotated, TypedDict
            from langgraph.graph import StateGraph, START
            from langchain_openai import ChatOpenAI

            llm = ChatOpenAI(model="gpt-4o-mini")

            class State(TypedDict):
                question: str
                answer: str

            async def node_a(state: State):
                response = await llm.ainvoke(state["question"])
                return {"answer": response.content}

            builder = StateGraph(State)
            builder.add_node("a", node_a)
            builder.add_edge(START, "a")
            graph = builder.compile()

            async for event in graph.astream({"question": "What is the capital of France?"}, stream_mode="messages"):
                print(event)

            # (AIMessageChunk(content='The', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], 'langgraph_path': ('__pregel_pull', 'a'), 'langgraph_checkpoint_ns': '...', 'checkpoint_ns': '...', 'ls_provider': 'openai', 'ls_model_name': 'gpt-4o-mini', 'ls_model_type': 'chat', 'ls_temperature': 0.7})
            # (AIMessageChunk(content=' capital', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], ...})
            # (AIMessageChunk(content=' of', additional_kwargs={}, response_metadata={}, id='...'), {...})
            # (AIMessageChunk(content=' France', additional_kwargs={}, response_metadata={}, id='...'), {...})
            # (AIMessageChunk(content=' is', additional_kwargs={}, response_metadata={}, id='...'), {...})
            # (AIMessageChunk(content=' Paris', additional_kwargs={}, response_metadata={}, id='...'), {...})
            ```
        Nc               3     K   	 	 j                         \  } }}rt        t              r| ||f n$t        t              r||f nr| |f n| T# t        j                  $ r Y y w xY wwrz   )
get_nowaitr   
QueueEmptyr   r   r  s      r   r  zPregel.astream.<locals>.output
  s     (.(9(9(;%Bg K!>tW--T2/)w-'!M  )) s'   A2A ?A2A/,A2.A//A2r  r  r  c              3  `   K   | ]&  }t        |t              rt        |t              sd  ( yw)TN)r   rt   rU   )r   hs     r   r   z!Pregel.astream.<locals>.<genexpr>
  s.      !!%>?&q*?@ s   ,.Fr  Tr  r  c                B    j                  j                  dd| f      S r  )call_soon_threadsafe
put_nowait)r   aioloopr  s    r   r  z Pregel.astream.<locals>.<lambda>
  s$    g::))B!+< r   r   r  r   r   r   r   r   r   r  r  r  r  r?  r   r  r   r  r   r   )r  r  use_astreamr  c                 B     j                  j                               S rz   )create_taskr  )r0  r  s   r   r  z"Pregel.astream.<locals>.get_waiter  s    &226;;=AAr   r   r  r  r  r  r  r	  r
  r  )r   zasyncio.Task[None])Err   r   get_running_loopr   r   rg   r
   r.  r/  ri   r   r   r  rK  r   rt   r0  handlersr  r   r   r&   r*   rn   r  rO  rU   r3   r(   rP   r   rQ   r   r~   r  r   r-  r   r   rY   r/   r  r  r  r  r,   r2   r   r  r   amatch_cached_writesr  r   rE  atickr4  r   r   aaccept_pushr  rB   r?   r  r@   r  r  r  shieldr  )!r   r  r   r   r  r  r  r  r   r   
stream_putr  r  r   	do_streamr!  r"  r#  r   r   r   ro  r  r$  r  rW  rG  rm  or%  r&  r0  r  s!      `     `                     @@r   astreamzPregel.astream	  sh    N **,k]D()G00&2C2CD


	" t{{F3@H,;;J8::h'	 < 
 
" )4 (11   	|	 ''!1 /  	!    D(#vd|,DEF9Mb9Qt56\)0077)*5 <' t56 !,=Nt9:& ,, &f&7&7F 	
   * jj mm ( !55 "3 !1 $ " %0 #4D\%%&BDI%& "&!6!6'( $(#;#;)* "..+, "..- C$ C$. %!$<++0'2D2DT[[2Q  '11$//B )"(,"2"23K"L ;I"L<DKK%&78
 %% !\1</B "&J ii4+>+>i?&*&?&?&A A N**477DKK*MN#)<<$(JJ$5$5$7HqqxxH $ 1 1#-&*&7&7	 $0 $ $ $a "( $A"#G$EC$ C$J X  {{n,*-f5F.G-H IM M  )>> *#..**4;;777[
tC$r !B I$ $ ii4+>+>i?qC$ C$ C$ C$f 8 	..!;!;A!>???	s   CT!R)T!6FS" 7R8S" ;C4S/R0ASR"
R"
S=R)R'R)SS" #S	$A/S" S S" T!S" S"S'R))S	S" SSSS" "	T+'TTTTT!r   r  r  r  r  r   c          
     p   ||n| j                   }d}
g }g } | j                  ||f||||||d|	D ]T  }|dk(  r<t        |t              r)|j	                  t
              x}|j                  |       A|}
D|j                  |       V |dk(  r(|r$t        |
t              ri |
t
        |iS t
        |iS |
S |S )a  Run the graph with a single input and config.

        Args:
            input: The input data for the graph. It can be a dictionary or any other type.
            config: Optional. The configuration for the graph run.
            stream_mode: Optional[str]. The stream mode for the graph run. Default is "values".
            output_keys: Optional. The output keys to retrieve from the graph run.
            interrupt_before: Optional. The nodes to interrupt the graph run before.
            interrupt_after: Optional. The nodes to interrupt the graph run after.
            debug: Optional. Enable debug mode for the graph run.
            **kwargs: Additional keyword arguments to pass to the graph run.

        Returns:
            The output of the graph run. If stream_mode is "values", it returns the latest output.
            If stream_mode is not "values", it returns a list of output chunks.
        Nr>  r   )r   r  r   rM  rK  r9   r  rO  r   r  r   r   r  r  r  r  r   r   latestchunksr5  chunkintss                  r   r  zPregel.invoke+  s    : &1%<k$BVBV-135&(
 T[[

 $#-+/

 

 	%E h&ud+!&9!55B%%d+"Fe$)	%, (" "&$/ 6v5y*5 $Z0
 MMr   c          
       K   ||n| j                   }d}
g }g } | j                  ||f||||||d|	2 3 d{   }|dk(  r<t        |t              r)|j	                  t
              x}|j                  |       G|}
J|j                  |       \7 W6 |dk(  r(|r$t        |
t              ri |
t
        |iS t
        |iS |
S |S w)a  Asynchronously invoke the graph on a single input.

        Args:
            input: The input data for the computation. It can be a dictionary or any other type.
            config: Optional. The configuration for the computation.
            stream_mode: Optional. The stream mode for the computation. Default is "values".
            output_keys: Optional. The output keys to include in the result. Default is None.
            interrupt_before: Optional. The nodes to interrupt before. Default is None.
            interrupt_after: Optional. The nodes to interrupt after. Default is None.
            debug: Optional. Whether to enable debug mode. Default is None.
            **kwargs: Additional keyword arguments.

        Returns:
            The result of the computation. If stream_mode is "values", it returns the latest value.
            If stream_mode is "chunks", it returns a list of chunks.
        Nr>  r   )r   r=  r   rM  rK  r9   r  rO  r@  s                  r   r  zPregel.ainvokeo  s    < &1%<k$BVBV-135&(
'4<<

 $#-+/

 

 	% 	%% h&ud+!&9!55B%%d+"Fe$)	% 

, (" "&$/ 6v5y*5 $Z0
 MMs&   2CBBBACB0Cc                8   | j                   st        d      |xs | j                  j                         }g }|D ]C  }|| j                  v s|j	                  t
        t        | j                  |         xs d|f       E | j                   j                  |       y)z$Clear the cache for the given nodes.3No cache is set for this graph. Cannot clear cache.__dynamic__N)r   r   r   r  rO  r%   rI   clearr   r   
namespacesr   s       r   clear_cachezPregel.clear_cache  s    zzRSS**,.
 	Dtzz!!!'#DJJt$45F	 	

$r   c                T  K   | j                   st        d      |xs | j                  j                         }g }|D ]C  }|| j                  v s|j	                  t
        t        | j                  |         xs d|f       E | j                   j                  |       d{    y7 w)z3Asynchronously clear the cache for the given nodes.rG  rH  N)r   r   r   r  rO  r%   rI   aclearrJ  s       r   aclear_cachezPregel.aclear_cache  s     zzRSS**,.
 	Dtzz!!!'#DJJt$45F	 jj
+++s   AB(AB( B&!B().r   r   r~   z0dict[str, BaseChannel | ManagedValueSpec] | Noner   r   r   rh   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   bool | Noner   BaseCheckpointSaver | Noner   r   r   r   r   z#RetryPolicy | Sequence[RetryPolicy]r   r   r   r   r   r   r   r   r   z"Mapping[str, Sequence[str]] | Noner   r   r   r   rz   )r   r   r   z
int | boolr   r   )r   r   r   dict[str, Any])r   dict[str, Any] | Noner   r   )r   r   r   r   r   r   )r   r   )r   zlist[ConfigurableFieldSpec])r   Sequence[str] | Noner   type[BaseModel])r   rT  r   rR  )r   r   )r   r   r   rU  )r   r   r   rR  )r   r   )r   r   )r!  r   r"  r   r   z$Iterator[tuple[str, PregelProtocol]])r!  r   r"  r   r   z)AsyncIterator[tuple[str, PregelProtocol]])r,  r"   r   r   )NF)
r   r   rQ  zCheckpointTuple | Noner"  rQ  rR  r   r   re   )r   r   r   r   r   re   )
r   r   rs  rS  rt  r   ru  
int | Noner   zIterator[StateSnapshot])
r   r   rs  rS  rt  r   ru  rV  r   zAsyncIterator[StateSnapshot])r   r   r  zSequence[Sequence[StateUpdate]]r   r   )r   r   r   zdict[str, Any] | Any | Noner  r   r   r   )r   r   r   dict[str, Any] | Anyr  r   r   r   )r   r   r   $StreamMode | list[StreamMode] | Noner  r   r  All | Sequence[str] | Noner  rY  r   rP  r   ztuple[bool, set[StreamMode], str | Sequence[str], All | Sequence[str], All | Sequence[str], BaseCheckpointSaver | None, BaseStore | None, BaseCache | None])r  rW  r   r   r   rX  r  r   r  rY  r  rY  r  rP  r   rP  r   r   r   zIterator[dict[str, Any] | Any])r  rW  r   r   r   rX  r  r   r  rY  r  rY  r  rP  r   rP  r   r   r   z#AsyncIterator[dict[str, Any] | Any])r  rW  r   r   r   rh   r  r   r  rY  r  rY  r  rP  r   rP  r   r   r   rW  )r   rT  r   r   ):r   r   r   __doc____annotations__r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   propertyr   r   r  r  r  r  r  r  r  r  r  r   r   r-  r]  rb  rN  r`  rx  r{  r  r  r  r  r  r  r=  r  r  rL  rO  __classcell__)r   s   @r   r   r      sW   @D ! 77&K&6L$0 )(26O/6O..//''!%L,%TKQ!%L,%K"E"A"E"B*,L',O'+L$+ %)K!(*.K'.$(F!(D#11 #"*"6:5768%)!37"&"&<>+/(,.2(,?C12 %2 C	2
 2  2 2 -2 42  32 !42 ,2 #2 2 12   !2"  #2$ :%2& )'2( &)2* ,+2, &-2. =/20 122 
32j /3
JO
+
=G
	
: /3(
JO(
+(
=G(
	(
T
'
 "
 "
H @D O( 26#.#	# * *" /3#+#	# ) ) /3+	" /3#+#	# 
 
 
 
 *.u%&%8<%	-%P *.u&8<	2  /3%*xx &x ,	x
 #x 
x| /3%*|| &| ,	|
 #| 
|~ <A(
$(
48(
	(
V <A(
$(
48(
	(
\ )-(, 11 &	1
 &1 1 
!1n )-(, 55 &	5
 &5 5 
&5nmm 4m 
	m^kk 4k 
	kb #	
P
P ,
P 	
P
 

P  #	
W
W %
W 	
W
 

W?
?
 :	?

 0?
 5?
 4?
 ?
	
?
H )-a
 =A267;6:)-!a#a &a
 :a 0a 5a 4a 'a a a 
(aL	 )-o
 =A267;6:)-!o#o &o
 :o 0o 5o 4o 'o o o 
-oh	 )-B
 #+267;6:)-!B#B &B
  B 0B 5B 4B 'B B B 
BN )-C
 #+267;6:)-!C#C &C
  C 0C 5C 4C 'C C C 
CJ%&, ,r   r   c                    t        t              }| j                         D ]*  \  }}|j                  D ]  }||   j	                  |        , t        |      S )z0Index from a trigger to nodes that depend on it.)r   r   r   r   rO  rM  )r   r   r   r   triggers        r   r   r     sZ    4?4Ekkm 3
d}} 	3GW%,,T2	33  !!r   )r   r   r   r   )
__future__r   r   
concurrentconcurrent.futuresr  r  collectionsr   r   collections.abcr   r   r   r	   	functoolsr
   typingr   r   r   r   r   r   uuidr   r   langchain_core.globalsr   langchain_core.runnablesr   langchain_core.runnables.baser   r   langchain_core.runnables.configr   r   r   langchain_core.runnables.graphr   langchain_core.runnables.utilsr   r   pydanticr   typing_extensionsr   langgraph.cache.baser   langgraph.channels.baser    langgraph.checkpoint.baser!   r"   r#   r$   langgraph.constantsr%   r&   r'   r(   r)   r*   r+   r,   r-   r.   r/   r0   r1   r2   r3   r4   r5   r6   r7   r8   r9   r:   r;   r<   r=   r>   langgraph.errorsr?   r@   rA   rB   langgraph.managed.baserC   langgraph.pregel.algorD   rE   rF   rG   rH   langgraph.pregel.callrI   langgraph.pregel.checkpointrJ   rK   langgraph.pregel.debugrL   langgraph.pregel.drawrM   langgraph.pregel.iorN   rO   langgraph.pregel.looprP   rQ   rR   langgraph.pregel.managerrS   rT   langgraph.pregel.messagesrU   langgraph.pregel.protocolrV   langgraph.pregel.readrW   langgraph.pregel.retryrX   langgraph.pregel.runnerrY   langgraph.pregel.utilsrZ   langgraph.pregel.validater[   r\   langgraph.pregel.writer]   r^   langgraph.store.baser_   langgraph.typesr`   ra   rb   rc   rd   re   rf   rg   rh   langgraph.utils.configri   rj   rk   rl   rm   rn   langgraph.utils.fieldsro   langgraph.utils.pydanticrp   rq   langgraph.utils.queuerr   rs   !langchain_core.tracers._streamingrt   ImportErrorr   rv   r   r   r{   r   r   <module>r     sa   "      * F F    , 8 
 1  " *       8  4  - K 1 , 8 Q Q J ; 4 , . 0 ; C B *
 
 
  ; K 7%K 8UGVO,c12
E
 E
PF,,^ F,,RX"oZ  % $%s    G GG