Update workflow.py

workflow.py

@@ -2,34 +2,33 @@
 
 import time
 from datetime import datetime
-from typing import Dict, Sequence, Any
+from typing import Dict, Any, Sequence
 
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 from langgraph.graph import END, StateGraph
 from langgraph.graph.message import add_messages
-import logging
+from typing_extensions import TypedDict, Annotated
 
 from processor import EnhancedCognitiveProcessor
 from config import ResearchConfig
 
-# Define the state schema for the workflow
-from typing_extensions import TypedDict, Annotated
+import logging
+logger = logging.getLogger(__name__)
 
+# Define the state schema
 class AgentState(TypedDict):
-    messages: Annotated[Sequence[Any], add_messages]
-    context: dict
-    metadata: dict
-
-logger = logging.getLogger(__name__)
+    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
+    context: Dict[str, Any]
+    metadata: Dict[str, Any]
 
 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()
+        # Provide the state schema to the StateGraph constructor
+        self.workflow = StateGraph(AgentState)
         self._build_workflow()
         self.app = self.workflow.compile()
 
@@ -134,7 +133,8 @@ class ResearchWorkflow:
             response = self.processor.process_query(validation_prompt)
             logger.info("Output validation completed.")
             return {
-                "messages": [AIMessage(content=analysis + f"\n\nValidation: {response.get('choices', [{}])[0].get('message', {}).get('content', '')}")]
+                "messages": [AIMessage(content=analysis + f"\n\nValidation: {response.get('choices', [{}])[0].get('message', {}).get('content', '')}")],
+                "context": state["context"]
             }
         except Exception as e:
             logger.exception("Error during output validation.")
@@ -196,7 +196,7 @@ class ResearchWorkflow:
     def _error_state(self, message: str) -> Dict:
         logger.error(message)
         return {
-            "messages": [{"content": f"❌ {message}"}],
+            "messages": [AIMessage(content=f"❌ {message}")],
             "context": {"error": True},
             "metadata": {"status": "error"}
         }
@@ -215,7 +215,7 @@ class ResearchWorkflow:
                     enhanced += f"{code}\n"
                 enhanced += "```"
             return {
-                "messages": [{"content": enhanced}],
+                "messages": [AIMessage(content=enhanced)],
                 "context": state["context"]
             }
         except Exception as e: