Fix graph state document typing and defer app import

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() 
+    main() 

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() 
+app = workflow.compile()