allow filenames to be passed through

planning_ai/nodes/map_node.py

@@ -9,20 +9,42 @@ from planning_ai.states import OverallState, SummaryState
 
 
 def generate_summary(state: SummaryState):
-    response = map_chain.invoke({"context": state["content"]})
-    return {"summaries": [response]}
+    response = map_chain.invoke(
+        {"context": state["content"], "filename": state["filename"]}
+    )
+    return {
+        "summaries": [
+            {
+                "response": response,
+                "original": state["content"],
+                "filename": state["filename"],
+            }
+        ]
+    }
 
 
 def map_summaries(state: OverallState):
     return [
-        Send("generate_summary", {"content": content}) for content in state["contents"]
+        Send("generate_summary", {"content": content, "filename": filename})
+        for content, filename in zip(state["contents"], state["filenames"])
     ]
 
 
 def collect_summaries(state: OverallState):
     return {
         "collapsed_summaries": [
-            Document(f"[{idx}]\n\n{summary}")
+            Document(
+                page_content=summary["response"].summary,
+                metadata={
+                    "stance": summary["response"].stance,
+                    "themes": summary["response"].themes,
+                    "places": summary["response"].places,
+                    "rating": summary["response"].rating,
+                    "original": summary["original"],
+                    "filename": summary["filename"],
+                    "index": idx,
+                },
+            )
             for idx, summary in enumerate(state["summaries"], start=1)
         ]
     }

planning_ai/states.py

@@ -1,11 +1,12 @@
-
 import operator
 from typing import Annotated, List, TypedDict
 
 from langchain_core.documents import Document
 
+
 class OverallState(TypedDict):
     contents: List[str]
+    filenames: List[str]
     summaries: Annotated[list, operator.add]
     collapsed_summaries: List[Document]
     final_summary: str
@@ -13,4 +14,4 @@ class OverallState(TypedDict):
 
 class SummaryState(TypedDict):
     content: str
-
+    filename: str