Update app.py

app.py

@@ -1,58 +1,35 @@
 import os
-import json
 import requests
 from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse, JSONResponse
-from fastapi.staticfiles import StaticFiles
+from fastapi.responses import JSONResponse, HTMLResponse
 from fastapi.templating import Jinja2Templates
 
 # =========================
-# APP INIT (HARDENED)
+# APP INIT
 # =========================
 
-app = FastAPI(title="PyRunner AI Space", version="2.0")
+app = FastAPI(title="PyRunner Agent OS", version="3.0")
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-
-TEMPLATE_DIR = os.path.join(BASE_DIR, "templates")
-STATIC_DIR = os.path.join(BASE_DIR, "static")
-
-# SAFE TEMPLATE LOADER (FIXES YOUR ERROR CLASS)
-templates = Jinja2Templates(directory=TEMPLATE_DIR)
-
-if os.path.exists(STATIC_DIR):
-    app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
+templates = Jinja2Templates(directory=os.path.join(BASE_DIR, "templates"))
 
 
 # =========================
-# ENV LOADER (SAFE DEFAULTS)
+# ENV
 # =========================
 
-def env(key, default=None):
-    value = os.getenv(key, default)
-    if value == "None":
-        return None
-    return value
-
-
-OLLAMA_URL = env("OLLAMA_URL", "http://localhost:11434")
-OLLAMA_FALLBACK_URL = env("OLLAMA_FALLBACK_URL")
+OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
 
-OLLAMA_MODEL = env("OLLAMA_MODEL", "qwen2.5:1.5b")
-OLLAMA_CODE_MODEL = env("OLLAMA_CODE_MODEL", "qwen2.5-coder:1.5b")
-OLLAMA_FAST_MODEL = env("OLLAMA_FAST_MODEL", "deepseek-coder:1.3b")
-OLLAMA_BACKUP_MODEL = env("OLLAMA_BACKUP_MODEL", "llama3.2:1b")
-
-OPENROUTER_API_KEY = env("OPENROUTER_API_KEY")
-OPENROUTER_MODEL = env("OPENROUTER_MODEL")
-
-SUPABASE_URL = env("SUPABASE_URL")
-SUPABASE_KEY = env("SUPABASE_KEY")
-N8N_WEBHOOK_URL = env("N8N_WEBHOOK_URL")
+MODELS = {
+    "orchestrator": os.getenv("OLLAMA_MODEL", "qwen2.5:1.5b"),
+    "coder": os.getenv("OLLAMA_CODE_MODEL", "qwen2.5-coder:1.5b"),
+    "fast": os.getenv("OLLAMA_FAST_MODEL", "deepseek-coder:1.3b"),
+    "fallback": os.getenv("OLLAMA_BACKUP_MODEL", "llama3.2:1b")
+}
 
 
 # =========================
-# OLLAMA CLIENT (HARDENED)
+# OLLAMA CORE CALL
 # =========================
 
 def call_ollama(prompt, model):
@@ -68,198 +45,148 @@ def call_ollama(prompt, model):
         )
         return r.json()
     except Exception as e:
-        return {"error": str(e), "source": "ollama_primary"}
+        return {"error": str(e)}
 
 
-def call_ollama_fallback(prompt):
-    if not OLLAMA_FALLBACK_URL:
-        return {"error": "No fallback URL configured"}
+# =========================
+# AGENT DEFINITIONS
+# =========================
 
-    try:
-        r = requests.post(
-            f"{OLLAMA_FALLBACK_URL}/api/generate",
-            json={
-                "model": OLLAMA_BACKUP_MODEL,
-                "prompt": prompt,
-                "stream": False
-            },
-            timeout=60
-        )
-        return r.json()
-    except Exception as e:
-        return {"error": str(e), "source": "ollama_fallback"}
+def orchestrator_agent(user_input: str):
+    """
+    Classifies task type
+    """
 
+    prompt = f"""
+You are an orchestrator AI.
 
-# =========================
-# OPENROUTER (OPTIONAL)
-# =========================
+Classify this request into one category ONLY:
 
-def call_openrouter(prompt):
-    if not OPENROUTER_API_KEY:
-        return {"error": "OpenRouter not configured"}
+- code
+- vision
+- fashion
+- hairstyle
+- general
+
+User request:
+{user_input}
+
+Return ONLY the category name.
+"""
+
+    res = call_ollama(prompt, MODELS["orchestrator"])
+
+    if "response" in res:
+        return res["response"].strip().lower()
+
+    return "general"
+
+
+def coder_agent(prompt):
+    return call_ollama(prompt, MODELS["coder"])
 
-    try:
-        r = requests.post(
-            "https://openrouter.ai/api/v1/chat/completions",
-            headers={
-                "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-                "Content-Type": "application/json"
-            },
-            json={
-                "model": OPENROUTER_MODEL,
-                "messages": [{"role": "user", "content": prompt}]
-            },
-            timeout=60
-        )
-        return r.json()
-    except Exception as e:
-        return {"error": str(e), "source": "openrouter"}
+
+def fashion_agent(prompt):
+    system_prompt = f"""
+You are a fashion expert AI.
+Analyze outfit, styling, colors, trends.
+
+User:
+{prompt}
+"""
+    return call_ollama(system_prompt, MODELS["fast"])
+
+
+def hairstyle_agent(prompt):
+    system_prompt = f"""
+You are a barber & hairstyle expert.
+Describe hairstyle changes, fades, textures, transformations.
+
+User:
+{prompt}
+"""
+    return call_ollama(system_prompt, MODELS["fast"])
+
+
+def fallback_agent(prompt):
+    return call_ollama(prompt, MODELS["fallback"])
 
 
 # =========================
