Update theme/app.py

theme/app.py

@@ -8,12 +8,11 @@ from fastapi.templating import Jinja2Templates
 # APP INIT
 # =========================
 
-app = FastAPI(title="PyRunner Agent OS", version="3.0")
+app = FastAPI(title="PyRunner Agent OS", version="3.1")
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 templates = Jinja2Templates(directory=os.path.join(BASE_DIR, "templates"))
 
-
 # =========================
 # ENV
 # =========================
@@ -27,12 +26,14 @@ MODELS = {
     "fallback": os.getenv("OLLAMA_BACKUP_MODEL", "llama3.2:1b")
 }
 
+TIMEOUT = (5, 120)  # connect, read
+
 
 # =========================
-# OLLAMA CORE CALL
+# SAFE OLLAMA CALL
 # =========================
 
-def call_ollama(prompt, model):
+def call_ollama(prompt: str, model: str):
     try:
         r = requests.post(
             f"{OLLAMA_URL}/api/generate",
@@ -41,71 +42,91 @@ def call_ollama(prompt, model):
                 "prompt": prompt,
                 "stream": False
             },
-            timeout=60
+            timeout=TIMEOUT
         )
-        return r.json()
+
+        if r.status_code != 200:
+            return {
+                "error": f"Ollama error {r.status_code}",
+                "raw": r.text
+            }
+
+        data = r.json()
+
+        # normalize output
+        return {
+            "response": data.get("response", "").strip(),
+            "raw": data
+        }
+
     except Exception as e:
         return {"error": str(e)}
 
 
 # =========================
-# AGENT DEFINITIONS
+# CLEAN OUTPUT PARSER
 # =========================
 
-def orchestrator_agent(user_input: str):
-    """
-    Classifies task type
-    """
+def clean_text(text: str) -> str:
+    if not text:
+        return "general"
 
+    return (
+        text.lower()
+        .replace("\n", "")
+        .replace(".", "")
+        .strip()
+    )
+
+
+# =========================
+# ORCHESTRATOR
+# =========================
+
+def orchestrator_agent(user_input: str):
     prompt = f"""
-You are an orchestrator AI.
+You are an AI classifier.
 
-Classify this request into one category ONLY:
+Return ONLY ONE word:
 
-- code
-- vision
-- fashion
-- hairstyle
-- general
+code
+vision
+fashion
+hairstyle
+general
 
-User request:
+User:
 {user_input}
-
-Return ONLY the category name.
 """
 
     res = call_ollama(prompt, MODELS["orchestrator"])
 
     if "response" in res:
-        return res["response"].strip().lower()
+        return clean_text(res["response"])
 
     return "general"
 
 
+# =========================
+# AGENTS
+# =========================
+
 def coder_agent(prompt):
     return call_ollama(prompt, MODELS["coder"])
 
 
 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"])
+    return call_ollama(
+        f"Fashion analysis:\n{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"])
+    return call_ollama(
+        f"Hairstyle analysis:\n{prompt}",
+        MODELS["fast"]
+    )
 
 
 def fallback_agent(prompt):
@@ -136,7 +157,10 @@ def route_agent(category: str, prompt: str):
 
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request})
+    return templates.TemplateResponse(
+        "index.html",
+        {"request": request}
+    )
 
 
 # =========================
@@ -147,7 +171,7 @@ async def home(request: Request):
 def health():
     return {
         "status": "ok",
-        "agents": list(MODELS.keys())
+        "models": MODELS
     }
 
 
@@ -161,32 +185,34 @@ async def agent(request: Request):
     body = await request.json()
     user_input = body.get("prompt", "")
 
-    # 1. Decide agent
+    if not user_input:
+        return JSONResponse({
+            "error": "No prompt provided"
+        }, status_code=400)
+
+    # 1. classify
     category = orchestrator_agent(user_input)
 
-    # 2. Route to agent
+    # 2. route
     result = route_agent(category, user_input)
 
-    # 3. Fallback safety
-    if "error" in result:
+    # 3. fallback safety
+    if isinstance(result, dict) and "error" in result:
         result = fallback_agent(user_input)
 
     return JSONResponse({
         "category": category,
-        "result": result
+        "result": result.get("response", result)
     })
 
 
 # =========================
-# IMAGE PIPELINE HOOK (future)
+# IMAGE PIPELINE HOOK
 # =========================
 
 @app.post("/api/image")
 async def image_pipeline(request: Request):
-    body = await request.json()
-
     return {
         "status": "queued",
-        "pipeline": "vision_stub",
-        "note": "Attach diffusion model later"
+        "pipeline": "vision_stub"
     }