From 6125ea68f55ba5156ffc7b8b9b230a8f5dd60c86 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 17 Apr 2024 09:24:27 -0700 Subject: [PATCH] WIP --- langgraph/checkpoint/base.py | 10 +-- langgraph/prebuilt/chat_agent_executor.py | 24 +++++-- langgraph/serde/__init__.py | 0 langgraph/serde/base.py | 17 +++++ langgraph/serde/jsonplus.py | 78 +++++++++++++++++++++++ 5 files changed, 115 insertions(+), 14 deletions(-) create mode 100644 langgraph/serde/__init__.py create mode 100644 langgraph/serde/base.py create mode 100644 langgraph/serde/jsonplus.py diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index a48cd2a8b..28b485845 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -8,12 +8,12 @@ from typing import ( Iterator, NamedTuple, Optional, - Protocol, TypedDict, ) from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig +from langgraph.serde.base import SerializerProtocol from langgraph.utils import StrEnum @@ -102,14 +102,6 @@ CheckpointThreadTs = ConfigurableFieldSpec( ) -class SerializerProtocol(Protocol): - def dumps(self, obj: Any) -> bytes: - ... - - def loads(self, data: bytes) -> Any: - ... - - class BaseCheckpointSaver(ABC): at: CheckpointAt = CheckpointAt.END_OF_STEP diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index fbc254b96..cd231e201 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -1,8 +1,8 @@ import json -from typing import Annotated, Sequence, TypedDict, Union +from typing import Annotated, Any, Optional, Sequence, TypedDict, Union from langchain_core.language_models import LanguageModelLike -from langchain_core.messages import BaseMessage, FunctionMessage +from langchain_core.messages import AIMessage, BaseMessage, FunctionMessage from langchain_core.runnables import RunnableLambda from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_function @@ -14,6 +14,13 @@ from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation from langgraph.prebuilt.tool_node import ToolNode +class RetriableError(Exception): + """An error that can be retried.""" + + def __init__(self, data: dict[str, Any]): + self.data = data + + # We create the AgentState that we will pass around # This simply involves a list of messages # We want steps to return messages to append to the list @@ -47,9 +54,16 @@ def create_function_calling_executor( return "continue" # Define the function that calls the model - def call_model(state: AgentState): - messages = state["messages"] - response = model.invoke(messages) + def call_model(state: AgentState, retry: Optional[RetriableError] = None): + if retry: + # TODO try to fix the invalid tool calls + pass + + # call the model + response: AIMessage = model.invoke(state["messages"]) + # If there are invalid tool calls, we raise a RetriableError + if response.invalid_tool_calls: + raise RetriableError({"invalid_tool_calls": response.invalid_tool_calls}) # We return a list, because this will get added to the existing list return {"messages": [response]} diff --git a/langgraph/serde/__init__.py b/langgraph/serde/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/langgraph/serde/base.py b/langgraph/serde/base.py new file mode 100644 index 000000000..2fbb0ab71 --- /dev/null +++ b/langgraph/serde/base.py @@ -0,0 +1,17 @@ +from typing import Any, Protocol + + +class SerializerProtocol(Protocol): + """Protocol for serialization and deserialization of objects. + + - `dumps`: Serialize an object to bytes. + - `loads`: Deserialize an object from bytes. + + Valid implementations include the `pickle`, `json` and `orjson` modules. + """ + + def dumps(self, obj: Any) -> bytes: + ... + + def loads(self, data: bytes) -> Any: + ... diff --git a/langgraph/serde/jsonplus.py b/langgraph/serde/jsonplus.py new file mode 100644 index 000000000..3df63d642 --- /dev/null +++ b/langgraph/serde/jsonplus.py @@ -0,0 +1,78 @@ +import importlib +import json +from datetime import datetime +from typing import Any, Optional +from uuid import UUID + +from langchain_core.load.load import Reviver +from langchain_core.load.serializable import Serializable +from langchain_core.pydantic_v1 import BaseModel as LcBaseModel +from pydantic import BaseModel + +from langgraph.serde.base import SerializerProtocol + +LC_REVIVER = Reviver() + + +class JsonPlusSerializer(SerializerProtocol): + def _encode_constructor_args( + self, + constructor: type[Any], + *, + method: Optional[str] = None, + args: Optional[list[Any]] = None, + kwargs: Optional[dict[str, Any]] = None, + ): + return { + "lc": 2, + "type": "constructor", + "id": [*constructor.__module__.split("."), constructor.__name__], + "method": method, + "args": args if args is not None else [], + "kwargs": kwargs if kwargs is not None else {}, + } + + def _default(self, obj): + if isinstance(obj, Serializable): + return obj.to_json() + elif isinstance(obj, (BaseModel, LcBaseModel)): + return self._encode_constructor_args(obj.__class__, kwargs=obj.dict()) + elif isinstance(obj, UUID): + return self._encode_constructor_args(UUID, args=[obj.hex]) + elif isinstance(obj, (set, frozenset)): + return self._encode_constructor_args(type(obj), args=[list(obj)]) + elif isinstance(obj, datetime): + return self._encode_constructor_args( + datetime, method="fromisoformat", args=[obj.isoformat(), obj.tzinfo] + ) + else: + raise TypeError( + f"Object of type {obj.__class__.__name__} is not JSON serializable" + ) + + def _reviver(self, value: dict[str, Any]) -> Any: + if ( + value.get("lc", None) == 2 + and value.get("type", None) == "constructor" + and value.get("id", None) is not None + ): + # Get module and class name + [*module, name] = value["id"] + # Import module + mod = importlib.import_module(".".join(module)) + # Import class + cls = getattr(mod, name) + # Instantiate class + if value["method"] is not None: + method = getattr(cls, value["method"]) + return method(*value["args"], **value["kwargs"]) + else: + return cls(*value["args"], **value["kwargs"]) + + return LC_REVIVER(value) + + def dumps(self, obj: Any) -> bytes: + return json.dumps(obj, default=self._default, sort_keys=True) + + def loads(self, data: bytes) -> Any: + return json.loads(data)