Update app.py

app.py

@@ -44,12 +44,25 @@ class ZephyrAPI:
         return "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
 
     def __call__(self, question: str) -> str:
-        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"
+        prompt = f"""<|system|>
+                You are a smart agent that can reason step-by-step and use tools when necessary.
+                
+                Available tools:
+                {self.tool_descriptions}
+                
+                Respond with this format:
+                Thought: do I need to use a tool?
+                Action: tool_name
+                Action Input: "input"
+                Observation: tool result
+                ... (repeat if needed)
+                Final Answer: your final answer to the user
+                
+                <|user|>
+                {question}
+                
+                <|assistant|>
+                {scratchpad}"""
         payload = {
             "inputs": prompt,
             "parameters": {
@@ -86,25 +99,27 @@ class LangGraphAgent:
             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)
+            match = re.search(r"Action:\s*(\w+)\s*Action Input:\s*\"(.+?)\"", response, re.DOTALL)
             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}")
+                        tool_output = tool_fn(tool_input)
+                        # Append tool call and result to scratchpad
+                        scratchpad = f"{response}\nObservation: {tool_output}"
+                        follow_up = self.model(content, scratchpad)
                         return {"messages": messages + [
-                            AIMessage(content=raw_response),
-                            AIMessage(content=tool_output),
+                            AIMessage(content=response),
+                            AIMessage(content=f"Observation: {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)]}
+        
+            return {"messages": messages + [AIMessage(content=response)]}
 
         builder.add_node("chat", call_model)
         builder.set_entry_point("chat")