WIP

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
 

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]}
 

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:
+        ...

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)