[docs]classLabelsDict(TypedDict):"""Dictionary of labels for nodes and edges in a graph."""nodes:dict[str,str]"""Labels for nodes."""edges:dict[str,str]"""Labels for edges."""
[docs]defis_uuid(value:str)->bool:"""Check if a string is a valid UUID. Args: value: The string to check. Returns: True if the string is a valid UUID, False otherwise. """try:UUID(value)returnTrueexceptValueError:returnFalse
[docs]classEdge(NamedTuple):"""Edge in a graph. Parameters: source: The source node id. target: The target node id. data: Optional data associated with the edge. Defaults to None. conditional: Whether the edge is conditional. Defaults to False. """source:strtarget:strdata:Optional[Stringifiable]=Noneconditional:bool=False
[docs]defcopy(self,*,source:Optional[str]=None,target:Optional[str]=None)->Edge:"""Return a copy of the edge with optional new source and target nodes. Args: source: The new source node id. Defaults to None. target: The new target node id. Defaults to None. Returns: A copy of the edge with the new source and target nodes. """returnEdge(source=sourceorself.source,target=targetorself.target,data=self.data,conditional=self.conditional,)
[docs]classNode(NamedTuple):"""Node in a graph. Parameters: id: The unique identifier of the node. name: The name of the node. data: The data of the node. metadata: Optional metadata for the node. Defaults to None. """id:strname:strdata:Union[Type[BaseModel],RunnableType]metadata:Optional[Dict[str,Any]]
[docs]defcopy(self,*,id:Optional[str]=None,name:Optional[str]=None)->Node:"""Return a copy of the node with optional new id and name. Args: id: The new node id. Defaults to None. name: The new node name. Defaults to None. Returns: A copy of the node with the new id and name. """returnNode(id=idorself.id,name=nameorself.name,data=self.data,metadata=self.metadata,)
[docs]classBranch(NamedTuple):"""Branch in a graph. Parameters: condition: A callable that returns a string representation of the condition. ends: Optional dictionary of end node ids for the branches. Defaults to None. """condition:Callable[...,str]ends:Optional[dict[str,str]]
[docs]classCurveStyle(Enum):"""Enum for different curve styles supported by Mermaid"""BASIS="basis"BUMP_X="bumpX"BUMP_Y="bumpY"CARDINAL="cardinal"CATMULL_ROM="catmullRom"LINEAR="linear"MONOTONE_X="monotoneX"MONOTONE_Y="monotoneY"NATURAL="natural"STEP="step"STEP_AFTER="stepAfter"STEP_BEFORE="stepBefore"
[docs]@dataclassclassNodeStyles:"""Schema for Hexadecimal color codes for different node types. Parameters: default: The default color code. Defaults to "fill:#f2f0ff,line-height:1.2". first: The color code for the first node. Defaults to "fill-opacity:0". last: The color code for the last node. Defaults to "fill:#bfb6fc". """default:str="fill:#f2f0ff,line-height:1.2"first:str="fill-opacity:0"last:str="fill:#bfb6fc"
[docs]classMermaidDrawMethod(Enum):"""Enum for different draw methods supported by Mermaid"""PYPPETEER="pyppeteer"# Uses Pyppeteer to render the graphAPI="api"# Uses Mermaid.INK API to render the graph
[docs]defnode_data_str(id:str,data:Union[Type[BaseModel],RunnableType])->str:"""Convert the data of a node to a string. Args: id: The node id. data: The node data. Returns: A string representation of the data. """fromlangchain_core.runnables.baseimportRunnableifnotis_uuid(id):returnidelifisinstance(data,Runnable):data_str=data.get_name()else:data_str=data.__name__returndata_strifnotdata_str.startswith("Runnable")elsedata_str[8:]
[docs]defnode_data_json(node:Node,*,with_schemas:bool=False)->Dict[str,Union[str,Dict[str,Any]]]:"""Convert the data of a node to a JSON-serializable format. Args: node: The node to convert. with_schemas: Whether to include the schema of the data if it is a Pydantic model. Defaults to False. Returns: A dictionary with the type of the data and the data itself. """fromlangchain_core.load.serializableimportto_json_not_implementedfromlangchain_core.runnables.baseimportRunnable,RunnableSerializableifisinstance(node.data,RunnableSerializable):json:Dict[str,Any]={"type":"runnable","data":{"id":node.data.lc_id(),"name":node_data_str(node.id,node.data),},}elifisinstance(node.data,Runnable):json={"type":"runnable","data":{"id":to_json_not_implemented(node.data)["id"],"name":node_data_str(node.id,node.data),},}elifinspect.isclass(node.data)andis_basemodel_subclass(node.data):json=({"type":"schema","data":node.data.schema(),}ifwith_schemaselse{"type":"schema","data":node_data_str(node.id,node.data),})else:json={"type":"unknown","data":node_data_str(node.id,node.data),}ifnode.metadataisnotNone:json["metadata"]=node.metadatareturnjson
[docs]@dataclassclassGraph:"""Graph of nodes and edges. Parameters: nodes: Dictionary of nodes in the graph. Defaults to an empty dictionary. edges: List of edges in the graph. Defaults to an empty list. """nodes:Dict[str,Node]=field(default_factory=dict)edges:List[Edge]=field(default_factory=list)
[docs]defto_json(self,*,with_schemas:bool=False)->Dict[str,List[Dict[str,Any]]]:"""Convert the graph to a JSON-serializable format. Args: with_schemas: Whether to include the schemas of the nodes if they are Pydantic models. Defaults to False. Returns: A dictionary with the nodes and edges of the graph. """stable_node_ids={node.id:iifis_uuid(node.id)elsenode.idfori,nodeinenumerate(self.nodes.values())}edges:List[Dict[str,Any]]=[]foredgeinself.edges:edge_dict={"source":stable_node_ids[edge.source],"target":stable_node_ids[edge.target],}ifedge.dataisnotNone:edge_dict["data"]=edge.dataifedge.conditional:edge_dict["conditional"]=Trueedges.append(edge_dict)return{"nodes":[{"id":stable_node_ids[node.id],**node_data_json(node,with_schemas=with_schemas),}fornodeinself.nodes.values()],"edges":edges,}
def__bool__(self)->bool:returnbool(self.nodes)
[docs]defnext_id(self)->str:"""Return a new unique node identifier that can be used to add a node to the graph."""returnuuid4().hex
[docs]defadd_node(self,data:Union[Type[BaseModel],RunnableType],id:Optional[str]=None,*,metadata:Optional[Dict[str,Any]]=None,)->Node:"""Add a node to the graph and return it. Args: data: The data of the node. id: The id of the node. Defaults to None. metadata: Optional metadata for the node. Defaults to None. Returns: The node that was added to the graph. Raises: ValueError: If a node with the same id already exists. """ifidisnotNoneandidinself.nodes:raiseValueError(f"Node with id {id} already exists")id=idorself.next_id()node=Node(id=id,data=data,metadata=metadata,name=node_data_str(id,data))self.nodes[node.id]=nodereturnnode
[docs]defremove_node(self,node:Node)->None:"""Remove a node from the graph and all edges connected to it. Args: node: The node to remove. """self.nodes.pop(node.id)self.edges=[edgeforedgeinself.edgesifedge.source!=node.idandedge.target!=node.id]
[docs]defadd_edge(self,source:Node,target:Node,data:Optional[Stringifiable]=None,conditional:bool=False,)->Edge:"""Add an edge to the graph and return it. Args: source: The source node of the edge. target: The target node of the edge. data: Optional data associated with the edge. Defaults to None. conditional: Whether the edge is conditional. Defaults to False. Returns: The edge that was added to the graph. Raises: ValueError: If the source or target node is not in the graph. """ifsource.idnotinself.nodes:raiseValueError(f"Source node {source.id} not in graph")iftarget.idnotinself.nodes:raiseValueError(f"Target node {target.id} not in graph")edge=Edge(source=source.id,target=target.id,data=data,conditional=conditional)self.edges.append(edge)returnedge
[docs]defextend(self,graph:Graph,*,prefix:str="")->Tuple[Optional[Node],Optional[Node]]:"""Add all nodes and edges from another graph. Note this doesn't check for duplicates, nor does it connect the graphs. Args: graph: The graph to add. prefix: The prefix to add to the node ids. Defaults to "". Returns: A tuple of the first and last nodes of the subgraph. """ifall(is_uuid(node.id)fornodeingraph.nodes.values()):prefix=""defprefixed(id:str)->str:returnf"{prefix}:{id}"ifprefixelseid# prefix each nodeself.nodes.update({prefixed(k):v.copy(id=prefixed(k))fork,vingraph.nodes.items()})# prefix each edge's source and targetself.edges.extend([edge.copy(source=prefixed(edge.source),target=prefixed(edge.target))foredgeingraph.edges])# return (prefixed) first and last nodes of the subgraphfirst,last=graph.first_node(),graph.last_node()return(first.copy(id=prefixed(first.id))iffirstelseNone,last.copy(id=prefixed(last.id))iflastelseNone,)
[docs]defreid(self)->Graph:"""Return a new graph with all nodes re-identified, using their unique, readable names where possible."""node_labels={node.id:node.namefornodeinself.nodes.values()}node_label_counts=Counter(node_labels.values())def_get_node_id(node_id:str)->str:label=node_labels[node_id]ifis_uuid(node_id)andnode_label_counts[label]==1:returnlabelelse:returnnode_idreturnGraph(nodes={_get_node_id(id):node.copy(id=_get_node_id(id))forid,nodeinself.nodes.items()},edges=[edge.copy(source=_get_node_id(edge.source),target=_get_node_id(edge.target),)foredgeinself.edges],)
[docs]deffirst_node(self)->Optional[Node]:"""Find the single node that is not a target of any edge. If there is no such node, or there are multiple, return None. When drawing the graph, this node would be the origin."""return_first_node(self)
[docs]deflast_node(self)->Optional[Node]:"""Find the single node that is not a source of any edge. If there is no such node, or there are multiple, return None. When drawing the graph, this node would be the destination."""return_last_node(self)
[docs]deftrim_first_node(self)->None:"""Remove the first node if it exists and has a single outgoing edge, i.e., if removing it would not leave the graph without a "first" node."""first_node=self.first_node()iffirst_nodeand_first_node(self,exclude=[first_node.id]):self.remove_node(first_node)
[docs]deftrim_last_node(self)->None:"""Remove the last node if it exists and has a single incoming edge, i.e., if removing it would not leave the graph without a "last" node."""last_node=self.last_node()iflast_nodeand_last_node(self,exclude=[last_node.id]):self.remove_node(last_node)
[docs]defdraw_ascii(self)->str:"""Draw the graph as an ASCII art string."""fromlangchain_core.runnables.graph_asciiimportdraw_asciireturndraw_ascii({node.id:node.namefornodeinself.nodes.values()},self.edges,)
[docs]defprint_ascii(self)->None:"""Print the graph as an ASCII art string."""print(self.draw_ascii())# noqa: T201
[docs]defdraw_png(self,output_file_path:Optional[str]=None,fontname:Optional[str]=None,labels:Optional[LabelsDict]=None,)->Union[bytes,None]:"""Draw the graph as a PNG image. Args: output_file_path: The path to save the image to. If None, the image is not saved. Defaults to None. fontname: The name of the font to use. Defaults to None. labels: Optional labels for nodes and edges in the graph. Defaults to None. Returns: The PNG image as bytes if output_file_path is None, None otherwise. """fromlangchain_core.runnables.graph_pngimportPngDrawerdefault_node_labels={node.id:node.namefornodeinself.nodes.values()}returnPngDrawer(fontname,LabelsDict(nodes={**default_node_labels,**(labels["nodes"]iflabelsisnotNoneelse{}),},edges=labels["edges"]iflabelsisnotNoneelse{},),).draw(self,output_file_path)
[docs]defdraw_mermaid(self,*,with_styles:bool=True,curve_style:CurveStyle=CurveStyle.LINEAR,node_colors:Optional[NodeStyles]=None,wrap_label_n_words:int=9,)->str:"""Draw the graph as a Mermaid syntax string. Args: with_styles: Whether to include styles in the syntax. Defaults to True. curve_style: The style of the edges. Defaults to CurveStyle.LINEAR. node_colors: The colors of the nodes. Defaults to NodeStyles(). wrap_label_n_words: The number of words to wrap the node labels at. Defaults to 9. Returns: The Mermaid syntax string. """fromlangchain_core.runnables.graph_mermaidimportdraw_mermaidgraph=self.reid()first_node=graph.first_node()last_node=graph.last_node()returndraw_mermaid(nodes=graph.nodes,edges=graph.edges,first_node=first_node.idiffirst_nodeelseNone,last_node=last_node.idiflast_nodeelseNone,with_styles=with_styles,curve_style=curve_style,node_styles=node_colors,wrap_label_n_words=wrap_label_n_words,)
[docs]defdraw_mermaid_png(self,*,curve_style:CurveStyle=CurveStyle.LINEAR,node_colors:Optional[NodeStyles]=None,wrap_label_n_words:int=9,output_file_path:Optional[str]=None,draw_method:MermaidDrawMethod=MermaidDrawMethod.API,background_color:str="white",padding:int=10,)->bytes:"""Draw the graph as a PNG image using Mermaid. Args: curve_style: The style of the edges. Defaults to CurveStyle.LINEAR. node_colors: The colors of the nodes. Defaults to NodeStyles(). wrap_label_n_words: The number of words to wrap the node labels at. Defaults to 9. output_file_path: The path to save the image to. If None, the image is not saved. Defaults to None. draw_method: The method to use to draw the graph. Defaults to MermaidDrawMethod.API. background_color: The color of the background. Defaults to "white". padding: The padding around the graph. Defaults to 10. Returns: The PNG image as bytes. """fromlangchain_core.runnables.graph_mermaidimportdraw_mermaid_pngmermaid_syntax=self.draw_mermaid(curve_style=curve_style,node_colors=node_colors,wrap_label_n_words=wrap_label_n_words,)returndraw_mermaid_png(mermaid_syntax=mermaid_syntax,output_file_path=output_file_path,draw_method=draw_method,background_color=background_color,padding=padding,)
def_first_node(graph:Graph,exclude:Sequence[str]=())->Optional[Node]:"""Find the single node that is not a target of any edge. Exclude nodes/sources with ids in the exclude list. If there is no such node, or there are multiple, return None. When drawing the graph, this node would be the origin."""targets={edge.targetforedgeingraph.edgesifedge.sourcenotinexclude}found:List[Node]=[]fornodeingraph.nodes.values():ifnode.idnotinexcludeandnode.idnotintargets:found.append(node)returnfound[0]iflen(found)==1elseNonedef_last_node(graph:Graph,exclude:Sequence[str]=())->Optional[Node]:"""Find the single node that is not a source of any edge. Exclude nodes/targets with ids in the exclude list. If there is no such node, or there are multiple, return None. When drawing the graph, this node would be the destination."""sources={edge.sourceforedgeingraph.edgesifedge.targetnotinexclude}found:List[Node]=[]fornodeingraph.nodes.values():ifnode.idnotinexcludeandnode.idnotinsources:found.append(node)returnfound[0]iflen(found)==1elseNone