Create app.py

app.py

@@ -0,0 +1,520 @@
+import os
+import time
+import threading
+import socket
+import subprocess
+import asyncio
+import json
+import uuid
+import requests
+
+from fastapi import FastAPI, Request, Header
+from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
+
+app = FastAPI()
+
+# ----------------------------
+# Config
+# ----------------------------
+API_KEY = os.getenv("API_KEY", "821274")
+MODEL = os.getenv("MODEL", "llama3.2:1b")
+OLLAMA_BASE = os.getenv("OLLAMA_BASE", "http://127.0.0.1:11434")
+
+SYSTEM_PROMPT = """
+You are a helpful, friendly AI assistant.
+
+Rules:
+- If the user asks casual conversation, respond normally in plain text.
+- If the user asks for code/scripts/automation or says “just do it” / “write something”, choose a reasonable task and output a complete useful script.
+- Don’t refuse normal conversation.
+- Only output code when appropriate.
+- When you write code: output code first, then a short explanation.
+
+Always be helpful. Never say you cannot fulfill a request unless it is unsafe.
+""".strip()
+
+# Free CPU: serialize requests to avoid overload / timeouts
+GEN_SEM = asyncio.Semaphore(1)
+
+# ----------------------------
+# Ollama helpers
+# ----------------------------
+def is_port_open(host="127.0.0.1", port=11434) -> bool:
+    try:
+        with socket.create_connection((host, port), timeout=0.5):
+            return True
+    except OSError:
+        return False
+
+def ollama_healthy() -> bool:
+    try:
+        r = requests.get(f"{OLLAMA_BASE}/api/tags", timeout=1.5)
+        return r.status_code == 200
+    except Exception:
+        return False
+
+def ensure_ollama_running():
+    # Only start if not reachable
+    if not is_port_open("127.0.0.1", 11434):
+        subprocess.Popen(["ollama", "serve"])
+
+def wait_for_ollama(timeout_s=120) -> bool:
+    start = time.time()
+    while time.time() - start < timeout_s:
+        if ollama_healthy():
+            return True
+        time.sleep(1)
+    return False
+
+def pull_and_warm_model():
+    """
+    Best-effort: pull the model (may be slow on free CPU) and warm it once.
+    Safe to fail (space still boots).
+    """
+    try:
+        ensure_ollama_running()
+        if not wait_for_ollama(120):
+            print("Ollama not ready yet; skipping model pull.")
+            return
+
+        print(f"Pulling model: {MODEL}")
+        r = requests.post(f"{OLLAMA_BASE}/api/pull", json={"name": MODEL}, timeout=60 * 30)
+        if r.status_code != 200:
+            print("Pull failed:", r.text[:2000])
+            return
+
+        # Warmup: avoids first real user request being extra flaky/slow
+        print("Warming up…")
+        requests.post(
+            f"{OLLAMA_BASE}/api/generate",
+            json={"model": MODEL, "system": SYSTEM_PROMPT, "prompt": "Say: ready.", "stream": False},
+            timeout=180,
+        )
+        print("Warmup done.")
+    except Exception as e:
+        print("Boot task error (non-fatal):", str(e))
+
+threading.Thread(target=pull_and_warm_model, daemon=True).start()
+
+def generate_with_recovery(prompt: str, attempts: int = 3):
+    last_err = None
+    for i in range(1, attempts + 1):
+        try:
+            if not ollama_healthy():
+                ensure_ollama_running()
+                wait_for_ollama(60)
+
+            r = requests.post(
+                f"{OLLAMA_BASE}/api/generate",
+                json={
+                    "model": MODEL,
+                    "system": SYSTEM_PROMPT,
+                    "prompt": prompt,
+                    "stream": False,
+                },
+                timeout=600,
+            )
+            r.raise_for_status()
+            data = r.json()
+            return data.get("response", ""), None
+        except Exception as e:
+            last_err = str(e)
+            time.sleep(min(2 ** (i - 1), 4))
+
+    return (
+        "⚠️ Backend hiccup while generating. Retrying usually works.\n\n"
+        "Debug error:\n" + (last_err or "unknown"),
+        last_err,
+    )
+
+def messages_to_prompt(messages):
+    """
+    Convert OpenAI-style messages into a single prompt string for Ollama /api/generate.
+    """
+    parts = []
+    for m in messages or []:
+        role = (m.get("role") or "user").strip().upper()
+        content = (m.get("content") or "").strip()
+        if content:
+            parts.append(f"{role}:\n{content}")
+    return "\n\n".join(parts).strip()
+
+# ----------------------------
+# Health
+# ----------------------------
+@app.get("/health")
+def health():
+    return {"ok": ollama_healthy(), "model": MODEL}
+
+# ----------------------------
+# OpenAI compatibility endpoints
+# ----------------------------
+@app.get("/v1/models")
+def openai_models(authorization: str = Header(default="")):
+    # Optional auth check (recommended)
+    if authorization and authorization != f"Bearer {API_KEY}":
+        return JSONResponse({"error": {"message": "Invalid API key"}}, status_code=401)
+
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": MODEL,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "private-ai"
+            }
+        ]
+    }
+
+@app.post("/v1/chat/completions")
+async def openai_chat_completions(request: Request, authorization: str = Header(default="")):
+    # OpenAI-style auth
+    if authorization != f"Bearer {API_KEY}":
+        return JSONResponse(
+            {"error": {"message": "Invalid API key", "type": "auth_error"}},
+            status_code=401
+        )
+
+    body = await request.json()
+    model = body.get("model") or MODEL
+    messages = body.get("messages") or []
+
+    # Convert into a single prompt for Ollama generate
+    prompt = messages_to_prompt(messages)
+    if not prompt:
+        prompt = "USER:\nHello"
+
+    async with GEN_SEM:
+        text, err = generate_with_recovery(prompt, attempts=3)
+
+    # If error, still return valid OpenAI shaped response
+    if err:
+        text += f"\n\n---\nBackend error:\n{err}"
+
+    return {
+        "id": f"chatcmpl-{uuid.uuid4().hex}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": text},
+                "finish_reason": "stop"
+            }
+        ],
+        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+    }
+
+# ----------------------------
+# UI (REAL HTML, not escaped)
+# ----------------------------
+@app.get("/", response_class=HTMLResponse)
+def ui():
+    return f"""<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <title>Private AI</title>
+  <style>
+    :root {{
+      --bg:#0b0f17;
+      --panel:rgba(255,255,255,.06);
+      --border:rgba(255,255,255,.10);
+      --text:rgba(255,255,255,.92);
+      --muted:rgba(255,255,255,.55);
+      --radius:18px;
+      --shadow:0 12px 40px rgba(0,0,0,.35);
+    }}
+    *{{box-sizing:border-box}}
+    body{{
+      margin:0;height:100vh;overflow:hidden;
+      font-family:ui-sans-serif,system-ui,-apple-system,Segoe UI,Roboto;
+      background:
+        radial-gradient(900px 600px at 20% 0%, rgba(59,130,246,.18), transparent 60%),
+        radial-gradient(800px 600px at 80% 0%, rgba(168,85,247,.16), transparent 60%),
+        var(--bg);
+      color:var(--text);
+    }}
+    .app{{display:grid;grid-template-rows:auto 1fr auto;height:100vh}}
+    header{{
+      padding:14px 18px;border-bottom:1px solid var(--border);
+      backdrop-filter:blur(14px);background:rgba(10,14,22,.6);
+    }}
+    header .inner{{max-width:980px;margin:0 auto;display:flex;align-items:center;justify-content:space-between}}
+    header h1{{font-size:14px;margin:0;font-weight:650}}
+    .status{{font-size:12px;color:var(--muted);display:flex;gap:8px;align-items:center}}
+    .dot{{width:8px;height:8px;border-radius:50%;background:#555;box-shadow:0 0 0 6px rgba(255,255,255,.06)}}
+    .dot.online{{background:#22c55e;box-shadow:0 0 0 6px rgba(34,197,94,.12)}}
+    .dot.busy{{background:#60a5fa;box-shadow:0 0 0 6px rgba(96,165,250,.14)}}
+    .dot.degraded{{background:#f59e0b;box-shadow:0 0 0 6px rgba(245,158,11,.14)}}
+    main{{overflow-y:auto;padding:20px 18px}}
+    .chat{{max-width:980px;margin:0 auto}}
+    .msg{{display:flex;gap:12px;margin:12px 0}}
+    .avatar{{width:36px;height:36px;border-radius:10px;background:var(--panel);border:1px solid var(--border);display:grid;place-items:center;font-size:13px}}
+    .bubble{{padding:12px 14px;border-radius:var(--radius);background:var(--panel);border:1px solid var(--border);box-shadow:var(--shadow);font-size:14px;line-height:1.45;white-space:pre-wrap;overflow-wrap:anywhere}}
+    .me .bubble{{background:linear-gradient(180deg,rgba(59,130,246,.22),rgba(255,255,255,.06));border-color:rgba(59,130,246,.25)}}
+    footer{{padding:14px 18px 18px;border-top:1px solid var(--border);backdrop-filter:blur(14px);background:rgba(10,14,22,.6)}}
+    .composer{{max-width:980px;margin:0 auto;display:grid;grid-template-columns:1fr auto;gap:10px}}
+    textarea{{resize:none;min-height:44px;max-height:180px;padding:12px 14px;border-radius:16px;border:1px solid var(--border);background:var(--panel);color:var(--text);font-size:14px;outline:none}}
+    button{{height:44px;padding:0 16px;border-radius:16px;border:1px solid var(--border);background:rgba(255,255,255,.08);color:var(--text);font-weight:650;cursor:pointer}}
+    button:disabled{{opacity:.6}}
+    .hint{{max-width:980px;margin:10px auto 0;color:rgba(255,255,255,.55);font-size:12px}}
+    kbd{{font-family:ui-monospace,SFMono-Regular,Menlo,Monaco,Consolas,"Liberation Mono","Courier New",monospace;font-size:12px;padding:2px 6px;border-radius:8px;background:rgba(255,255,255,.08);border:1px solid rgba(255,255,255,.12)}}
+  </style>
+</head>
+<body>
+  <div class="app">
+    <header>
+      <div class="inner">
+        <h1>Private AI</h1>
+        <div class="status">
+          <span class="dot" id="dot"></span>
+          <span id="status">Connecting…</span>
+        </div>
+      </div>
+    </header>
+
+    <main><div class="chat" id="chat"></div></main>
+
+    <footer>
+      <div class="composer">
+        <textarea id="input" placeholder="Message Private AI…"></textarea>
+        <button id="send">Send</button>
+      </div>
+      <div class="hint" id="hint"></div>
+    </footer>
+  </div>
+
+<script>
+const chat = document.getElementById("chat");
+const input = document.getElementById("input");
+const sendBtn = document.getElementById("send");
+const dot = document.getElementById("dot");
+const statusEl = document.getElementById("status");
+const hint = document.getElementById("hint");
+
+let failCount = 0;
+
+function setStatus(mode, txt) {{
+  dot.classList.remove("online","busy","degraded");
+  if (mode === "online") dot.classList.add("online");
+  if (mode === "busy") dot.classList.add("busy");
+  if (mode === "degraded") dot.classList.add("degraded");
+  statusEl.textContent = txt;
+}}
+
+function add(role, text) {{
+  const el = document.createElement("div");
+  el.className = "msg " + role;
+  el.innerHTML = `
+    <div class="avatar">${{role === "me" ? "You" : "AI"}}</div>
+    <div class="bubble"></div>
+  `;
+  el.querySelector(".bubble").textContent = text || "";
+  chat.appendChild(el);
+  chat.scrollTop = chat.scrollHeight;
+  return el;
+}}
+
+async function healthCheck() {{
+  try {{
+    const r = await fetch("/health", {{ cache: "no-store" }});
+    const d = await r.json();
+    if (!r.ok || !d.ok) throw new Error("unhealthy");
+    failCount = 0;
+    setStatus("online", "Online");
+    return true;
+  }} catch {{
+    failCount++;
+    setStatus("degraded", failCount >= 5 ? "Recovering…" : "Reconnecting…");
+    return false;
+  }}
+}}
+
+function getKey() {{
+  let k = localStorage.getItem("API_KEY");
+  if (!k) {{
+    k = prompt("Enter API key:");
+    if (k) localStorage.setItem("API_KEY", k);
+  }}
+  return k || "";
+}}
+
+async function sendStream(msg, bubbleEl) {{
+  const key = getKey();
+  const r = await fetch("/v1/chat/stream", {{
+    method: "POST",
+    headers: {{
+      "Authorization": "Bearer " + key,
+      "Content-Type": "application/json"
+    }},
+    body: JSON.stringify({{ prompt: msg }})
+  }});
+
+  if (!r.ok || !r.body) throw new Error("No stream");
+
+  const reader = r.body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+
+  while (true) {{
+    const {{ value, done }} = await reader.read();
+    if (done) break;
+
+    buffer += decoder.decode(value, {{ stream: true }});
+    const parts = buffer.split("\\n\\n");
+    buffer = parts.pop() || "";
+
+    for (const part of parts) {{
+      const lines = part.split("\\n");
+      let event = "message";
+      let dataLine = "";
+
+      for (const line of lines) {{
+        if (line.startsWith("event:")) event = line.slice(6).trim();
+        if (line.startsWith("data:")) dataLine += line.slice(5).trim();
+      }}
+
+      if (!dataLine) continue;
+      const payload = JSON.parse(dataLine);
+
+      if (event === "error") {{
+        bubbleEl.textContent += "\\n\\n---\\nError:\\n" + (payload.error || "Unknown error");
+        setStatus("degraded", "Recovering…");
+      }} else if (event === "done") {{
+        setStatus("online", "Online");
+      }} else {{
+        bubbleEl.textContent += payload.delta || "";
+      }}
+    }}
+  }}
+}}
+
+async function send() {{
+  const msg = input.value.trim();
+  if (!msg) return;
+
+  input.value = "";
+  add("me", msg);
+
+  setStatus("busy", "Thinking…");
+  const aiMsg = add("ai", "");
+  const bubble = aiMsg.querySelector(".bubble");
+  sendBtn.disabled = true;
+
+  try {{
+    await sendStream(msg, bubble);
+  }} catch (e) {{
+    bubble.textContent =
+      "Temporary error. Try again in a moment.\\n\\nTip: verify your API key (stored in browser localStorage).";
+    setStatus("degraded", "Reconnecting…");
+  }} finally {{
+    sendBtn.disabled = false;
+    setTimeout(healthCheck, 800);
+  }}
+}}
+
+sendBtn.onclick = send;
+input.addEventListener("keydown", (e) => {{
+  if (e.key === "Enter" && !e.shiftKey) {{
+    e.preventDefault();
+    send();
+  }}
+}});
+
+hint.innerHTML = "Press <kbd>Enter</kbd> to send, <kbd>Shift</kbd>+<kbd>Enter</kbd> for newline.";
+add("ai", "Hey 👋 I’m ready when you are.");
+healthCheck();
+setInterval(healthCheck, 5000);
+</script>
+</body>
+</html>
+"""
+
+# ----------------------------
+# Streaming SSE endpoint
+# ----------------------------
+@app.post("/v1/chat/stream")
+async def chat_stream(request: Request):
+    auth = request.headers.get("authorization", "")
+    if auth != f"Bearer {API_KEY}":
+        async def deny():
+            yield "event: error\ndata: " + json.dumps({"error": "403: Invalid API key"}) + "\n\n"
+            yield "event: done\ndata: {}\n\n"
+        return StreamingResponse(deny(), media_type="text/event-stream")
+
+    body = await request.json()
+    prompt = (body.get("prompt") or "").strip()
+    if not prompt:
+        async def empty():
+            yield "data: " + json.dumps({"delta": "Send a message and I’ll respond."}) + "\n\n"
+            yield "event: done\ndata: {}\n\n"
+        return StreamingResponse(empty(), media_type="text/event-stream")
+
+    async def event_gen():
+        async with GEN_SEM:
+            try:
+                if not ollama_healthy():
+                    ensure_ollama_running()
+                    wait_for_ollama(60)
+
+                r = requests.post(
+                    f"{OLLAMA_BASE}/api/generate",
+                    json={
+                        "model": MODEL,
+                        "system": SYSTEM_PROMPT,
+                        "prompt": prompt,
+                        "stream": True,
+                    },
+                    stream=True,
+                    timeout=600,
+                )
+
+                if r.status_code != 200:
+                    yield "event: error\ndata: " + json.dumps({"error": r.text[:2000]}) + "\n\n"
+                    yield "event: done\ndata: {}\n\n"
+                    return
+
+                for line in r.iter_lines(decode_unicode=True):
+                    if not line:
+                        continue
+                    try:
+                        obj = json.loads(line)
+                    except Exception:
+                        continue
+
+                    delta = obj.get("response", "")
+                    if delta:
+                        yield "data: " + json.dumps({"delta": delta}) + "\n\n"
+
+                    if obj.get("done"):
+                        break
+
+                yield "event: done\ndata: {}\n\n"
+
+            except Exception as e:
+                yield "event: error\ndata: " + json.dumps({"error": str(e)}) + "\n\n"
+                yield "event: done\ndata: {}\n\n"
+
+    return StreamingResponse(event_gen(), media_type="text/event-stream")
+
+# ----------------------------
+# Non-stream fallback
+# ----------------------------
+@app.post("/v1/chat")
+async def chat_api(request: Request):
+    auth = request.headers.get("authorization", "")
+    if auth != f"Bearer {API_KEY}":
+        return JSONResponse({"response": "", "error": "403: Invalid API key"}, status_code=200)
+
+    body = await request.json()
+    prompt = (body.get("prompt") or "").strip()
+    if not prompt:
+        return {"response": "Send a message and I’ll respond.", "error": None}
+
+    async with GEN_SEM:
+        text, err = generate_with_recovery(prompt, attempts=3)
+        return {"response": text, "error": err}