Update app.py

app.py

@@ -10,12 +10,19 @@ from tools import tools
 from langchain_core.messages import HumanMessage
 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
 
 
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
+@tool
+def dummy_tool(query: str) -> str:
+    return f"You asked me to look up: {query}"
+
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
@@ -29,44 +36,49 @@ class BasicAgent:
 
 class LangGraphAgent:
     def __init__(self):
-        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")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype="auto",
+        )
 
-        pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, max_new_tokens=512)
-        self.llm = HuggingFacePipeline(pipeline=pipe)
+        self.pipe = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            return_full_text=False
+        )
 
-        self.graph = self._build_graph()
+        self.tools = [dummy_tool]  # Add more tools later if needed
+        self.tool_node = ToolNode(tools=self.tools)
 
-    def _build_graph(self):
-        agent_node = create_react_agent(model=self.llm, tools=tools)
-        tool_node = ToolNode(tools)
+        # LangGraph states are dicts with a "messages" key
+        builder = StateGraph()
 
-        def run_agent_node(state):
-            return {"messages": agent_node.invoke(state)["messages"]}
+        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 run_tool_node(state):
-            return tool_node.invoke({"messages": state["messages"]})
+        self.app = builder.compile()
 
-        builder = StateGraph(input_schema={"messages": list})
-        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 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()}]}
 
     def __call__(self, question: str) -> str:
-        print(f"LangGraphAgent processing: {question[:50]}...")
-        try:
-            messages = [{"role": "user", "content": question}]
-            output = self.graph.invoke({"messages": messages})
-            print("LangGraphAgent result:", output)
-            return output["messages"][-1]["content"]
-        except Exception as e:
-            print(f"LangGraphAgent error: {e}")
-            return "⚠️ Error during LangGraph agent processing."
+        result = self.app.invoke({"messages": [{"role": "user", "content": question}]})
+        messages = result["messages"]
+        return messages[-1]["content"] if messages else "❌ No response generated."
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """