Update agent.py

agent.py

@@ -1,4 +1,3 @@
-"""LangGraph agent using retriever fallback to Qwen2.5-Coder-32B-Instruct (no tools)."""
 import os
 import pandas as pd
 from langchain_core.messages import HumanMessage, AIMessage
@@ -13,17 +12,15 @@ qa_dict = {
     for _, row in qa_pairs.iterrows()
 }
 
-# --- Define LangGraph with fallback to LLM ---
 def build_graph():
     """Construct a LangGraph agent with retriever and fallback LLM."""
-    
-    # Initialize HuggingFace Qwen model as fallback
+
     llm = HuggingFaceEndpoint(
         repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
+        task="chat-completion",  
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
-    # Node: Retriever
     def retriever_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
         if query in qa_dict:
@@ -32,13 +29,11 @@ def build_graph():
         print("🔍 No match found. Falling back to LLM.")
         return {"messages": state["messages"]}
 
-    # Node: Fallback LLM (Qwen)
     def assistant_node(state: MessagesState):
-        query = state["messages"][-1].content.strip()
-        response = llm.invoke(query)
-        return {"messages": [AIMessage(content=response)]}
+        # <--- Use invoke with list of messages to get correct "chat" format
+        result = llm.invoke(state["messages"])
+        return {"messages": [result]}
 
-    # Build graph
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever_node)
     builder.add_node("assistant", assistant_node)