Update app.py

app.py

@@ -7,12 +7,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 from langchain_community.llms import HuggingFacePipeline
 from tools import tools
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.prebuilt import ToolNode, create_react_agent
 from langgraph.graph import StateGraph, END
 from langchain.agents import tool
 from langchain_core.runnables import Runnable
 from langchain_core.tools import Tool
+from langchain_community.chat_models import ChatHuggingFace
 
 
 # (Keep Constants as is)
@@ -32,49 +33,41 @@ class BasicAgent:
 
 class LangGraphAgent:
     def __init__(self):
-        model_id = "HuggingFaceH4/zephyr-7b-beta"
-
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype="auto",
-        )
-
-        self.pipe = pipeline(
-            "text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
-            return_full_text=False
+        # Load Zephyr with correct config
+        self.model = ChatHuggingFace.from_model_id(
+            model_id="HuggingFaceH4/zephyr-7b-beta",
+            task="text-generation",
+            model_kwargs={"temperature": 0.7, "max_new_tokens": 512}
         )
 
-        self.tools = [tools]  # Add more tools later if needed
-        self.tool_node = ToolNode(tools=self.tools)
-
-        # LangGraph states are dicts with a "messages" key
+        # Define a simple graph with just one node
         builder = StateGraph()
 
-        builder.add_node("invoke_model", self.invoke_model)
-        builder.add_node("tools", self.tool_node)
-        builder.set_entry_point("invoke_model")
-        builder.add_edge("invoke_model", "tools")
-        builder.add_edge("tools", END)
+        def call_model(state):
+            messages = state.get("messages", [])
+            response = self.model.invoke(messages)
+            return {"messages": messages + [response]}
 
-        self.app = builder.compile()
+        builder.add_node("chat", call_model)
+        builder.set_entry_point("chat")
+        builder.add_edge("chat", END)
 
-    def invoke_model(self, state: dict) -> dict:
-        messages = state["messages"]
-        if isinstance(messages, str):
-            messages = [{"role": "user", "content": messages}]
-        prompt = self.tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        response = self.pipe(prompt, max_new_tokens=256, temperature=0.7)[0]["generated_text"]
-        return {"messages": messages + [{"role": "assistant", "content": response.strip()}]}
+        # Compile the graph
+        self.graph = builder.compile()
 
     def __call__(self, question: str) -> str:
-        result = self.app.invoke({"messages": [{"role": "user", "content": question}]})
-        messages = result["messages"]
-        return messages[-1]["content"] if messages else "❌ No response generated."
+        # Wrap input in HumanMessage format
+        result = self.graph.invoke({
+            "messages": [HumanMessage(content=question)]
+        })
+
+        # Extract final response
+        messages = result.get("messages", [])
+        for msg in reversed(messages):
+            if isinstance(msg, AIMessage):
+                return msg.content
+
+        return "❌ No response generated."
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """