Update tools/tool_agent.py

tools/tool_agent.py

@@ -3,25 +3,29 @@ import json
 
 class ToolCallingAgent:
     def __init__(self):
-        # Force CPU and smaller model
+        # Small CPU-friendly model
         self.model = pipeline(
             "text-generation",
-            model="cognitivecomputations/dolphin-2.1-mistral-7b",  # Smaller than llama3
-            device=-1,  # Force CPU
-            torch_dtype="float32"  # Better for CPU
+            model="gpt2",  # Replace with small model you want
+            device=-1,
+            torch_dtype="float32"
         )
         
     def generate(self, prompt, tools):
         try:
             tools_json = json.dumps(tools, ensure_ascii=False)
-            prompt = f"""Respond with JSON for one tool call. Tools: {tools_json}\nInput: {prompt}"""
+            prompt = f"""Respond ONLY with JSON for one tool call from the following list: {tools_json}\nUser input: {prompt}"""
             
             response = self.model(
                 prompt,
-                max_new_tokens=150,  # Shorter responses
-                do_sample=False  # More deterministic
+                max_new_tokens=100,
+                do_sample=False
             )
             
-            return json.loads(response[0]['generated_text'])
-        except:
-            return {"tool_name": "error", "parameters": {"message": "Failed to process request"}}
+            # Try to find JSON in output
+            text = response[0]['generated_text']
+            json_start = text.find("{")
+            json_end = text.rfind("}") + 1
+            return json.loads(text[json_start:json_end])
+        except Exception as e:
+            return {"tool_name": "error", "parameters": {"message": f"Failed to process request: {str(e)}"}}