-# SMART ROUTER (1B–2B ONLY SAFE)
+# ROUTER
 # =========================
 
-def route_model(mode: str):
-    mode = (mode or "").lower()
+def route_agent(category: str, prompt: str):
 
-    if "code" in mode:
-        return OLLAMA_CODE_MODEL
+    if category == "code":
+        return coder_agent(prompt)
 
-    if "fast" in mode:
-        return OLLAMA_FAST_MODEL
+    if category == "fashion":
+        return fashion_agent(prompt)
 
-    if "backup" in mode or "fallback" in mode:
-        return OLLAMA_BACKUP_MODEL
+    if category == "hairstyle":
+        return hairstyle_agent(prompt)
 
-    return OLLAMA_MODEL
+    return fallback_agent(prompt)
 
 
 # =========================
-# UI ROUTE (SAFE)
+# UI
 # =========================
 
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
-    try:
-        return templates.TemplateResponse(
-            "index.html",
-            {"request": request}
-        )
-    except Exception as e:
-        # HARD FALLBACK (prevents total crash)
-        return HTMLResponse(f"""
-        <h1>PyRunner Running</h1>
-        <p>Template error fallback active</p>
-        <pre>{str(e)}</pre>
-        """)
+    return templates.TemplateResponse("index.html", {"request": request})
 
 
 # =========================
-# HEALTH CHECK
+# HEALTH
 # =========================
 
 @app.get("/health")
 def health():
     return {
         "status": "ok",
-        "ollama": OLLAMA_URL,
-        "models": {
-            "default": OLLAMA_MODEL,
-            "code": OLLAMA_CODE_MODEL,
-            "fast": OLLAMA_FAST_MODEL,
-            "backup": OLLAMA_BACKUP_MODEL
-        }
+        "agents": list(MODELS.keys())
     }
 
 
 # =========================
-# CHAT ENDPOINT (IMPROVED FLOW)
+# MAIN AGENT ENDPOINT
 # =========================
 
-@app.post("/api/chat")
-async def chat(request: Request):
-    body = await request.json()
+@app.post("/api/agent")
+async def agent(request: Request):
 
-    prompt = body.get("prompt", "")
-    mode = body.get("mode", "default")
+    body = await request.json()
+    user_input = body.get("prompt", "")
 
-    model = route_model(mode)
+    # 1. Decide agent
+    category = orchestrator_agent(user_input)
 
-    # 1. Try primary Ollama
-    result = call_ollama(prompt, model)
+    # 2. Route to agent
+    result = route_agent(category, user_input)
 
-    # 2. Fallback to local backup
+    # 3. Fallback safety
     if "error" in result:
-        result = call_ollama_fallback(prompt)
+        result = fallback_agent(user_input)
 
-    # 3. Optional OpenRouter fallback (last resort)
-    if "error" in result:
-        result = call_openrouter(prompt)
-
-    return JSONResponse(result)
+    return JSONResponse({
+        "category": category,
+        "result": result
+    })
 
 
 # =========================
-# IMAGE PIPELINE (STUB READY)
+# IMAGE PIPELINE HOOK (future)
 # =========================
 
 @app.post("/api/image")
 async def image_pipeline(request: Request):
     body = await request.json()
-    prompt = body.get("prompt", "")
 
     return {
         "status": "queued",
-        "pipeline": "not_connected",
-        "prompt": prompt,
-        "note": "Connect to diffusion/vision service next"
-    }
-
-
-# =========================
-# N8N WEBHOOK
-# =========================
-
-@app.post("/api/webhook/n8n")
-async def trigger_n8n(request: Request):
-    if not N8N_WEBHOOK_URL:
-        return {"error": "N8N webhook not set"}
-
-    body = await request.json()
-
-    try:
-        r = requests.post(N8N_WEBHOOK_URL, json=body, timeout=30)
-        return {
-            "status": "sent",
-            "response": r.text
-        }
-    except Exception as e:
-        return {"error": str(e)}
-
-
-# =========================
-# DEBUG CONFIG
-# =========================
-
-@app.get("/config")
-def config():
-    return {
-        "ollama_url": OLLAMA_URL,
-        "supabase": bool(SUPABASE_URL and SUPABASE_KEY),
-        "openrouter": bool(OPENROUTER_API_KEY),
-        "n8n": bool(N8N_WEBHOOK_URL),
-        "models": {
-            "default": OLLAMA_MODEL,
-            "code": OLLAMA_CODE_MODEL
-        }
-    }
-
-
-# =========================
-# SIMPLE TEST ROUTE (IMPORTANT FOR DEBUGGING)
-# =========================
-
-@app.get("/ping")
-def ping():
-    return {"ping": "ok"}
+        "pipeline": "vision_stub",
+        "note": "Attach diffusion model later"
+    }