Update agent.py

agent.py

@@ -2,7 +2,7 @@ import os
 import pandas as pd
 from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.graph import StateGraph, MessagesState
-from langchain_community.llms import HuggingFaceHub  
+from langchain_huggingface import HuggingFaceEndpoint  
 from tools import TOOLS  
 
 # --- Read local QA data for retriever ---
@@ -14,18 +14,18 @@ qa_dict = {
 }
 
 def build_graph():
-    # Initialize Mistral model using HuggingFaceHub
-    llm = HuggingFaceHub(
-        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        model_kwargs={
-            "max_new_tokens": 512,
-            "temperature": 0.1,
-            "do_sample": False
-        },
+    # Initialize Mistral model with proper configuration
+    llm = HuggingFaceEndpoint(
+        endpoint_url="https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3",
+        task="text-generation",
+        max_new_tokens=512,
+        temperature=0.1,
+        top_k=50,
+        top_p=0.95,
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
-    # Retriever node (unchanged)
+    # Retriever node
     def retriever_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
         if query in qa_dict:
@@ -34,7 +34,7 @@ def build_graph():
         print("🔍 No match. Sending to LLM.")
         return {"messages": state["messages"]}
 
-    # Assistant node (LLM) - updated for HuggingFaceHub
+    # Assistant node (LLM)
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
 
@@ -49,19 +49,22 @@ def build_graph():
         )
         
         # Format prompt for Mistral model
-        prompt = f"<s>[INST] {system_prompt}\n\n{query} [/INST]"
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": query}
+        ]
         
         # Generate response
-        response = llm.invoke(prompt).strip()
+        response = llm.invoke(messages).strip()
 
         # Clean up response
-        for tag in ("Final answer:", "Answer:", "assistant:"):
+        for tag in ("Final answer:", "Answer:", "assistant:", "assistant:"):
             if response.lower().startswith(tag.lower()):
                 response = response[len(tag):].strip()
         
         return {"messages": [AIMessage(content=response.strip())]}
 
-    # Tool node (unchanged)
+    # Tool node
     def tool_node(state: MessagesState):
         try:
             tool_signal = state.get("tool_call", "")
@@ -82,7 +85,7 @@ def build_graph():
             print(f"⚠️ Tool error: {e}")
             return {"messages": [AIMessage(content="Unknown")]}  # fail-safe
 
-    # Build LangGraph (unchanged)
+    # Build LangGraph
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever_node)
     builder.add_node("assistant", assistant_node)
@@ -96,7 +99,7 @@ def build_graph():
 
     return builder.compile()
 
-# Agent class (unchanged)
+# Agent class
 class BasicAgent:
     def __init__(self):
         print("✅ BasicAgent initialized with retriever + LLM + tools")