Ahhh

app.py

@@ -1,365 +1,325 @@
-# === IMPORTS: Standard, secure, minimal ===
+# app.py
 import gradio as gr
 import json
 import re
 from datetime import datetime
 from contextlib import contextmanager
-from typing import Dict, Any, Optional
+from typing import Dict, Any
 
-# SQLAlchemy for robust, transactional data persistence
 from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, ForeignKey, JSON
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, relationship
-
-# Discord client (unofficial but functional; replace with discord.py if preferred)
 import discum
+from smolagents import Tool
 
-# smolagents for agent-tool orchestration (correctly integrated)
-from smolagents import Tool, Agent
-
-
-# === DATABASE CONFIGURATION: Persistent, local, atomic ===
-# Use SQLite for simplicity; replace with PostgreSQL/MySQL in production
-DB_URL = "sqlite:///puppets.db"
+# === DATABASE SETUP ===
+DB_PATH = "sqlite:///puppets.db"
 Base = declarative_base()
-engine = create_engine(DB_URL, connect_args={"check_same_thread": False})  # Safe for Gradio
+engine = create_engine(DB_PATH, connect_args={"check_same_thread": False})
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
-
-# === DATA MODELS: Immutable schema with validation ===
+# === MODELS ===
 class Personality(Base):
-    """
-    Represents a reusable AI personality.
-    - `base_prompt`: Core behavior definition (e.g., "You are a sarcastic chef...")
-    - `fine_tune_config`: LLM parameters (model, temperature) as JSON
-    """
     __tablename__ = "personalities"
     id = Column(Integer, primary_key=True, autoincrement=True)
-    name = Column(String(255), nullable=False, unique=False)  # Allow duplicate names
+    name = Column(String(255), nullable=False)
     base_prompt = Column(Text, nullable=False)
     fine_tune_config = Column(JSON, default=dict, nullable=False)
     created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
 
-
 class ChannelConfig(Base):
-    """
-    Binds a personality to a Discord channel with optional overrides.
-    - `channel_id`: Discord channel snowflake (string to avoid overflow)
-    - `channel_instructions`: Channel-specific rules (e.g., "Always use emojis")
-    """
     __tablename__ = "channel_configs"
     id = Column(Integer, primary_key=True, autoincrement=True)
-    channel_id = Column(String(64), nullable=False)  # Discord IDs are up to 19 digits
+    channel_id = Column(String(64), nullable=False)
     personality_id = Column(Integer, ForeignKey("personalities.id"), nullable=False)
     channel_instructions = Column(Text, nullable=True)
     created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
-    # Relationship for easy access to personality data
     personality = relationship("Personality", lazy="joined")
 
+# === ENSURE SCHEMA COMPATIBILITY ===
+def add_missing_columns():
+    import sqlite3
+    conn = sqlite3.connect("puppets.db")
+    cursor = conn.cursor()
+    try:
+        cursor.execute("ALTER TABLE personalities ADD COLUMN fine_tune_config TEXT DEFAULT '{}';")
+    except sqlite3.OperationalError:
+        pass  # Column already exists
+    conn.commit()
+    conn.close()
 
-# === INITIALIZE DATABASE: Idempotent, safe on startup ===
-# Creates tables if they don't exist; no-op if they do
+add_missing_columns()
 Base.metadata.create_all(bind=engine)
 
-
-# === GLOBAL STATE: Isolated, managed puppet instances ===
-# Maps integer puppet_id → live discum.Client
-# ⚠️ In production, add TTL cleanup or explicit shutdown hooks
+# === GLOBAL STATE ===
 _discum_clients: Dict[int, discum.Client] = {}
 
-
-# === DATABASE SESSION MANAGER: Transaction-safe context ===
+# === UTILS ===
 @contextmanager
 def get_db_session():
-    """
-    Provides a transactional database session.
-    - Commits on success
-    - Rolls back on exception
-    - Always closes session (prevents leaks)
-    Usage: `with get_db_session() as session: ...`
-    """
     session = SessionLocal()
     try:
         yield session
