update _assistant_node

agent.py

@@ -266,20 +266,27 @@ class AgentState(TypedDict):
     input_file: Optional[List[str]]  # Contains file path (PDF/PNG)
     messages: Annotated[List[BaseMessage], add_messages]
 
+
 # === Agent Class ===
 class MyAgent:
     def __init__(
         self,
-        model_name: str = "anthropic:claude-3-5-sonnet-latest",
+        model_name: str = "anthropic:claude-3-5-sonnet-latest",  # <-- Use a valid model name
         temperature: float = 0.0
     ):
-        # Initialize LLM
-        self.llm = init_chat_model(model_name, temperature=temperature)
-        # Base tools
-        self.tools = tools
-        # RAG components
-        self.docs: List[Any] = []
-        self.retriever: Optional[BM25Retriever] = None
+        try:
+            self.llm = init_chat_model(
+                model_name,
+                temperature=temperature
+                )
+            # Base tools
+            self.tools = tools
+            # RAG components
+            self.docs: List[Any] = []
+            self.retriever: Optional[BM25Retriever] = None
+        except Exception as e:
+            print(f"Error initializing LLM: {e}")
+            raise
 
     def add_files(self, file_paths: List[str]):
         """
@@ -391,23 +398,25 @@ class MyAgent:
             if file_paths:
                 state["input_file"] = file_paths
 
-            # Build graph
+            # Build graph with proper conditional edge to prevent loops
             builder = StateGraph(dict)
             builder.add_node("assistant", self._assistant_node)
             builder.add_node("tools", ToolNode(self.tools))
             builder.add_edge(START, "assistant")
             
-            # Always allow the assistant to hand off to the tools node
+            # Fix conditional edges with better check
             builder.add_conditional_edges(
                 "assistant",
-                lambda s: any(t.name in s["messages"][-1].content for t in self.tools),
-                {True: "tools", False: "assistant"}
+                tools_condition,  # Use built-in tools_condition
+                "tools"
             )
             builder.add_edge("tools", "assistant")
+            
+            # Add recursion_limit to prevent infinite loops
             graph = builder.compile()
-
-            # Use invoke() instead of run()
-            out = graph.invoke(state)
+            
+            # Use invoke() with higher recursion limit
+            out = graph.invoke(state, {"recursion_limit": 10})  # Lower limit
             last_message = out["messages"][-1].content
             
             # Extract only the FINAL ANSWER part
@@ -423,19 +432,31 @@ class MyAgent:
     def _assistant_node(self, state: dict) -> dict:
         """Process messages with the LLM."""
         try:
-            # Check if messages exist
-            if not state["messages"]:
+            # Check if messages exist and ensure proper format
+            if not state.get("messages") or len(state["messages"]) == 0:
                 # Add a system message if empty
-                state["messages"].append(SystemMessage(content=SYSTEM_PROMPT))
-                
+                state["messages"] = [SystemMessage(content=SYSTEM_PROMPT)]
+        
+            # Ensure we have at least a system and user message
+            has_system = any(isinstance(m, SystemMessage) for m in state["messages"])
+            has_human = any(isinstance(m, HumanMessage) for m in state["messages"])
+        
+            if not has_system:
+                state["messages"].insert(0, SystemMessage(content=SYSTEM_PROMPT))
+        
+            if not has_human:
+                state["messages"].append(HumanMessage(content="Hello"))
+            
             # Invoke the chat model with our BaseMessage list
-            resp = self.llm(state["messages"])
+            resp = self.llm.invoke(state["messages"])
             state["messages"].append(resp)
             return state
         except Exception as e:
-            # Handle errors by adding an error message
             error_msg = f"Error calling LLM: {str(e)}"
             print(error_msg)
+            print(f"Message count: {len(state.get('messages', []))}")
+            if state.get("messages"):
+                print(f"Message types: {[type(m).__name__ for m in state['messages']]}")
             return state