handle document serialization (#335)

langchain-ai · Jun 26, 2024 · 9ec47d1 · 9ec47d1
1 parent 330d0f8
commit 9ec47d1
Show file tree

Hide file tree

Showing 3 changed files with 459 additions and 303 deletions.
diff --git a/backend/graph.py b/backend/graph.py
@@ -104,9 +104,24 @@
 COHERE_MODEL_KEY = "cohere_command"
 
 
+def update_documents(
+    _: list[Document], right: list[Document] | list[dict]
+) -> list[Document]:
+    res: list[Document] = []
+
+    for item in right:
+        if isinstance(item, dict):
+            res.append(Document(**item))
+        elif isinstance(item, Document):
+            res.append(item)
+        else:
+            raise TypeError(f"Got unknown document type '{type(item)}'")
+    return res
+
+
 class AgentState(TypedDict):
     query: str
-    documents: list[Document]
+    documents: Annotated[list[Document], update_documents]
     messages: Annotated[list[BaseMessage], add_messages]