Update tools/tool_agent.py

tools/tool_agent.py

@@ -1,31 +1,32 @@
-from transformers import pipeline
-import json
-
-class ToolCallingAgent:
-    def __init__(self):
-        # Small CPU-friendly model
-        self.model = pipeline(
-            "text-generation",
-            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"""أنت مساعد ذكي. اختر أداة واحدة من القائمة التالية ثم أعد النتيجة كـ JSON فقط:
+{tools_json}
+
+مدخل المستخدم: {prompt}
+
+صيغة الإخراج المطلوبة:
+{{"tool_name": "اسم_الأداة", "parameters": {{"param": "القيمة"}}}}"""
+
+        response = self.model(
+            prompt,
+            max_new_tokens=150,
+            do_sample=False
         )
-        
-    def generate(self, prompt, tools):
-        try:
-            tools_json = json.dumps(tools, ensure_ascii=False)
-            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=100,
-                do_sample=False
-            )
-            
-            # 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)}"}}
+
+        text = response[0]['generated_text']
+
+        # Extract JSON from text
+        json_start = text.find("{")
+        json_end = text.rfind("}") + 1
+        parsed = json.loads(text[json_start:json_end])
+
+        # Fallback if JSON missing keys
+        if "tool_name" not in parsed:
+            parsed = {"tool_name": "error", "parameters": {"message": "لم يتم التعرف على الأداة"}}
+
+        return parsed
+
+    except Exception as e:
+        return {"tool_name": "error", "parameters": {"message": f"فشل المعالجة: {str(e)}"}}