Update tools/tool_agent.py

tools/tool_agent.py

@@ -1,4 +1,3 @@
-# tools/tool_agent.py
 from transformers import pipeline
 import json
 
@@ -6,29 +5,27 @@ class ToolCallingAgent:
     def __init__(self):
         self.model = pipeline(
             "text-generation",
-            model="aubmindlab/aragpt2-base",  # Arabic-friendly, small
-            device=-1
+            model="cognitivecomputations/dolphin-2.9-llama3-8b",
+            device_map="auto"
         )
-
+        
     def generate(self, prompt, tools):
+        # Format the tools specification
         tools_json = json.dumps(tools, ensure_ascii=False)
-        full_prompt = f"""أنت مساعد ذكي. اختر أداة واحدة من القائمة التالية وأعد النتيجة كـ JSON فقط:
-{tools_json}
-
-مدخل المستخدم: {prompt}
-
-صيغة الإخراج المطلوبة:
-{{"tool_name": "اسم_الأداة", "parameters": {{"param": "القيمة"}}}}"""
-
+        
+        # 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:
-            response = self.model(
-                full_prompt,
-                max_new_tokens=150,
-                do_sample=False
-            )
-            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": str(e)}}
+            return json.loads(response[0]['generated_text'])
+        except:
+            return {"error": "Failed to parse tool call"}