Update workflow.py

workflow.py

@@ -2,16 +2,24 @@
 
 import time
 from datetime import datetime
-from typing import Dict
+from typing import Dict, Sequence, Any
 
 from langchain_core.messages import AIMessage, HumanMessage
 from langgraph.graph import END, StateGraph
 from langgraph.graph.message import add_messages
+import logging
 
 from processor import EnhancedCognitiveProcessor
 from config import ResearchConfig
 
-import logging
+# Define the state schema for the workflow
+from typing_extensions import TypedDict, Annotated
+
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[Any], add_messages]
+    context: dict
+    metadata: dict
+
 logger = logging.getLogger(__name__)
 
 class ResearchWorkflow:
@@ -19,8 +27,9 @@ class ResearchWorkflow:
     Defines a multi-step research workflow using a state graph.
     """
     def __init__(self) -> None:
+        # Provide the state schema to StateGraph
+        self.workflow = StateGraph(AgentState)
         self.processor = EnhancedCognitiveProcessor()
-        self.workflow = StateGraph()
         self._build_workflow()
         self.app = self.workflow.compile()