From deab0ed8ebf5fc3cd7890c736e92e10f641bfaa9 Mon Sep 17 00:00:00 2001
From: Andrew Nguonly <andrewnguonly@gmail.com>
Date: Thu, 18 Apr 2024 11:55:18 -0700
Subject: [PATCH] Update input/output types.

---
 backend/app/agent.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/backend/app/agent.py b/backend/app/agent.py
index 5b688665..2658c9b7 100644
--- 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],
     )
 )