Update tools/tool_agent.py

tools/tool_agent.py

@@ -10,22 +10,36 @@ class ToolCallingAgent:
         )
         
     def generate(self, prompt, tools):
-        # Format the tools specification
         tools_json = json.dumps(tools, ensure_ascii=False)
-        
-        # Create the tool-calling prompt
         system_msg = f"""You are an AI assistant that can call tools.
-        Available tools: {tools_json}
-        Respond with JSON containing 'tool_name' and 'parameters'."""
-        
-        # Generate the response
+Available tools: {tools_json}
+Respond ONLY with a valid JSON containing keys 'tool_name' and 'parameters'."""
+
+        # Construct prompt with system and user tokens (assuming model supports these)
+        full_prompt = f"<|system|>{system_msg}</s><|user|>{prompt}</s>"
+
         response = self.model(
-            f"<|system|>{system_msg}</s><|user|>{prompt}</s>",
+            full_prompt,
             max_new_tokens=200,
-            do_sample=True
+            do_sample=False  # deterministic output for better JSON consistency
         )
-        
+
+        text = response[0]['generated_text']
+
+        # Extract JSON substring between first '{' and last '}'
+        json_start = text.find("{")
+        json_end = text.rfind("}") + 1
+        if json_start == -1 or json_end == -1:
+            return {"error": "No JSON found in model output", "raw_output": text}
+
+        json_text = text[json_start:json_end]
+
         try:
-            return json.loads(response[0]['generated_text'])
-        except:
-            return {"error": "Failed to parse tool call"}
+            return json.loads(json_text)
+        except json.JSONDecodeError as e:
+            return {
+                "error": "Failed to parse JSON",
+                "message": str(e),
+                "raw_output": text,
+                "extracted_json": json_text
+            }