+        session.commit()
     except Exception:
         session.rollback()
         raise
     finally:
         session.close()
 
-
-# === VALIDATION HELPERS: Prevent invalid inputs ===
 def validate_discord_id(discord_id: str) -> bool:
-    """
-    Validates Discord snowflake IDs.
-    - Must be a string of digits (1-20 characters)
-    - Rejects empty, non-numeric, or malformed IDs
-    """
-    if not isinstance(discord_id, str) or not discord_id:
-        return False
-    return bool(re.fullmatch(r"^\d{1,20}$", discord_id))
-
+    return bool(re.fullmatch(r"^\d{1,20}$", discord_id)) if discord_id else False
 
-# === TOOL DEFINITIONS: smolagents-compatible, structured I/O ===
-# Each tool:
-# - Inherits from smolagents.Tool
-# - Defines strict inputs/outputs
-# - Returns JSON (not raw strings) for agent consumption
-# - Handles errors gracefully (no crashes)
+def build_final_prompt(personality, channel_instr=""):
+    parts = [personality.base_prompt]
+    if channel_instr:
+        parts.append(f"\n\n[CHANNEL RULES]\n{channel_instr}")
+    return "\n".join(parts)
 
+# === TOOLS ===
 class ListPersonalitiesTool(Tool):
-    """Tool for agent to discover available personalities."""
     name = "list_personalities"
-    description = "Retrieve all defined personalities with IDs and names."
-    inputs = {}  # No inputs required
-    output_type = "string"  # smolagents requires string; we serialize JSON
-
-    def forward(self) -> str:
-        """Returns JSON array of personalities."""
+    description = "List all personalities"
+    inputs = {}
+    output_type = "string"
+    def forward(self):
         try:
-            with get_db_session() as session:
-                personalities = session.query(Personality).all()
-                result = [
-                    {"id": p.id, "name": p.name}
-                    for p in personalities
-                ]
-                return json.dumps(result)
+            with get_db_session() as s:
+                ps = s.query(Personality).all()
+                return json.dumps([{"id": p.id, "name": p.name} for p in ps])
         except Exception as e:
-            # Never crash the agent; return structured error
-            return json.dumps({"error": "db_query_failed", "message": str(e)})
-
+            return json.dumps({"error": str(e)})
 
 class GetChannelPromptTool(Tool):
-    """Tool to build the final prompt for a channel using personality + overrides."""
     name = "get_channel_prompt"
-    description = "Generate the full prompt for a channel by combining personality and channel instructions."
+    description = "Get final prompt for channel and personality"
     inputs = {
-        "channel_id": {"type": "string", "description": "Discord channel ID"},
-        "personality_id": {"type": "integer", "description": "Personality ID"},
+        "channel_id": {"type": "string"},
+        "personality_id": {"type": "integer"}
     }
     output_type = "string"
-
-    def forward(self, channel_id: str, personality_id: int) -> str:
-        """Returns JSON with prompt and LLM config."""
-        # Validate inputs early
-        if not validate_discord_id(channel_id):
-            return json.dumps({"error": "invalid_channel_id"})
-        if personality_id <= 0:
-            return json.dumps({"error": "invalid_personality_id"})
-
+    def forward(self, channel_id: str, personality_id: int):
+        if not validate_discord_id(channel_id) or personality_id <= 0:
+            return json.dumps({"error": "invalid inputs"})
         try:
