fix: remove score from state

planning_ai/nodes/map_node.py

@@ -25,10 +25,9 @@ def retrieve_themes(state: DocumentState) -> DocumentState:
     except Exception as e:
         logger.error(f"Theme selection error: {e}")
         themes = []
-    state["themes"] = themes
-    state["themes"] = [d for d in state["themes"] if d["score"] > 2]
-    state["score"] = np.mean([theme["score"] for theme in state["themes"]])
-    if state["score"] < 3:
+    state["themes"] = [d for d in themes if d["score"] > 2]
+    score = np.mean([theme["score"] for theme in state["themes"]])
+    if score < 3:
         state["processed"] = True
         state["failed"] = True
 

planning_ai/states.py

@@ -6,6 +6,7 @@ from langchain_core.documents import Document
 from pydantic import BaseModel
 
 from planning_ai.chains.hallucination_chain import HallucinationChecker
+from planning_ai.chains.themes_chain import ThemeScore
 from planning_ai.common.utils import filename_reducer
 
 
@@ -14,7 +15,7 @@ class DocumentState(TypedDict):
     filename: int
 
     entities: list[dict]
-    themes: set[str]
+    themes: set[ThemeScore]
 
     summary: BaseModel
     hallucination: HallucinationChecker