Update app.py

app.py

@@ -5,11 +5,11 @@ import inspect
 import pandas as pd
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_community.llms import HuggingFacePipeline
 from tools import tools
 from langchain_core.messages import HumanMessage
-from langgraph.prebuilt import ToolExecutor, chat_agent_executor
+from langgraph.prebuilt import ToolNode, create_react_agent
+from langgraph.graph import StateGraph, END
 
 
 # (Keep Constants as is)
@@ -27,28 +27,47 @@ class BasicAgent:
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer
 
-class ZephyrAgent:
+class LangGraphAgent:
     def __init__(self):
-        print("Initializing local Zephyr model with tools...")
+        print("Initializing LangGraphAgent...")
+        model_id = "HuggingFaceH4/zephyr-7b-beta"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 
-        tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
-        model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/zephyr-7b-beta", torch_dtype=torch.float16, device_map="auto")
+        pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, max_new_tokens=512)
+        self.llm = HuggingFacePipeline(pipeline=pipe)
 
-        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
-        llm = HuggingFacePipeline(pipeline=pipe)
+        self.graph = self._build_graph()
 
-        # Wrap tools
-        tool_executor = ToolExecutor(tools)
-        self.agent_executor = chat_agent_executor.create_chat_agent_executor(llm=llm, tools=tools)
+    def _build_graph(self):
+        agent_node = create_react_agent(model=self.llm, tools=tools)
+        tool_node = ToolNode(tools)
+
+        def run_agent_node(state):
+            question = state["question"]
+            messages = [{"role": "user", "content": question}]
+            result = agent_node.invoke({"messages": messages})
+            return {"messages": result["messages"]}
+
+        def run_tool_node(state):
+            return tool_node.invoke({"messages": state["messages"]})
+
+        builder = StateGraph(input_schema={"question": str})
+        builder.add_node("agent", run_agent_node)
+        builder.add_node("tools", run_tool_node)
+        builder.set_entry_point("agent")
+        builder.add_edge("agent", "tools")
+        builder.add_edge("tools", END)
+        return builder.compile()
 
     def __call__(self, question: str) -> str:
+        print(f"LangGraphAgent processing: {question[:50]}...")
         try:
-            message = HumanMessage(content=question)
-            response = self.agent_executor.invoke({"messages": [message]})
-            return response.content
+            output = self.graph.invoke({"question": question})
+            return output["messages"][-1].content
         except Exception as e:
-            print(f"Tool-augmented ZephyrAgent error: {e}")
-            return "⚠️ Agent failed to answer due to tool or model error."
+            print(f"LangGraphAgent error: {e}")
+            return "⚠️ Error during LangGraph agent processing."
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
@@ -211,7 +230,7 @@ with gr.Blocks() as demo:
 
     def test_agent_response(question: str) -> str:
         # agent = BasicAgent()
-        agent = ZephyrAgent()
+        agent = LangGraphAgent()
         return agent(question)
 
     test_button.click(fn=test_agent_response, inputs=question_input, outputs=answer_output)