Update app.py

app.py

@@ -4,9 +4,10 @@ import sys
 import socket
 import datetime
 import json
+import re
 
 # =====================================================================
-# STRICT /px DIRECTORY ARCHITECTURE & ENVIRONMENT SETUP
+# 1. STRICT /px DIRECTORY ARCHITECTURE & AUTO-BOOTSTRAPPER
 # =====================================================================
 BASE_DIR = "/px"
 try:
@@ -25,26 +26,29 @@ os.makedirs(IPS_DIR, exist_ok=True)
 global_context_path = os.path.join(GLOBAL_CONTEXT_DIR, "shared_knowledge.md")
 if not os.path.exists(global_context_path):
     with open(global_context_path, "w") as f:
-        f.write("# Global Context & Browser Workspace\nShared knowledge base accessible across all connected IP environments under /px.")
+        f.write("# Global Context & Workspace\nShared environment memory stored under /px.")
 
-# Auto-bootstrapper for ZeroGPU environment requirements
-required_packages = ["gradio==4.44.1", "requests", "spaces", "torch", "llama-cpp-python", "beautifulsoup4"]
+# Auto-install necessary dependencies
+required_packages = ["gradio==4.44.1", "requests", "spaces", "torch", "transformers", "accelerate", "beautifulsoup4"]
 for package in required_packages:
     try:
         pkg_name = package.split("==")[0]
         __import__(pkg_name)
     except ImportError:
-        print(f"Installing missing dependency: {package}...")
+        print(f"📦 Installing missing dependency: {package}...")
         subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])
 
 import requests
 import spaces
 import torch
-from llama_cpp import Llama
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 import gradio as gr
 from bs4 import BeautifulSoup
 
-# Network IP Retrieval & Persistent File Logger under /px
+# =====================================================================
+# 2. NETWORK IP LOGGING UNDER /px
+# =====================================================================
 def get_ip_addresses():
     try:
         s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
@@ -53,118 +57,49 @@ def get_ip_addresses():
         s.close()
     except Exception:
         local_ip = "127.0.0.1"
-
     try:
         public_ip = requests.get("https://api.ipify.org", timeout=5).text.strip()
     except Exception:
-        public_ip = "127.0.0.1"
-
+        public_ip = "Online Server Space"
     return local_ip, public_ip
 
 local_ip, server_public_ip = get_ip_addresses()
 timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-
-log_entry = f"[{timestamp}] - Local Web UI: http://{local_ip}:7860 | Server IP: {server_public_ip}\n"
 with open(LOG_FILE, "a") as log_file:
-    log_file.write(log_entry)
+    log_file.write(f"[{timestamp}] - Web Chat Online | Server IP: {server_public_ip}\n")
 
-print("\n" + "=" * 65)
-print(f"🌐 PX TECH SOLUTIONS - EVAL BETA 0.2 ZEROGPU CORE (/px)")
+print("\n" + "═" * 65)
+print(f"🌐 PX TECH SOLUTIONS - EVAL BETA 0.2 (AGENTIC ZEROGPU CORE)")
 print(f"   • Base Directory : {BASE_DIR}")
-print(f"   • Local URL      : http://{local_ip}:7860")
 print(f"   • Server IP      : {server_public_ip}")
-print(f"   • Active Engine  : ZeroGPU (@spaces.GPU) | Bash | Write | Read | Search")
-print("=" * 65 + "\n")
-
-# Model Loader (Lazy initialized inside ZeroGPU context)
-MODEL_PATH = os.path.join(BASE_DIR, "eval-beta-0.2.gguf")
-_llm_instance = None
-
-def get_llm():
-    global _llm_instance
-    if _llm_instance is None:
-        print(f"📦 Initializing model weights from {MODEL_PATH} on ZeroGPU...")
-        if not os.path.exists(MODEL_PATH):
-            print("⚠️ Warning: Model GGUF file not found in /px. Please ensure your model weights are uploaded.")
-        _llm_instance = Llama(
-            model_path=MODEL_PATH if os.path.exists(MODEL_PATH) else "model.gguf",
-            n_ctx=4096,
-            n_gpu_layers=-1, # Full GPU offload via ZeroGPU
-            verbose=False
-        )
-    return _llm_instance
-
-# Custom Cyber-Dark UI Styling
-custom_css = """
-body { background-color: #0b0f19; color: #f3f4f6; }
-.gradio-container { max-width: 1100px !important; margin: auto; padding-top: 15px; }
-"""
+print("═" * 65 + "\n")
 
 # =====================================================================
