Update input/output types.

langchain-ai · Apr 18, 2024 · deab0ed · deab0ed
1 parent b1c7702
commit deab0ed
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/backend/app/agent.py b/backend/app/agent.py
@@ -8,6 +8,7 @@
     RunnableBinding,
 )
 from langgraph.checkpoint import CheckpointAt
+from langgraph.graph.message import Messages
 from langgraph.pregel import Pregel
 
 from app.agent_types.tools_agent import get_tools_agent_executor
@@ -247,8 +248,8 @@ def __init__(
         system_message=ConfigurableField(id="system_message", name="Instructions"),
     )
     .with_types(
-        input_type=Union[Sequence[AnyMessage], Dict[str, Any]],
-        output_type=Union[Sequence[AnyMessage], Dict[str, Any]],
+        input_type=Messages,
+        output_type=Sequence[AnyMessage],
     )
 )
 
@@ -312,8 +313,8 @@ def __init__(
         thread_id=ConfigurableField(id="thread_id", name="Thread ID", is_shared=True),
     )
     .with_types(
-        input_type=Union[Sequence[AnyMessage], Dict[str, Any]],
-        output_type=Union[Sequence[AnyMessage], Dict[str, Any]],
+        input_type=Dict[str, Any],
+        output_type=Dict[str, Any],
     )
 )
 
@@ -352,8 +353,8 @@ def __init__(
         chat_retrieval=chat_retrieval,
     )
     .with_types(
-        input_type=Union[Sequence[AnyMessage], Dict[str, Any]],
-        output_type=Union[Sequence[AnyMessage], Dict[str, Any]],
+        input_type=Messages,
+        output_type=Sequence[AnyMessage],
     )
 )