From c29e3e44023694ff6e335495f2aafda6b8c2ee1a Mon Sep 17 00:00:00 2001
From: ofeye <117448686+ofeye@users.noreply.github.com>
Date: Fri, 8 Aug 2025 11:52:58 +0300
Subject: [PATCH] fix: correct graph state types and defer app import

---
 main.py            | 10 ++++------
 src/graph/agent.py |  7 ++++---
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/main.py b/main.py
index 9e362dd..9c87797 100644
--- a/main.py
+++ b/main.py
@@ -1,16 +1,14 @@
 from langchain_core.messages import HumanMessage
 from dotenv import load_dotenv
 
-# Load environment variables from .env file before anything else
-load_dotenv()
-
-from src.graph.agent import app
-
 
 def main():
     """
     Main function to run the agent in a loop.
     """
+    load_dotenv()
+    from src.graph.agent import app
+
     print("Welcome to the LangGraph Agent! Type 'exit' or 'quit' to end the conversation.")
     while True:
         user_input = input("You: ")
@@ -31,4 +29,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main() 
\ No newline at end of file
+    main() 
diff --git a/src/graph/agent.py b/src/graph/agent.py
index 3d69783..c1f82ce 100644
--- a/src/graph/agent.py
+++ b/src/graph/agent.py
@@ -1,6 +1,7 @@
-from typing import List, TypedDict, Annotated
+from typing import TypedDict, Annotated
 import operator
 from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
+from langchain_core.documents import Document
 from langgraph.graph import StateGraph, END
 from langgraph.prebuilt import ToolNode # Use ToolNode instead of ToolExecutor
 from langchain_groq import ChatGroq
@@ -29,7 +30,7 @@ class GraphState(TypedDict):
         retrieved_docs: A list of documents retrieved from the vector store.
     """
     messages: Annotated[list[BaseMessage], operator.add]
-    retrieved_docs: list[str]
+    retrieved_docs: list[Document]
 
 
 def retrieve_from_memory(state: GraphState):
@@ -149,4 +150,4 @@ def should_call_tool(state: GraphState):
 workflow.add_edge("updater", END)
 
 # Compile the graph into a runnable app
-app = workflow.compile() 
\ No newline at end of file
+app = workflow.compile()