Update pygmyclaw.py

pygmyclaw.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 """
-PygmyClaw – Compact AI Agent with multi-instance speculative decoding,
-persistent queue, and Hugging Face-backed persistent storage.
+PygmyClaw – Compact AI Agent with queue, AI Agent + HF tools support.
 """
 
 import os
@@ -10,80 +9,69 @@ import json
 import time
 import queue
 import threading
-from pathlib import Path
 import urllib.request
-
+from pathlib import Path
 from huggingface_hub import hf_hub_download, upload_file
 
-# Optional Redis support
-try:
-    import redis
-    REDIS_AVAILABLE = True
-except ImportError:
-    REDIS_AVAILABLE = False
-
 # -------------------- Globals --------------------
 SCRIPT_DIR = Path(__file__).parent.resolve()
 DEFAULT_MODEL = "qwen2.5:0.5b"
 DEFAULT_ENDPOINT = "http://localhost:11434/api/generate"
 
-TASK_QUEUE = queue.Queue()
-QUEUE_PROCESSOR_EVENT = threading.Event()
-
-# HF storage
 HF_TOKEN = os.environ.get("HF_TOKEN")
 HF_REPO = "rahul7star/pyclaw"
 HF_LOCAL_DIR = SCRIPT_DIR / "pyclaw_hf"
 FILES_TO_DOWNLOAD = ["memory.json", "tools.json"]
 
+TASK_QUEUE = queue.Queue()
+QUEUE_PROCESSOR_EVENT = threading.Event()
+
 # -------------------- HF Download --------------------
 def download_hf_files():
     HF_LOCAL_DIR.mkdir(parents=True, exist_ok=True)
-    for fname in FILES_TO_DOWNLOAD:
-        local_path = HF_LOCAL_DIR / fname
+    for file_name in FILES_TO_DOWNLOAD:
+        local_path = HF_LOCAL_DIR / file_name
         if not local_path.exists() or local_path.stat().st_size == 0:
             try:
-                path = hf_hub_download(
+                hf_hub_download(
                     repo_id=HF_REPO,
-                    filename=fname,
+                    filename=file_name,
                     token=HF_TOKEN,
                     local_dir=str(HF_LOCAL_DIR)
                 )
-                print(f"Downloaded {fname} -> {path}")
+                print(f"Downloaded {file_name}")
             except Exception as e:
-                print(f"⚠️ Failed to download {fname}: {e}")
+                print(f"⚠️ Failed to download {file_name}: {e}")
                 local_path.write_text("{}")
-                print(f"Created empty {fname}")
+                print(f"Created empty {file_name}")
 
-# -------------------- PygmyClaw Class --------------------
+# -------------------- PygmyClaw Agent --------------------
 class PygmyClaw:
     def __init__(self):
-
-         # Use given workspace or default
-        if workspace:
-            self.workspace = Path(workspace).resolve()
-        else:
-            self.workspace = Path("/workspace/data").resolve()
-        self.workspace.mkdir(parents=True, exist_ok=True)
         self.model = DEFAULT_MODEL
         self.endpoint = DEFAULT_ENDPOINT
         self.memory_data = {}
         self.tools_data = {}
+        self.python_tools = ["Python Script"]  # default code execution tool
 
         # Ensure HF files exist
         download_hf_files()
         self._load_hf_memory()
         self._load_hf_tools()
 
+        # Combine tools from HF + Python + AI Agent
+        self.python_tools += list(self.tools_data.keys())
+        self.python_tools.append("AI Agent")  # default AI Agent
+
         # Ensure model is ready
         self._ensure_model_ready()
         self._warmup_model()
 
-        # Start queue processor
+        # Start queue processor in background
         QUEUE_PROCESSOR_EVENT.set()
         threading.Thread(target=self._process_queue, daemon=True).start()
 
-    # -------------------- HF Memory / Tools --------------------
+    # -------------------- Memory & Tools --------------------
     def _load_hf_memory(self):
         mem_file = HF_LOCAL_DIR / "memory.json"
         mem_file.parent.mkdir(parents=True, exist_ok=True)
@@ -94,26 +82,22 @@ class PygmyClaw:
                 self.memory_data = json.load(f)
             except json.JSONDecodeError:
                 self.memory_data = {}
-                print("⚠️ memory.json was empty or invalid, initialized with {}")
+                print("⚠️ memory.json invalid, initialized with {}")
 
     def _save_hf_memory(self):
         mem_file = HF_LOCAL_DIR / "memory.json"
         with open(mem_file, "w") as f:
             json.dump(self.memory_data, f, indent=2)
-
-        # Upload to Hugging Face if token provided
-        if HF_TOKEN:
-            try:
-                upload_file(
-                    path_or_fileobj=str(mem_file),
-                    path_in_repo="memory.json",
-                    repo_id=HF_REPO,
-                    token=HF_TOKEN,
-                    repo_type="model"
-                )
-                print("💾 Memory saved to HF repo")
-            except Exception as e:
-                print(f"⚠️ Failed to push memory to HF: {e}")
+        try:
+            upload_file(
+                path_or_fileobj=str(mem_file),
+                path_in_repo="memory.json",
+                repo_id=HF_REPO,
+                token=HF_TOKEN,
+                repo_type="model"
+            )
+        except Exception as e:
+            print(f"⚠️ Failed to push memory to HF: {e}")
 
     def _load_hf_tools(self):
         tools_file = HF_LOCAL_DIR / "tools.json"
@@ -125,13 +109,13 @@ class PygmyClaw:
                 self.tools_data = json.load(f)
             except json.JSONDecodeError:
                 self.tools_data = {}
-                print("⚠️ tools.json was empty or invalid, initialized with {}")
+                print("⚠️ tools.json invalid, initialized with {}")
 
     # -------------------- Model --------------------
     def _ensure_model_ready(self):
         print(f"⏳ Ensuring model '{self.model}' is ready...")
+        payload = {"model": self.model, "prompt": "hello", "stream": False, "options": {"num_predict": 1}}
         try:
-            payload = {"model": self.model, "prompt": "hello", "stream": False, "options": {"num_predict": 1}}
             req = urllib.request.Request(
                 self.endpoint,
                 data=json.dumps(payload).encode("utf-8"),
@@ -155,18 +139,16 @@ class PygmyClaw:
                 headers={"Content-Type": "application/json"},
                 method="POST"
             )
-            with urllib.request.urlopen(req, timeout=5) as resp:
+            with urllib.request.urlopen(req, timeout=5):
                 pass
         except Exception:
             pass
 
-    # -------------------- Queue Processor --------------------
-    def add_task(self, prompt, tool=None, callback=None):
-        """Add a task to the queue with optional tool and callback"""
+    # -------------------- Queue --------------------
+    def add_task(self, prompt, tool="AI Agent", callback=None):
         task_id = str(time.time())
-        task = {"id": task_id, "prompt": prompt, "tool": tool or "AI Agent", "callback": callback}
-        TASK_QUEUE.put(task)
-        print(f"Task {task_id} queued with tool={task['tool']}.")
+        TASK_QUEUE.put({"id": task_id, "prompt": prompt, "tool": tool, "callback": callback})
+        print(f"Task {task_id} queued with tool={tool}")
         return task_id
 
     def _process_queue(self):
@@ -177,59 +159,56 @@ class PygmyClaw:
             except queue.Empty:
                 continue
 
-            print(f"Processing task {task['id']} -> {task['prompt']}")
+            task_id = task["id"]
+            prompt = task["prompt"]
+            tool = task.get("tool", "AI Agent")
+            callback = task.get("callback", None)
+
+            print(f"Processing task {task_id} with tool={tool} -> {prompt}")
             try:
-                result = self.generate_with_ssd(task["prompt"])
-                self.memory_data[task["id"]] = {
-                    "prompt": task["prompt"],
+                if tool == "Python Script":
+                    # Run code dynamically
+                    local_vars = {}
+                    exec(prompt, {}, local_vars)
+                    result = str(local_vars)
+                else:
+                    # AI Agent or HF tool
+                    result = self.generate_with_ssd(prompt)
+
+                self.memory_data[task_id] = {
+                    "prompt": prompt,
                     "response": result,
                     "timestamp": time.time(),
-                    "tool": task.get("tool", "AI Agent")
+                    "tool": tool
                 }
                 self._save_hf_memory()
-                print(f"✅ Task {task['id']} processed.")
-
-                # Call callback if present
-                if task.get("callback"):
-                    try:
-                        task["callback"](result)
-                    except Exception as e:
-                        print(f"⚠️ Callback failed for task {task['id']}: {e}")
-
+                if callback:
+                    callback(result)
+                print(f"✅ Task {task_id} done.")
             except Exception as e:
-                print(f"❌ Failed to process task {task['id']}: {e}")
+                print(f"❌ Task {task_id} failed: {e}")
             finally:
                 TASK_QUEUE.task_done()
 
     # -------------------- Submit Prompt --------------------
-    def submit_prompt(self, prompt):
-        """Add prompt to queue and wait for completion."""
-        task_id = self.add_task(prompt)
+    def submit_prompt(self, prompt, tool="AI Agent"):
+        task_id = self.add_task(prompt, tool=tool)
         while task_id not in self.memory_data:
             time.sleep(0.1)
         return self.memory_data[task_id]["response"]
 
-    # -------------------- Model Call --------------------
+    # -------------------- Model / Ollama call --------------------
     def generate_with_ssd(self, prompt):
-        payload = {"model": self.model, "prompt": prompt, "stream": False, "options": {"num_predict": 1}}
+        payload = {"model": self.model, "prompt": prompt, "stream": False, "options": {"num_predict": 50}}
         req = urllib.request.Request(
             self.endpoint,
             data=json.dumps(payload).encode("utf-8"),
             headers={"Content-Type": "application/json"},
             method="POST"
         )
-        with urllib.request.urlopen(req, timeout=60) as resp:
-            resp_data = json.loads(resp.read())
-            return resp_data.get("response", "")
-
-# -------------------- Main --------------------
-def main():
-    agent = PygmyClaw()
-    prompt = input("Enter prompt: ")
-    outcome = agent.submit_prompt(prompt)
-    print("\n=== Outcome ===")
-    print(outcome)
-
-
-if __name__ == "__main__":
-    main()
+        try:
+            with urllib.request.urlopen(req, timeout=60) as resp:
+                data = json.loads(resp.read())
+                return data.get("response", "No response")
+        except Exception as e:
+            return f"❌ Model request failed: {e}"