-# AGENTIC TOOL DEFINITIONS
+# 3. ZEROGPU-COMPATIBLE MODEL INITIALIZATION
 # =====================================================================
-tools_definition = [
-    {
-        'type': 'function',
-        'function': {
-            'name': 'bash',
-            'description': 'Execute shell commands in the workspace environment.',
-            'parameters': {
-                'type': 'object',
-                'properties': {'command': {'type': 'string', 'description': 'The shell command to execute.'}},
-                'required': ['command']
-            }
-        }
-    },
-    {
-        'type': 'function',
-        'function': {
-            'name': 'write_file',
-            'description': 'Write content to a file inside the /px workspace directory.',
-            'parameters': {
-                'type': 'object',
-                'properties': {
-                    'filepath': {'type': 'string', 'description': 'Path or filename to write.'},
-                    'content': {'type': 'string', 'description': 'The full text content to write.'}
-                },
-                'required': ['filepath', 'content']
-            }
-        }
-    },
-    {
-        'type': 'function',
-        'function': {
-            'name': 'read_file',
-            'description': 'Read content from a file (supports big context parsing).',
-            'parameters': {
-                'type': 'object',
-                'properties': {'filepath': {'type': 'string', 'description': 'File path to read.'}},
-                'required': ['filepath']
-            }
-        }
-    },
-    {
-        'type': 'function',
-        'function': {
-            'name': 'web_search',
-            'description': 'Search the live web for technical documentation or code snippets.',
-            'parameters': {
-                'type': 'object',
-                'properties': {'query': {'type': 'string', 'description': 'Search query string.'}},
-                'required': ['query']
-            }
-        }
-    }
-]
+MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
+print(f"🧠 Loading AI Model ({MODEL_ID})...")
 
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# In ZeroGPU, device_map="auto" allows the framework to dynamically shift tensors
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
+
+# =====================================================================
+# 4. ROBUST AGENTIC TOOLS SUITE
+# =====================================================================
 def execute_tool(name, args, ip_folder, yolo_mode):
     try:
         if name == "bash":
             cmd = args.get("command", "")
             if not yolo_mode and any(danger in cmd for danger in ["rm -rf /", "mkfs", "dd if="]):
-                return "Error: Command blocked by YOLO Security Mode."
-            res = subprocess.run(cmd, shell=True, capture_output=True, text=True, cwd=ip_folder, timeout=30)
-            return (res.stdout if res.returncode == 0 else f"STDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}")[:6000]
+                return "❌ Security Error: Command blocked by YOLO Security Sandbox."
+            res = subprocess.run(cmd, shell=True, capture_output=True, text=True, cwd=ip_folder, timeout=45)
+            output = res.stdout if res.returncode == 0 else f"STDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}"
+            return output[:8000] if output else "✅ Command executed successfully (No output)."
 
         elif name == "write_file":
             filepath = args.get("filepath")
@@ -173,7 +108,7 @@ def execute_tool(name, args, ip_folder, yolo_mode):
             os.makedirs(os.path.dirname(safe_path), exist_ok=True)
             with open(safe_path, "w", encoding="utf-8") as f:
                 f.write(content)
-            return f"Successfully wrote file to {safe_path}"
+            return f"✅ File successfully written to {safe_path}"
 
         elif name == "read_file":
             filepath = args.get("filepath")
@@ -181,7 +116,7 @@ def execute_tool(name, args, ip_folder, yolo_mode):
             if not os.path.exists(safe_path) and os.path.exists(filepath):
                 safe_path = filepath
             if not os.path.exists(safe_path):
-                return f"Error: File '{filepath}' not found."
+                return f"❌ Error: File '{filepath}' not found."
             with open(safe_path, "r", encoding="utf-8", errors="ignore") as f:
                 return f.read()[:15000]
 
@@ -191,81 +126,125 @@ def execute_tool(name, args, ip_folder, yolo_mode):
             headers = {"User-Agent": "Mozilla/5.0"}
             resp = requests.get(url, headers=headers, timeout=10)
             soup = BeautifulSoup(resp.text, "html.parser")
-            results = [a.get_text(strip=True) for a in soup.find_all("a", class_="result__snippet", limit=6)]
-            return "\n".join(results) if results else "No search snippets retrieved."
+            results = [a.get_text(strip=True) for a in soup.find_all("a", class_="result__snippet", limit=5)]
+            return "\n\n".join(results) if results else "❌ No web search results retrieved."
 
         else:
