Update app.py

app.py

@@ -37,10 +37,19 @@ class ZephyrAPI:
         self.headers = {
             "Authorization": f"Bearer {os.getenv('HF_TOKEN')}"
         }
+        self.tool_descriptions = self.format_tools(tools or [])
         print("ZephyrAPI initialized using Inference API.")
 
+    def format_tools(self, tools):
+        return "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
+
     def __call__(self, question: str) -> str:
-        prompt = f"<|system|>\nYou are a helpful assistant.\n<|user|>\n{question}\n<|assistant|>\n"
+        prompt = f"<|system|>\n"
+                 f"You are a helpful AI agent. You can use the following tools when needed:\n"
+                 f"{self.tool_descriptions}\n"
+                 f"\nRespond to the user. If a tool is needed, use this format:\n"
+                 f"Action: tool_name\nAction Input: input_for_tool\n"
+                 f"\n<|user|>\n{question}\n<|assistant|>\n"
         payload = {
             "inputs": prompt,
             "parameters": {
@@ -61,8 +70,9 @@ class ZephyrAPI:
 
 
 class LangGraphAgent:
-    def __init__(self):
-        self.model = ZephyrAPI()
+    def __init__(self, tools=None):
+        self.tools = {tool.name: tool for tool in tools} if tools else {}
+        self.model = ZephyrAPI(tools=tools)
 
         builder = StateGraph(dict)
 
@@ -72,8 +82,29 @@ class LangGraphAgent:
             if not user_msg:
                 return {"messages": messages + [AIMessage(content="❌ No user input found.")]}
             
-            response = self.model(user_msg.content)
-            return {"messages": messages + [AIMessage(content=response)]}
+            content = user_msg.content.strip()
+            raw_response = self.model(content)
+
+            # Check if model issued a tool call
+            match = re.search(r"Action:\s*(\w+)\s*Action Input:\s*(.+)", raw_response, re.IGNORECASE)
+            if match:
+                tool_name, tool_input = match.groups()
+                tool_fn = self.tools.get(tool_name)
+                if tool_fn:
+                    try:
+                        tool_output = tool_fn(tool_input.strip('"'))
+                        follow_up = self.model(f"User asked: {content}\nTool [{tool_name}] returned: {tool_output}")
+                        return {"messages": messages + [
+                            AIMessage(content=raw_response),
+                            AIMessage(content=tool_output),
+                            AIMessage(content=follow_up),
+                        ]}
+                    except Exception as e:
+                        return {"messages": messages + [AIMessage(content=f"⚠️ Tool error: {e}")]}
+                else:
+                    return {"messages": messages + [AIMessage(content=f"⚠️ Unknown tool: {tool_name}")]}
+
+            return {"messages": messages + [AIMessage(content=raw_response)]}
 
         builder.add_node("chat", call_model)
         builder.set_entry_point("chat")
@@ -254,7 +285,7 @@ with gr.Blocks() as demo:
 
     def test_agent_response(question: str) -> str:
         # agent = BasicAgent()
-        agent = LangGraphAgent()
+        agent = LangGraphAgent(tools=tools)
         print(agent("What's the capital of France?"))
         return agent(question)