Update tools/tool_agent.py

tools/tool_agent.py

@@ -1,7 +1,18 @@
-def generate(self, prompt, tools):
-    try:
+# tools/tool_agent.py
+from transformers import pipeline
+import json
+
+class ToolCallingAgent:
+    def __init__(self):
+        self.model = pipeline(
+            "text-generation",
+            model="aubmindlab/aragpt2-base",  # Arabic-friendly, small
+            device=-1
+        )
+
+    def generate(self, prompt, tools):
         tools_json = json.dumps(tools, ensure_ascii=False)
-        prompt = f"""أنت مساعد ذكي. اختر أداة واحدة من القائمة التالية ثم أعد النتيجة كـ JSON فقط:
+        full_prompt = f"""أنت مساعد ذكي. اختر أداة واحدة من القائمة التالية وأعد النتيجة كـ JSON فقط:
 {tools_json}
 
 مدخل المستخدم: {prompt}
@@ -9,24 +20,15 @@ def generate(self, prompt, tools):
 صيغة الإخراج المطلوبة:
 {{"tool_name": "اسم_الأداة", "parameters": {{"param": "القيمة"}}}}"""
 
-        response = self.model(
-            prompt,
-            max_new_tokens=150,
-            do_sample=False
-        )
-
-        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)}"}}
+        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)}}