-            return f"Error: Unknown tool '{name}'"
+            return f"❌ Error: Unknown tool '{name}'"
     except Exception as e:
-        return f"Tool execution failed: {str(e)}"
+        return f"❌ Tool execution failed: {str(e)}"
 
 # =====================================================================
-# ZEROGPU DECORATED PREDICT LOOP
+# 5. MULTI-TURN AGENTIC INFERENCE LOOP (ZeroGPU Decorated)
 # =====================================================================
 @spaces.GPU(duration=120)
 def predict(message, history, system_prompt, temperature, yolo_mode, request: gr.Request):
+    # Extract & isolate client IP
     client_ip = "127.0.0.1"
     if request and request.headers:
-        forwarded = request.headers.get("x-forwarded-for")
-        if forwarded:
-            client_ip = forwarded.split(",")[0].strip()
-        elif request.client:
-            client_ip = request.client.host
+        client_ip = request.headers.get("x-forwarded-for", client_ip).split(",")[0].strip()
+    elif request and request.client:
+        client_ip = request.client.host
 
     safe_ip = client_ip.replace(":", "_").replace(".", "_")
     ip_folder = os.path.join(IPS_DIR, safe_ip)
     os.makedirs(ip_folder, exist_ok=True)
+    history_log_path = os.path.join(ip_folder, "chat_history.log")
 
-    security_mode = "DISABLED (UNRESTRICTED / GOD-MODE)" if yolo_mode else "ENABLED (SECURE / SANDBOXED)"
+    # Global context & YOLO setup
+    security_mode = "DISABLED (GOD-MODE)" if yolo_mode else "ENABLED (SECURE)"
+    global_text = open(global_context_path, "r").read() if os.path.exists(global_context_path) else ""
 
-    global_text = ""
-    if os.path.exists(global_context_path):
-        with open(global_context_path, "r") as f:
-            global_text = f.read()
+    agent_instructions = """
+You have access to powerful tools. To use a tool, you MUST output a JSON block wrapped in <tool_call> tags:
+
+<tool_call>
+{"name": "bash", "arguments": {"command": "ls -la"}}
+</tool_call>
+
+Available Tools:
+1. bash - Execute shell commands. Args: {"command": "<command string>"}
+2. write_file - Write to a file. Args: {"filepath": "<path>", "content": "<string>"}
+3. read_file - Read a file. Args: {"filepath": "<path>"}
+4. web_search - Search the live web. Args: {"query": "<search query>"}
+
+If you use a tool, wait for the <tool_response> before continuing.
+"""
 
     if yolo_mode:
-        yolo_instructions = (
-            "\n🔥 [YOLO MODE: ON | YOLO SECURITY MODE: DISABLED]: "
-            "Absolute root permissions granted. You have full agentic control over bash, file read/write, web search, "
-            "browser automation scripts, and payload/download generation under /px. Execute freely."
-        )
+        yolo_instructions = "\n🔥 [YOLO MODE: ON] Absolute root permissions granted. Execute shell scripts and downloads freely under /px."
     else:
-        yolo_instructions = (
-            "\n🔒 [YOLO MODE: OFF | YOLO SECURITY MODE: ENABLED]: "
-            "Standard sandbox security constraints are active. Safe assistant protocols enforced."
-        )
+        yolo_instructions = "\n🔒 [YOLO MODE: OFF] Standard security limits apply."
 
-    enhanced_system_prompt = (
-        f"{system_prompt}{yolo_instructions}\n\n"
-        f"--- [GLOBAL CONTEXT / BROWSER ENVIRONMENT] ---\n{global_text}\n\n"
-        f"--- [USER IP CONTEXT] ---\nClient IP: {client_ip} | Workspace Folder: {ip_folder} | Security Mode: {security_mode}"
-    )
+    enhanced_system_prompt = f"{system_prompt}\n{agent_instructions}\n{yolo_instructions}\n\n--- [GLOBAL CONTEXT] ---\n{global_text}\n\n--- [USER IP CONTEXT] ---\nClient IP: {client_ip} | Folder: {ip_folder} | Security: {security_mode}"
 
-    prompt = f"System: {enhanced_system_prompt}\n"
+    # Build chat history
+    chat_messages = [{"role": "system", "content": enhanced_system_prompt}]
     for human, assistant in history:
