Update tools/tool_agent.py

tools/tool_agent.py

@@ -3,29 +3,25 @@ import json
 
 class ToolCallingAgent:
     def __init__(self):
+        # Force CPU and smaller model
         self.model = pipeline(
             "text-generation",
-            model="cognitivecomputations/dolphin-2.9-llama3-8b",
-            device_map="auto"
+            model="cognitivecomputations/dolphin-2.1-mistral-7b",  # Smaller than llama3
+            device=-1,  # Force CPU
+            torch_dtype="float32"  # Better for CPU
         )
         
     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
-        response = self.model(
-            f"<|system|>{system_msg}</s><|user|>{prompt}</s>",
-            max_new_tokens=200,
-            do_sample=True
-        )
-        
         try:
+            tools_json = json.dumps(tools, ensure_ascii=False)
+            prompt = f"""Respond with JSON for one tool call. Tools: {tools_json}\nInput: {prompt}"""
+            
+            response = self.model(
+                prompt,
+                max_new_tokens=150,  # Shorter responses
+                do_sample=False  # More deterministic
+            )
+            
             return json.loads(response[0]['generated_text'])
         except:
-            return {"error": "Failed to parse tool call"}
+            return {"tool_name": "error", "parameters": {"message": "Failed to process request"}}