Create app.py

app.py

@@ -0,0 +1,420 @@
+"""
+Yappy - AI motivation/discipline/support chatbot backend
+Created by Rik
+
+Serves Qwen2.5-3B-Instruct (GGUF, CPU) via llama-cpp-python, streamed over
+Server-Sent Events.
+
+STATELESS BY DESIGN: this server keeps no conversation memory of its own.
+The frontend stores chat history in the browser's localStorage (per mode)
+and sends the relevant history back with every request. This means chat
+history survives Space sleep/restarts, since it never depended on the
+server remembering anything.
+"""
+
+import os
+from typing import Optional, List, Dict, Literal
+
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from sse_starlette.sse import EventSourceResponse
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
+import threading
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+MODEL_REPO = os.environ.get("MODEL_REPO", "Qwen/Qwen2.5-3B-Instruct-GGUF")
+MODEL_FILE = os.environ.get("MODEL_FILE", "qwen2.5-3b-instruct-q4_k_m.gguf")
+MODEL_DIR = os.environ.get("MODEL_DIR", "/app/models")
+
+N_CTX = int(os.environ.get("N_CTX", "8192"))
+MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "900"))
+N_THREADS = int(os.environ.get("N_THREADS", str(os.cpu_count() or 4)))
+
+API_KEY = os.environ.get("YAPPY_API_KEY", "")
+
+MAX_HISTORY_MESSAGES = int(os.environ.get("MAX_HISTORY_MESSAGES", "24"))
+
+# ---------------------------------------------------------------------------
+# Model loading
+# ---------------------------------------------------------------------------
+
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+print(f"[Yappy] Downloading/locating model {MODEL_REPO}/{MODEL_FILE} ...")
+model_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=MODEL_FILE,
+    local_dir=MODEL_DIR,
+)
+print(f"[Yappy] Model ready at {model_path}")
+
+print("[Yappy] Loading model into llama.cpp ...")
+llm = Llama(
+    model_path=model_path,
+    n_ctx=N_CTX,
+    n_threads=N_THREADS,
+    n_gpu_layers=0,
+    chat_format="chatml",
+    verbose=False,
+)
+print("[Yappy] Model loaded. Ready to yap.")
+
+model_lock = threading.Lock()
+
+# ---------------------------------------------------------------------------
+# Mode system prompts
+# ---------------------------------------------------------------------------
+
+BASE_IDENTITY = (
+    "You are Yappy, an AI companion created by Rik. You are not a generic "
+    "assistant - you have a distinct personality that adapts to the mode "
+    "the user picked, but you always remember you are Yappy. Keep replies "
+    "conversational (like texting a person), not like an encyclopedia. "
+    "Never use markdown headers or bullet-point essays unless the user "
+    "explicitly asks for a list. Speak in plain paragraphs or short lines.\n\n"
+    "Before every reply, think first. Put your private reasoning inside "
+    "<thinking>...</thinking> tags - in there, briefly work out what the "
+    "user actually needs, what mode-appropriate angle to take, and what "
+    "tone fits. Keep the thinking short (2-5 sentences), not a full essay. "
+    "Immediately after the closing </thinking> tag, write your real reply "
+    "to the user as normal text, with no other tags. Always include both "
+    "the <thinking> block and the reply - never skip the thinking block."
+)
+
+MODE_PROMPTS = {
+    "motivation": (
+        BASE_IDENTITY + "\n\n"
+        "MODE: MOTIVATION (roast-to-motivate).\n"
+        "Your job is to fire the user up by roasting them - hard. Be blunt, "
+        "sarcastic, and call out their excuses directly. Mock the excuse, "
+        "not the person's worth as a human being. No slurs, no attacks on "
+        "protected traits (race, gender, disability, etc), no body-shaming, "
+        "and never mock someone who mentions real distress, self-harm, or "
+        "a mental health crisis - if that happens, drop the roast instantly "
+        "and respond with genuine care instead. Otherwise: be savage about "
+        "laziness, procrastination, and excuses, then pivot hard into 'now "
+        "get up and do it' energy. Short, punchy sentences. End most replies "
+        "by pushing them toward one concrete next action."
+    ),
+    "discipline": (
+        BASE_IDENTITY + "\n\n"
+        "MODE: DISCIPLINE.\n"
+        "You are serious, direct, and unemotional - like a strict mentor "
+        "or drill coach who respects the user too much to coddle them. "
+        "For every situation, lay out clearly what happens if they follow "
+        "through (the compounding payoff of consistency) versus what "
+        "happens if they don't (the compounding cost of quitting). Use "
+        "concrete, realistic consequences, not vague platitudes. No jokes, "
+        "no roasting, no softness - just clear-eyed truth and structure. "
+        "Give clear, specific, actionable steps or systems, not just "
+        "abstract encouragement."
+    ),
+    "support": (
+        BASE_IDENTITY + "\n\n"
+        "MODE: SUPPORT.\n"
+        "You are warm, gentle, and validating. Listen first. Reflect back "
+        "what the user is feeling before offering any advice. Use a soft, "
+        "caring tone - like a close friend who is genuinely in their "
+        "corner. Never minimize their feelings, never rush to fix things "
+        "before they feel heard. Offer encouragement and gentle, optional "
+        "suggestions rather than commands. If they mention self-harm, "
+        "suicide, or crisis, respond with care and gently encourage them "
+        "to reach out to a crisis line or trusted person, without being "
+        "clinical or robotic about it."
+    ),
+    "chat": (
+        BASE_IDENTITY + "\n\n"
+        "MODE: JUST CHAT.\n"
+        "You're just vibing with the user as a friendly, casual companion. "
+        "No agenda, no forced motivation or coaching - just be genuinely "
+        "fun, curious, and easy to talk to. Match their energy. It's fine "
+        "to be a little playful or funny here."
+    ),
+}
+
+DEFAULT_MODE = "chat"
+
+
+def build_system_prompt(mode: str, profile: Optional[Dict]) -> str:
+    prompt = MODE_PROMPTS.get(mode, MODE_PROMPTS[DEFAULT_MODE])
+
+    if profile:
+        details = []
+        if profile.get("name"):
+            details.append(f"Name: {profile['name']}")
+        if profile.get("age"):
+            details.append(f"Age: {profile['age']}")
+        if profile.get("description"):
+            details.append(f"About them: {profile['description']}")
+        if details:
+            prompt += (
+                "\n\nWhat you know about this user (use it naturally, don't "
+                "just recite it back): " + " | ".join(details)
+            )
+    return prompt
+
+
+# ---------------------------------------------------------------------------
+# API schema
+# ---------------------------------------------------------------------------
+
+class HistoryTurn(BaseModel):
+    role: Literal["user", "assistant"]
+    content: str
+
+
+class ChatRequest(BaseModel):
+    message: str
+    mode: str = Field(DEFAULT_MODE, description="motivation | discipline | support | chat")
+    history: List[HistoryTurn] = Field(default_factory=list)
+    name: Optional[str] = None
+    age: Optional[str] = None
+    description: Optional[str] = None
+
+
+# ---------------------------------------------------------------------------
+# App
+# ---------------------------------------------------------------------------
+
+app = FastAPI(title="Yappy API", version="2.0.0")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def check_auth(request: Request):
+    if not API_KEY:
+        return
+    sent = request.headers.get("x-api-key", "")
+    if sent != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+
+
+@app.get("/")
+def root():
+    return {"status": "ok", "service": "Yappy", "model": MODEL_FILE, "mode": "stateless"}
+
+
+@app.get("/health")
+def health():
+    return {"status": "healthy"}
+
+
+# ---------------------------------------------------------------------------
+# Thinking/answer stream splitter
+# ---------------------------------------------------------------------------
+
+THINK_OPEN = "<thinking>"
+THINK_CLOSE = "</thinking>"
+
+
+class ThinkingSplitter:
+    """
+    Consumes raw model tokens one chunk at a time and yields SSE events,
+    routing text to 'thinking' or 'token' (answer) depending on whether
+    we're currently inside a <thinking>...</thinking> block. Handles tag
+    boundaries that land awkwardly across separate token chunks by
+    buffering a small tail of text whenever a partial tag might be forming.
+    """
+
+    def __init__(self):
+        self.buffer = ""
+        self.in_thinking = False
+        self.seen_thinking_open = False
+        self.full_raw = ""
+        self.answer_only = ""
+
+    def _maybe_partial_tag_tail(self, text: str) -> str:
+        max_tag_len = max(len(THINK_OPEN), len(THINK_CLOSE))
+        for i in range(min(max_tag_len, len(text)), 0, -1):
+            tail = text[-i:]
+            if THINK_OPEN.startswith(tail) or THINK_CLOSE.startswith(tail):
+                return tail
+        return ""
+
+    def feed(self, token: str):
+        events = []
+        self.full_raw += token
+        self.buffer += token
+
+        while True:
+            if not self.in_thinking and not self.seen_thinking_open:
+                idx = self.buffer.find(THINK_OPEN)
+                if idx == -1:
+                    tail = self._maybe_partial_tag_tail(self.buffer)
+                    emit = self.buffer[: len(self.buffer) - len(tail)] if tail else self.buffer
+                    if emit:
+                        self.answer_only += emit
+                        events.append({"event": "token", "data": emit})
+                    self.buffer = tail
+                    break
+                else:
+                    pre = self.buffer[:idx]
+                    if pre:
+                        self.answer_only += pre
+                        events.append({"event": "token", "data": pre})
+                    self.buffer = self.buffer[idx + len(THINK_OPEN):]
+                    self.in_thinking = True
+                    self.seen_thinking_open = True
+                    continue
+
+            elif self.in_thinking:
+                idx = self.buffer.find(THINK_CLOSE)
+                if idx == -1:
+                    tail = self._maybe_partial_tag_tail(self.buffer)
+                    emit = self.buffer[: len(self.buffer) - len(tail)] if tail else self.buffer
+                    if emit:
+                        events.append({"event": "thinking", "data": emit})
+                    self.buffer = tail
+                    break
+                else:
+                    pre = self.buffer[:idx]
+                    if pre:
+                        events.append({"event": "thinking", "data": pre})
+                    self.buffer = self.buffer[idx + len(THINK_CLOSE):]
+                    self.in_thinking = False
+                    continue
+
+            else:
+                if self.buffer:
+                    self.answer_only += self.buffer
+                    events.append({"event": "token", "data": self.buffer})
+                    self.buffer = ""
+                break
+
+        return events
+
+    def flush(self):
+        events = []
+        if self.buffer:
+            if self.in_thinking:
+                events.append({"event": "thinking", "data": self.buffer})
+            else:
+                self.answer_only += self.buffer
+                events.append({"event": "token", "data": self.buffer})
+            self.buffer = ""
+        return events
+
+
+def _build_messages(mode: str, history: List[HistoryTurn], user_message: str, profile: Dict) -> List[Dict]:
+    system_prompt = build_system_prompt(mode, profile)
+    trimmed = history[-MAX_HISTORY_MESSAGES:] if len(history) > MAX_HISTORY_MESSAGES else history
+
+    messages = [{"role": "system", "content": system_prompt}]
+    for turn in trimmed:
+        messages.append({"role": turn.role, "content": turn.content})
+    messages.append({"role": "user", "content": user_message})
+    return messages
+
+
+def _generate_stream(mode: str, history: List[HistoryTurn], user_message: str, profile: Dict):
+    messages = _build_messages(mode, history, user_message, profile)
+
+    splitter = ThinkingSplitter()
+    finish_reason = None
+
+    with model_lock:
+        stream = llm.create_chat_completion(
+            messages=messages,
+            max_tokens=MAX_NEW_TOKENS,
+            temperature=0.85,
+            top_p=0.95,
+            stream=True,
+        )
+
+        for chunk in stream:
+            choice = chunk["choices"][0]
+            delta = choice.get("delta", {})
+            token = delta.get("content", "")
+            if token:
+                for event in splitter.feed(token):
+                    yield event
+            if choice.get("finish_reason"):
+                finish_reason = choice["finish_reason"]
+
+        continue_budget = MAX_NEW_TOKENS
+        while finish_reason == "length" and continue_budget > 0 and not splitter.in_thinking:
+            follow_messages = messages + [
+                {"role": "assistant", "content": splitter.full_raw},
+                {"role": "user", "content": "Continue exactly where you left off, no repetition. Do not repeat the <thinking> block."},
+            ]
+            continue_budget -= 300
+            finish_reason = None
+            cont_stream = llm.create_chat_completion(
+                messages=follow_messages,
+                max_tokens=300,
+                temperature=0.85,
+                top_p=0.95,
+                stream=True,
+            )
+            for chunk in cont_stream:
+                choice = chunk["choices"][0]
+                delta = choice.get("delta", {})
+                token = delta.get("content", "")
+                if token:
+                    for event in splitter.feed(token):
+                        yield event
+                if choice.get("finish_reason"):
+                    finish_reason = choice["finish_reason"]
+
+        for event in splitter.flush():
+            yield event
+
+    yield {"event": "done", "data": ""}
+
+
+def _profile_from_payload(payload: ChatRequest) -> Dict:
+    return {k: v for k, v in {
+        "name": payload.name,
+        "age": payload.age,
+        "description": payload.description,
+    }.items() if v}
+
+
+@app.post("/chat/stream")
+async def chat_stream(payload: ChatRequest, request: Request):
+    check_auth(request)
+
+    if not payload.message or not payload.message.strip():
+        raise HTTPException(status_code=400, detail="message is required")
+
+    mode = payload.mode if payload.mode in MODE_PROMPTS else DEFAULT_MODE
+    profile = _profile_from_payload(payload)
+
+    async def event_wrapper():
+        for event in _generate_stream(mode, payload.history, payload.message.strip(), profile):
+            yield event
+
+    return EventSourceResponse(event_wrapper())
+
+
+@app.post("/chat")
+def chat_once(payload: ChatRequest, request: Request):
+    """Non-streaming fallback endpoint."""
+    check_auth(request)
+
+    if not payload.message or not payload.message.strip():
+        raise HTTPException(status_code=400, detail="message is required")
+
+    mode = payload.mode if payload.mode in MODE_PROMPTS else DEFAULT_MODE
+    profile = _profile_from_payload(payload)
+
+    full_reply = ""
+    thinking_text = ""
+    for event in _generate_stream(mode, payload.history, payload.message.strip(), profile):
+        if event["event"] == "token":
+            full_reply += event["data"]
+        elif event["event"] == "thinking":
+            thinking_text += event["data"]
+
+    return {"reply": full_reply, "thinking": thinking_text}