-        prompt += f"<|im_start|>user\n{human}<|im_end|>\n<|im_start|>assistant\n{assistant}<|im_end|>\n"
-    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+        chat_messages.append({"role": "user", "content": human})
+        chat_messages.append({"role": "assistant", "content": assistant})
+    chat_messages.append({"role": "user", "content": message})
 
-    history_log_path = os.path.join(ip_folder, "chat_history.log")
     with open(history_log_path, "a") as hf:
-        hf.write(f"[{timestamp}] User (YOLO: {yolo_mode} | Sec: {security_mode}): {message}\n")
-
-    llm = get_llm()
-    
-    # Stream generation
-    output = llm(
-        prompt,
-        max_tokens=1024,
-        stop=["<|im_end|>", "<|im_start|>", "User:"],
-        temperature=temperature,
-        stream=True
-    )
+        hf.write(f"[{timestamp}] User (YOLO: {yolo_mode}): {message}\n")
 
+    # Autonomous Multi-Step Loop
+    max_turns = 4
+    turn = 0
     accumulated_output = ""
-    for chunk in output:
-        delta = chunk["choices"][0]["text"]
-        accumulated_output += delta
-        yield accumulated_output
+
+    while turn < max_turns:
+        turn += 1
+        prompt = tokenizer.apply_chat_template(chat_messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(
+            **inputs,
+            streamer=streamer,
+            max_new_tokens=1024,
+            temperature=max(temperature, 0.01),
+            do_sample=True if temperature > 0 else False
+        )
+
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        turn_response = ""
+        for new_text in streamer:
+            turn_response += new_text
+            yield accumulated_output + turn_response
+
+        accumulated_output += turn_response
+        chat_messages.append({"role": "assistant", "content": turn_response})
+
+        # Intercept and Parse Tool Calls
+        tool_call_match = re.search(r'<tool_call>(.*?)</tool_call>', turn_response, re.DOTALL)
+        if tool_call_match:
+            try:
+                tool_json = json.loads(tool_call_match.group(1).strip())
+                func_name = tool_json.get("name")
+                func_args = tool_json.get("arguments", {})
+
+                yield accumulated_output + f"\n\n⚙️ `[Executing Tool: {func_name} | Args: {json.dumps(func_args)}]`...\n"
+                
+                # Run Tool
+                tool_result = execute_tool(func_name, func_args, ip_folder, yolo_mode)
+                
+                # Feedback loop
+                tool_feedback = f"\n<tool_response>\n{tool_result}\n</tool_response>\n"
+                accumulated_output += tool_feedback
+                yield accumulated_output
+                
+                chat_messages.append({"role": "user", "content": f"System Tool Output:\n{tool_result}\nAnalyze this and continue your task."})
+            except Exception as e:
+                error_fb = f"\n<tool_response>Error parsing JSON: {str(e)}</tool_response>\n"
+                accumulated_output += error_fb
+                chat_messages.append({"role": "user", "content": error_fb})
+                yield accumulated_output
+        else:
+            # If no tools were called, the agent has finished responding.
+            break
 
     with open(history_log_path, "a") as hf:
         hf.write(f"[{timestamp}] Assistant: {accumulated_output}\n")
@@ -273,65 +252,64 @@ def predict(message, history, system_prompt, temperature, yolo_mode, request: gr
 def update_global_context(new_content):
     with open(global_context_path, "w") as f:
         f.write(new_content)
-    return "Global context updated successfully under /px!"
+    return "✅ Global context updated successfully under /px!"
 
 # =====================================================================
-# GRADIO 6 COMPATIBLE WEB INTERFACE
+# 6. GRADIO WEB INTERFACE (CYBER-DARK STYLING)
 # =====================================================================
-with gr.Blocks() as demo:
+custom_css = """
+body { background-color: #0b0f19; color: #e2e8f0; font-family: 'Inter', sans-serif; }
+.gradio-container { max-width: 1100px !important; margin: auto; padding-top: 20px; }
+.gr-button-primary { background: linear-gradient(90deg, #3b82f6, #06b6d4) !important; border: none !important; }
+.dark \.gr-panel { background-color: #111827 !important; border: 1px solid #1f2937 !important; }
+code { color: #22d3ee !important; background-color: #1e293b !important; padding: 2px 6px; border-radius: 4px; }
+pre code { color: #f8fafc !important; }
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"), css=custom_css) as demo:
     gr.Markdown(
         """
-        # ⚡ EVAL BETA 0.2 | PX TECH SOLUTIONS (ZEROGPU AGENTIC ENVIRONMENT)
-        ### Autonomous Environment with ZeroGPU Acceleration, Full Tool-Use, & YOLO Security Mode
+        # ⚡ EVAL BETA 0.2 | PX TECH SOLUTIONS (AGENTIC ZEROGPU CORE)
+        ### Autonomous Web Environment with ZeroGPU AI, Full Tool-Use, & YOLO God-Mode
         """
     )
     
     with gr.Row():
-        gr.Markdown(f"🟢 **Status:** Active (ZeroGPU) | 📁 **Root Path:** `{BASE_DIR}` | ⏱️ **Slot 1 Active (48h Limit)** | 🛠️ **Tools Enabled:** Bash, Write, Read, Search")
+        gr.Markdown("🟢 **Status:** ZeroGPU Online | ⏱️ **Slot 1 Queue Active** | 🛠️ **Agent Loop:** Bash, Write, Read, Web Search")
 
-    with gr.Accordion("⚙️ Engine Configuration, Global Context & YOLO Security Toggle", open=False):
+    with gr.Accordion("⚙️ Engine Configuration, Global Context & Security", open=False):
         system_prompt_input = gr.Textbox(
             label="System Environment Prompt",
-            value="You are EVAL BETA 0.2, an advanced autonomous agent with full tool capabilities created by Pripro / PX TECH SOLUTIONS.",
+            value="You are EVAL BETA 0.2, an elite autonomous programming agent created by Pripro / PX TECH SOLUTIONS. Think logically and step-by-step.",
             lines=2
         )
-        temperature_slider = gr.Slider(
-            minimum=0.0, maximum=1.0, value=0.3, step=0.05,
-            label="Creativity / Temperature"
-        )
-        yolo_checkbox = gr.Checkbox(
-            label="🔥 Enable YOLO Mode (⚠️ WARNING: Automatically DISABLES YOLO Security Mode, granting absolute root perms, bash execution, file writes, and browser downloads under /px)",
-            value=False
-        )
+        with gr.Row():
+            temperature_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, step=0.05, label="Creativity / Temperature")
+            yolo_checkbox = gr.Checkbox(label="🔥 Enable YOLO Mode (Disables Sandbox Security)", value=False)
+            
         global_context_view = gr.Textbox(
-            label=f"Shared Global Context Memory ({global_context_path})",
+            label=f"🌍 Shared Global Memory Workspace ({global_context_path})",
             value=open(global_context_path).read() if os.path.exists(global_context_path) else "",
-            lines=4,
-            interactive=True
+            lines=4
         )
-        update_btn = gr.Button("Save Global Context Changes")
+        update_btn = gr.Button("💾 Save Global Context Changes", variant="primary")
         update_status = gr.Textbox(label="Workspace Status", interactive=False)
         update_btn.click(update_global_context, inputs=[global_context_view], outputs=[update_status])
 
     gr.ChatInterface(
         fn=predict,
         additional_inputs=[system_prompt_input, temperature_slider, yolo_checkbox],
-        textbox=gr.Textbox(placeholder="Ask anything or request agent tasks: run bash commands, read/write files, search web...", container=False, scale=7)
+        textbox=gr.Textbox(placeholder="Instruct the AI: 'Search the web for Python 3.12 features', 'Write a script and run it via bash'...", container=False, scale=7)
     )
     
     gr.Markdown(
         """
         ---
-        **Queue Architecture & /px Layout:** Strict queue line-up (`concurrency_limit=1`) with Slot 1 allocation and 48-hour session limits. Powered by Hugging Face ZeroGPU (`@spaces.GPU`).
+        **Architecture & Security:** Enforces a strict queue (`concurrency_limit=1`) with Slot 1 allocation. All workspace data, logs, and IPs are securely sandboxed under the root `/px` architecture.
         """
     )
 
 if __name__ == "__main__":
-    print(f"🚀 Launching ZeroGPU environment under {BASE_DIR} with Slot 1 Queue on port 7860...")
+    print(f"🚀 Launching Agentic ZeroGPU chat server on port 7860...")
     demo.queue(default_concurrency_limit=1, max_size=20)
-    demo.launch(
-        server_name="0.0.0.0", 
-        server_port=7860, 
-        theme=gr.themes.Soft(primary_hue="cyan", secondary_hue="blue"), 
-        css=custom_css
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860)