-            with get_db_session() as session:
-                # Fetch personality
-                personality = session.query(Personality).filter(
-                    Personality.id == personality_id
-                ).first()
-                if not personality:
-                    return json.dumps({"error": "personality_not_found"})
-
-                # Fetch channel config (if exists)
-                config = session.query(ChannelConfig).filter(
-                    ChannelConfig.channel_id == channel_id,
-                    ChannelConfig.personality_id == personality_id
+            with get_db_session() as s:
+                p = s.query(Personality).filter_by(id=personality_id).first()
+                if not p:
+                    return json.dumps({"error": "personality not found"})
+                cfg = s.query(ChannelConfig).filter_by(
+                    channel_id=channel_id, personality_id=personality_id
                 ).first()
-
-                # Build final prompt: base + channel instructions
-                prompt_parts = [personality.base_prompt]
-                if config and config.channel_instructions:
-                    prompt_parts.append(f"\n\n[CHANNEL RULES]\n{config.channel_instructions}")
-                final_prompt = "\n".join(prompt_parts)
-
+                instr = cfg.channel_instructions if cfg else ""
+                prompt = build_final_prompt(p, instr)
                 return json.dumps({
-                    "prompt": final_prompt,
-                    "fine_tune_config": personality.fine_tune_config or {}
+                    "prompt": prompt,
+                    "fine_tune_config": p.fine_tune_config or {}
                 })
         except Exception as e:
-            return json.dumps({"error": "prompt_generation_failed", "message": str(e)})
-
+            return json.dumps({"error": str(e)})
 
 class InitializePuppetTool(Tool):
-    """Tool to spawn a new Discord puppet (client) with a token."""
     name = "initialize_puppet"
-    description = "Initialize a Discord client (puppet) with a token."
+    description = "Initialize a Discord puppet"
     inputs = {
-        "puppet_id": {"type": "integer", "description": "Unique ID for this puppet (e.g., 1)"},
-        "token": {"type": "string", "description": "Discord account or bot token"},
+        "puppet_id": {"type": "integer"},
+        "token": {"type": "string"}
     }
     output_type = "string"
-
-    def forward(self, puppet_id: int, token: str) -> str:
-        """Returns success status or error."""
-        if puppet_id <= 0:
-            return json.dumps({"error": "puppet_id_must_be_positive"})
-        if not isinstance(token, str) or len(token) < 50:
-            return json.dumps({"error": "invalid_token_format"})
-
+    def forward(self, puppet_id: int, token: str):
+        if puppet_id <= 0 or len(token) < 50:
+            return json.dumps({"error": "invalid puppet_id or token"})
         try:
-            # Clean up existing puppet if present
             if puppet_id in _discum_clients:
-                try:
-                    _discum_clients[puppet_id].close()
-                except:
-                    pass  # Ignore cleanup errors
-
-            # Initialize new client
-            client = discum.Client(token=token, log=False)
-            _discum_clients[puppet_id] = client
-
+                _discum_clients[puppet_id].close()
+            _discum_clients[puppet_id] = discum.Client(token=token, log=False)
             return json.dumps({"status": "success", "puppet_id": puppet_id})
         except Exception as e:
-            return json.dumps({"error": "puppet_initialization_failed", "message": str(e)})
-
+            return json.dumps({"error": str(e)})
 
 class ConfigureChannelTool(Tool):
-    """Tool to bind a personality to a channel for a specific puppet."""
     name = "configure_channel"
-    description = "Assign a personality to a Discord channel with optional instructions."
+    description = "Assign personality to channel"
     inputs = {
-        "puppet_id": {"type": "integer", "description": "Puppet ID (must be initialized)"},
-        "channel_id": {"type": "string", "description": "Discord channel ID"},
-        "personality_id": {"type": "integer", "description": "Personality ID"},
-        "instructions": {"type": "string", "description": "Channel-specific instructions (optional)"},
+        "puppet_id": {"type": "integer"},
+        "channel_id": {"type": "string"},
+        "personality_id": {"type": "integer"},
+        "instructions": {"type": "string"}
     }
     output_type = "string"
-
-    def forward(self, puppet_id: int, channel_id: str, personality_id: int, instructions: str = "") -> str:
-        """Returns success status or error."""
-        # Validate inputs
-        if puppet_id <= 0:
-            return json.dumps({"error": "invalid_puppet_id"})
-        if not validate_discord_id(channel_id):
-            return json.dumps({"error": "invalid_channel_id"})
-        if personality_id <= 0:
-            return json.dumps({"error": "invalid_personality_id"})
-
-        # Verify puppet is alive
-        if puppet_id not in _discum_clients:
-            return json.dumps({"error": "puppet_not_initialized"})
-
+    def forward(self, puppet_id: int, channel_id: str, personality_id: int, instructions: str):
+        if (puppet_id <= 0 or not validate_discord_id(channel_id) or 
+            personality_id <= 0 or puppet_id not in _discum_clients):
+            return json.dumps({"error": "invalid inputs or puppet not initialized"})
         try:
-            with get_db_session() as session:
-                # Verify personality exists
-                personality = session.query(Personality).filter(
-                    Personality.id == personality_id
-                ).first()
-                if not personality:
-                    return json.dumps({"error": "personality_not_found"})
-
-                # Save channel config
-                config = ChannelConfig(
+            with get_db_session() as s:
+                p = s.query(Personality).filter_by(id=personality_id).first()
+                if not p:
+                    return json.dumps({"error": "personality not found"})
+                cfg = ChannelConfig(
                     channel_id=channel_id,
                     personality_id=personality_id,
                     channel_instructions=instructions.strip() or None
                 )
-                session.add(config)
-
+                s.add(cfg)
             return json.dumps({"status": "success"})
         except Exception as e:
-            return json.dumps({"error": "configuration_failed", "message": str(e)})
-
+            return json.dumps({"error": str(e)})
 
 class SendMessageTool(Tool):
-    """Tool for a puppet to send a message to a channel."""
     name = "send_message"
-    description = "Send a message to a Discord channel as a puppet."
+    description = "Send message as puppet"
     inputs = {
-        "puppet_id": {"type": "integer", "description": "Puppet ID"},
-        "channel_id": {"type": "string", "description": "Discord channel ID"},
-        "message": {"type": "string", "description": "Message content"},
+        "puppet_id": {"type": "integer"},
+        "channel_id": {"type": "string"},
+        "message": {"type": "string"}
     }
     output_type = "string"
-
-    def forward(self, puppet_id: int, channel_id: str, message: str) -> str:
-        """Returns success status or error."""
-        if puppet_id <= 0:
-            return json.dumps({"error": "invalid_puppet_id"})
-        if not validate_discord_id(channel_id):
-            return json.dumps({"error": "invalid_channel_id"})
-        if not message or not message.strip():
-            return json.dumps({"error": "empty_message"})
-
-        # Verify puppet is alive
-        if puppet_id not in _discum_clients:
-            return json.dumps({"error": "puppet_not_initialized"})
-
+    def forward(self, puppet_id: int, channel_id: str, message: str):
+        if (puppet_id <= 0 or not validate_discord_id(channel_id) or 
+            not message.strip() or puppet_id not in _discum_clients):
+            return json.dumps({"error": "invalid inputs or puppet not initialized"})
         try:
-            client = _discum_clients[puppet_id]
-            client.sendMessage(channel_id, message.strip())
+            _discum_clients[puppet_id].sendMessage(channel_id, message.strip())
             return json.dumps({"status": "success"})
         except Exception as e:
-            return json.dumps({"error": "message_send_failed", "message": str(e)})
-
+            return json.dumps({"error": str(e)})
 
-# === TOOL REGISTRY: All tools available to the agent ===
 TOOLS = [
     ListPersonalitiesTool(),
     GetChannelPromptTool(),
     InitializePuppetTool(),
     ConfigureChannelTool(),
-    SendMessageTool(),
+    SendMessageTool()
 ]
 
-
-# === AGENT EXECUTOR: smolagents integration ===
-def execute_agent_goal(goal: str) -> str:
-    """
-    Runs the smolagents-powered agent with the user's goal.
-    - Uses gpt-4 by default (change model as needed)
-    - Returns raw agent output (string)
-    - Safe: catches all exceptions
-    """
+# === GRADIO UI ===
+def list_personalities_ui():
+    tool = ListPersonalitiesTool()
+    result = json.loads(tool.forward())
+    if "error" in result:
+        return f"❌ {result['error']}"
+    return "\n".join([f"ID: {p['id']} | {p['name']}" for p in result])
+
+def create_personality_ui(name, base_prompt):
+    if not name.strip() or not base_prompt.strip():
+        return "❌ Name and base prompt required", ""
     try:
-        # Initialize agent with tools
-        agent = Agent(tools=TOOLS, model="gpt-4")  # Or your preferred model
-        result = agent.run(goal)
-        return str(result)
+        with get_db_session() as s:
+            p = Personality(name=name.strip(), base_prompt=base_prompt.strip())
+            s.add(p)
+            s.flush()
+            return f"✅ Created! ID: {p.id}", str(p.id)
     except Exception as e:
-        return f"Agent execution failed: {str(e)}"
+        return f"❌ Failed: {str(e)}", ""
 
+def edit_personality_ui(pid, name, base_prompt, ft_json):
+    if not pid:
+        return "❌ Personality ID required"
+    try:
+        pid = int(pid)
+        ft = json.loads(ft_json) if ft_json.strip() else {}
+        with get_db_session() as s:
+            p = s.query(Personality).filter_by(id=pid).first()
+            if not p:
+                return "❌ Not found"
+            if name.strip(): p.name = name.strip()
+            if base_prompt.strip(): p.base_prompt = base_prompt.strip()
+            p.fine_tune_config = ft
+        return "✅ Updated"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
 
-# === GRADIO INTERFACE: Clean, minimal, functional ===
-with gr.Blocks(title="Puppet Master") as demo:
-    gr.Markdown("# 🎭 Puppet Master — smolagents-Powered Discord Orchestrator")
-    gr.Markdown("Describe your goal. The agent will use tools to manage puppets, personalities, and channels.")
-
-    # User input: natural language goal
-    goal_input = gr.Textbox(
-        label="Your Goal",
-        placeholder="Example: Initialize puppet 1 with my token, assign personality 'Helper' to channel 123456789, and send 'Hello!'",
-        lines=3
-    )
-
-    # Output: agent's actions and results
-    output = gr.Textbox(label="Agent Output", lines=10, interactive=False)
+def initialize_puppet_ui(puppet_id, token):
+    try:
+        pid = int(puppet_id)
+        tool = InitializePuppetTool()
+        result = json.loads(tool.forward(pid, token))
+        if "error" in result:
+            return f"❌ {result['error']}"
+        return f"✅ Puppet {pid} initialized"
+    except Exception as e:
+        return f"❌ {str(e)}"
 
-    # Trigger agent execution
-    gr.Button("Execute Goal").click(
-        fn=execute_agent_goal,
-        inputs=goal_input,
-        outputs=output
-    )
+def configure_channel_ui(puppet_id, channel_id, personality_id, instructions):
+    try:
+        pid = int(puppet_id)
+        per_id = int(personality_id)
+        tool = ConfigureChannelTool()
+        result = json.loads(tool.forward(pid, channel_id, per_id, instructions))
+        if "error" in result:
+            return f"❌ {result['error']}"
+        return "✅ Channel configured"
+    except Exception as e:
+        return f"❌ {str(e)}"
 
-    # Documentation panel (hidden by default in prod; useful for dev)
-    with gr.Accordion("Tool Documentation", open=False):
-        gr.Markdown("""
-        ### Available Tools
-        - `list_personalities()`: Get all personalities
-        - `get_channel_prompt(channel_id, personality_id)`: Build final prompt
-        - `initialize_puppet(puppet_id, token)`: Spawn Discord client
-        - `configure_channel(puppet_id, channel_id, personality_id, instructions)`: Bind personality to channel
-        - `send_message(puppet_id, channel_id, message)`: Send message as puppet
-        """)
+def send_message_ui(puppet_id, channel_id, message):
+    try:
+        pid = int(puppet_id)
+        tool = SendMessageTool()
+        result = json.loads(tool.forward(pid, channel_id, message))
+        if "error" in result:
+            return f"❌ {result['error']}"
+        return "✅ Message sent"
+    except Exception as e:
+        return f"❌ {str(e)}"
 
+with gr.Blocks(title="Puppet Master") as demo:
+    gr.Markdown("# 🎭 Puppet Master — Discord Puppet Orchestrator")
+    
+    with gr.Tabs():
+        with gr.Tab("🧠 Personalities"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Create")
+                    cp_name = gr.Textbox(label="Name")
+                    cp_prompt = gr.Textbox(label="Base Prompt", lines=3)
+                    cp_btn = gr.Button("Create")
+                    cp_status = gr.Textbox(label="Status", interactive=False)
+                    cp_id = gr.Textbox(label="ID", interactive=False)
+                with gr.Column():
+                    gr.Markdown("### Edit")
+                    ep_id = gr.Textbox(label="Personality ID")
+                    ep_name = gr.Textbox(label="New Name (optional)")
+                    ep_prompt = gr.Textbox(label="New Prompt (optional)", lines=3)
+                    ep_ft = gr.Textbox(label="Fine-tune Config (JSON)", lines=2)
+                    ep_btn = gr.Button("Update")
+                    ep_status = gr.Textbox(label="Status", interactive=False)
+            gr.Markdown("### All Personalities")
+            list_btn = gr.Button("Refresh")
+            list_out = gr.Textbox(label="List", lines=6, interactive=False)
+            cp_btn.click(create_personality_ui, [cp_name, cp_prompt], [cp_status, cp_id])
+            ep_btn.click(edit_personality_ui, [ep_id, ep_name, ep_prompt, ep_ft], ep_status)
+            list_btn.click(list_personalities_ui, outputs=list_out)
+
+        with gr.Tab("🤖 Puppets"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Initialize Puppet")
+                    ip_id = gr.Textbox(label="Puppet ID (e.g., 1)")
+                    ip_token = gr.Textbox(label="Discord Token", type="password")
+                    ip_btn = gr.Button("Initialize")
+                    ip_status = gr.Textbox(label="Status", interactive=False)
+                with gr.Column():
+                    gr.Markdown("### Configure Channel")
+                    cc_puppet = gr.Textbox(label="Puppet ID")
+                    cc_ch = gr.Textbox(label="Channel ID")
+                    cc_pid = gr.Textbox(label="Personality ID")
+                    cc_instr = gr.Textbox(label="Channel Instructions", lines=2)
+                    cc_btn = gr.Button("Assign")
+                    cc_status = gr.Textbox(label="Status", interactive=False)
+            ip_btn.click(initialize_puppet_ui, [ip_id, ip_token], ip_status)
+            cc_btn.click(configure_channel_ui, [cc_puppet, cc_ch, cc_pid, cc_instr], cc_status)
+
+        with gr.Tab("📤 Send Message"):
+            with gr.Row():
+                with gr.Column():
+                    sm_puppet = gr.Textbox(label="Puppet ID")
+                    sm_ch = gr.Textbox(label="Channel ID")
+                    sm_msg = gr.Textbox(label="Message", lines=3)
+                    sm_btn = gr.Button("Send")
+                    sm_status = gr.Textbox(label="Status", interactive=False)
+            sm_btn.click(send_message_ui, [sm_puppet, sm_ch, sm_msg], sm_status)
 
-# === ENTRY POINT: Production-ready launch ===
 if __name__ == "__main__":
-    # Launch with security best practices:
-    # - Bind to localhost only (change to "0.0.0.0" for remote access)
-    # - Disable sharing (public links are dangerous)
-    # - Enable error display for debugging
-    demo.launch(
-        server_name="127.0.0.1",  # Localhost only
-        server_port=7860,
-        show_error=True,
-        share=False  # Never enable in production
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)