From 4e373b8af2647095e3e81034f8f192af1a273014 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sun, 31 Mar 2024 20:19:10 -0700 Subject: [PATCH] Remove noisy frames from stack traces - The less frames in library code the less confused the user is - This doesn't remove all library frames (not possible), but it's a start --- langgraph/graph/graph.py | 94 ++++++++++++++++----- langgraph/pregel/__init__.py | 36 ++++---- tests/__snapshots__/test_pregel.ambr | 120 +++++---------------------- tests/test_pregel.py | 10 +-- 4 files changed, 116 insertions(+), 144 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index efc9d8530..ea89ac3c9 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -1,13 +1,20 @@ +import asyncio import logging from collections import defaultdict -from typing import Any, Awaitable, Callable, Dict, NamedTuple, Optional, Sequence, Union +from typing import ( + Any, + Awaitable, + Callable, + Coroutine, + Dict, + NamedTuple, + Optional, + Sequence, + Union, +) from langchain_core.runnables import Runnable -from langchain_core.runnables.base import ( - RunnableLambda, - RunnableLike, - coerce_to_runnable, -) +from langchain_core.runnables.base import RunnableLike, coerce_to_runnable from langchain_core.runnables.config import RunnableConfig from langchain_core.runnables.graph import ( Graph as RunnableGraph, @@ -28,23 +35,57 @@ START = "__start__" END = "__end__" +class RunnableCallable(Runnable): + def __init__( + self, + func: Callable[..., Optional[Runnable]], + afunc: Callable[..., Awaitable[Optional[Runnable]]], + name: str, + writer: Callable[[str], Optional[Runnable]], + ) -> None: + self.name = name + self.func = func + self.afunc = afunc + self.writer = writer + + def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any: + ret = self._call_with_config(self.func, input, config, writer=self.writer) + if isinstance(ret, Runnable): + return ret.invoke(input, config) + return ret + + async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any: + ret = await self._acall_with_config( + self.afunc, input, config, writer=self.writer + ) + if isinstance(ret, Runnable): + return await ret.ainvoke(input, config) + return ret + + class Branch(NamedTuple): - condition: Runnable[Any, str] + condition: Union[Runnable[Any, str], Callable[..., str], Coroutine[Any, Any, str]] ends: Optional[dict[str, str]] def run(self, writer: Callable[[str], Optional[Runnable]]) -> None: return ChannelWrite.register_writer( - RunnableLambda( - self._route, - self._aroute, - name=self.condition.name, - ).bind(writer=writer) + RunnableCallable( + func=self._route, + afunc=self._aroute, + writer=writer, + name=self.condition.name + if isinstance(self.condition, Runnable) + else self.condition.__name__, + ) ) def _route( self, input: Any, *, writer: Callable[[str], Optional[Runnable]] ) -> Runnable: - result = self.condition.invoke(input, {"run_name": "condition"}) + if isinstance(self.condition, Runnable): + result = self.condition.invoke(input, {"run_name": "condition"}) + else: + result = self.condition(input) if self.ends: destination = self.ends[result] else: @@ -54,7 +95,12 @@ class Branch(NamedTuple): async def _aroute( self, input: Any, *, writer: Callable[[str], Optional[Runnable]] ) -> Runnable: - result = await self.condition.ainvoke(input, {"run_name": "condition"}) + if isinstance(self.condition, Runnable): + result = await self.condition.ainvoke(input, {"run_name": "condition"}) + elif asyncio.iscoroutinefunction(self.condition): + result = await self.condition(input) + else: + result = self.condition(input) if self.ends: destination = self.ends[result] else: @@ -122,6 +168,14 @@ class Graph: "Adding an edge to a graph that has already been compiled. This will " "not be reflected in the compiled graph." ) + # find a name for the condition + try: + name = ( + condition.__name__ if condition.__name__ != "" else "condition" + ) + except AttributeError: + name = "condition" + # validate the condition if start_key not in self.nodes and start_key != START: raise ValueError(f"Need to add_node `{start_key}` first") if conditional_edge_mapping and set( @@ -133,19 +187,13 @@ class Graph: f"{list(conditional_edge_mapping.values())}. Possible nodes are " f"{list(self.nodes.keys())}." ) - if not isinstance(condition, Runnable): - condition = RunnableLambda(condition) - if condition.name is None: - condition.name = "condition" - if condition.name in self.branches[start_key]: + if name in self.branches[start_key]: raise ValueError( f"Branch with name `{condition.name}` already exists for node " f"`{start_key}`" ) - - self.branches[start_key][condition.name] = Branch( - condition, conditional_edge_mapping - ) + # save it + self.branches[start_key][name] = Branch(condition, conditional_edge_mapping) def set_entry_point(self, key: str) -> None: return self.add_edge(START, key) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index e8fa02ba9..8dd8deac1 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -952,8 +952,9 @@ class Pregel( output_keys = output_keys if output_keys is not None else self.output_channels output_is_dict = not isinstance(output_keys, str) latest: Union[dict[str, Any], Any] = {} if output_is_dict else None - for chunk in self.stream( - input, + for chunk in self._transform_stream_with_config( + iter([input]), + self._transform, config, stream_mode="values", output_keys=output_keys, @@ -979,8 +980,9 @@ class Pregel( debug: Optional[bool] = None, **kwargs: Any, ) -> Iterator[Union[dict[str, Any], Any]]: - return self.transform( + return self._transform_stream_with_config( iter([input]), + self._transform, config, stream_mode=stream_mode, output_keys=output_keys, @@ -1004,7 +1006,7 @@ class Pregel( debug: Optional[bool] = None, **kwargs: Any, ) -> Iterator[Union[dict[str, Any], Any]]: - for chunk in self._transform_stream_with_config( + return self._transform_stream_with_config( input, self._transform, config, @@ -1015,8 +1017,7 @@ class Pregel( interrupt_after_nodes=interrupt_after_nodes, debug=debug, **kwargs, - ): - yield chunk + ) async def ainvoke( self, @@ -1030,11 +1031,15 @@ class Pregel( debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: + async def input_stream() -> AsyncIterator[Union[dict[str, Any], Any]]: + yield input + output_keys = output_keys if output_keys is not None else self.output_channels output_is_dict = not isinstance(output_keys, str) latest: Union[dict[str, Any], Any] = {} if output_is_dict else None - async for chunk in self.astream( - input, + async for chunk in self._atransform_stream_with_config( + input_stream(), + self._atransform, config, stream_mode="values", output_keys=output_keys, @@ -1047,7 +1052,7 @@ class Pregel( latest = {**latest, **chunk} if output_is_dict else chunk return latest - async def astream( + def astream( self, input: Union[dict[str, Any], Any], config: Optional[RunnableConfig] = None, @@ -1063,8 +1068,9 @@ class Pregel( async def input_stream() -> AsyncIterator[Union[dict[str, Any], Any]]: yield input - async for chunk in self.atransform( + return self._atransform_stream_with_config( input_stream(), + self._atransform, config, stream_mode=stream_mode, output_keys=output_keys, @@ -1073,10 +1079,9 @@ class Pregel( interrupt_after_nodes=interrupt_after_nodes, debug=debug, **kwargs, - ): - yield chunk + ) - async def atransform( + def atransform( self, input: AsyncIterator[Union[dict[str, Any], Any]], config: Optional[RunnableConfig] = None, @@ -1089,7 +1094,7 @@ class Pregel( debug: Optional[bool] = None, **kwargs: Any, ) -> AsyncIterator[Union[dict[str, Any], Any]]: - async for chunk in self._atransform_stream_with_config( + return self._atransform_stream_with_config( input, self._atransform, config, @@ -1100,8 +1105,7 @@ class Pregel( interrupt_after_nodes=interrupt_after_nodes, debug=debug, **kwargs, - ): - yield chunk + ) def _panic_or_proceed( diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 90a0378a3..874c691f7 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -51,29 +51,13 @@ }, { "id": "__start___should_start", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_start" - } + "type": "unknown", + "data": "__start___should_start" }, { "id": "left_condition", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "condition" - } + "type": "unknown", + "data": "left_condition" } ], "edges": [ @@ -235,29 +219,13 @@ }, { "id": "__start___should_start", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_start" - } + "type": "unknown", + "data": "__start___should_start" }, { "id": "left_condition", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "condition" - } + "type": "unknown", + "data": "left_condition" } ], "edges": [ @@ -385,16 +353,8 @@ }, { "id": "agent_should_continue", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_continue" - } + "type": "unknown", + "data": "agent_should_continue" } ], "edges": [ @@ -632,16 +592,8 @@ }, { "id": "agent_should_continue", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_continue" - } + "type": "unknown", + "data": "agent_should_continue" } ], "edges": [ @@ -1006,16 +958,8 @@ }, { "id": "agent_should_continue", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_continue" - } + "type": "unknown", + "data": "agent_should_continue" } ], "edges": [ @@ -1825,16 +1769,8 @@ }, { "id": "agent_should_continue", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_continue" - } + "type": "unknown", + "data": "agent_should_continue" } ], "edges": [ @@ -2068,16 +2004,8 @@ }, { "id": "agent_should_continue", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_continue" - } + "type": "unknown", + "data": "agent_should_continue" } ], "edges": [ @@ -2311,16 +2239,8 @@ }, { "id": "agent_should_continue", - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "should_continue" - } + "type": "unknown", + "data": "agent_should_continue" } ], "edges": [ diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 38c0fdef1..977ec74c1 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2269,7 +2269,7 @@ def test_message_graph( FunctionMessage( content="result for query", name="search_api", - id="00000000-0000-4000-8000-000000000019", + id="00000000-0000-4000-8000-000000000018", ), AIMessage( content="", @@ -2281,7 +2281,7 @@ def test_message_graph( FunctionMessage( content="result for another", name="search_api", - id="00000000-0000-4000-8000-000000000033", + id="00000000-0000-4000-8000-000000000031", ), AIMessage(content="answer", id="ai3"), ] @@ -2291,7 +2291,7 @@ def test_message_graph( "__start__": [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000045", + id="00000000-0000-4000-8000-000000000042", ) ] }, @@ -2308,7 +2308,7 @@ def test_message_graph( "action": FunctionMessage( content="result for query", name="search_api", - id="00000000-0000-4000-8000-000000000059", + id="00000000-0000-4000-8000-000000000055", ) }, { @@ -2324,7 +2324,7 @@ def test_message_graph( "action": FunctionMessage( content="result for another", name="search_api", - id="00000000-0000-4000-8000-000000000073", + id="00000000-0000-4000-8000-000000000068", ) }, {"agent": AIMessage(content="answer", id="ai3")},