"""Kith Local — extraction brain (MiniCPM, OpenBMB). Two backends behind one extract() seam: - "local" (default): MiniCPM5-1B GGUF (~688 MB) running in-process on CPU via llama.cpp. No API, no cloud — this is what earns the Off-the-Grid badge and makes "your conversations never leave your machine" literally true. - "hosted": OpenBMB's free MiniCPM4.1-8B (vLLM, OpenAI-compatible). Higher quality, but a cloud call. Kept as a fallback for the demo via KITH_BRAIN=hosted. Switch with the KITH_BRAIN env var. Callers only ever touch extract(). """ import json import os import re import requests KITH_BRAIN = os.environ.get("KITH_BRAIN", "hosted") # "gpu" | "local" | "hosted" # Backends: # "gpu" — MiniCPM4.1-8B run LOCALLY on the Space GPU (the real product: # best quality + private + fast). app.py injects the GPU generate # fn via set_gpu_chat(); selected automatically when CUDA exists. # "local" — MiniCPM5-1B GGUF on CPU via llama.cpp (private, but weak/slow). # "hosted" — OpenBMB's MiniCPM4.1-8B cloud endpoint (dev fallback on no-GPU # boxes; a cloud call, so NOT used as the shipped path). # app.py sets KITH_BRAIN = "gpu" when a GPU is present, else "hosted". _gpu_chat_fn = None # injected by app.py: (messages, temperature) -> str def set_gpu_chat(fn) -> None: """app.py hands us its @spaces.GPU generate fn so extract() can use it.""" global _gpu_chat_fn _gpu_chat_fn = fn # hosted backend (fallback) OPENBMB_BASE = os.environ.get("OPENBMB_BASE", "http://35.203.155.71:8001") OPENBMB_MODEL = os.environ.get("OPENBMB_MODEL", "MiniCPM4.1-8B") OPENBMB_API_KEY = os.environ.get("OPENBMB_API_KEY", "") # local backend (default) LOCAL_REPO = os.environ.get("KITH_LOCAL_REPO", "openbmb/MiniCPM5-1B-GGUF") LOCAL_FILE = os.environ.get("KITH_LOCAL_FILE", "MiniCPM5-1B-Q4_K_M.gguf") _llm = None # lazily-loaded llama.cpp model (module-level: load once, reuse) VALID_RELATIONSHIPS = {"work", "personal", "health", "other"} VALID_SIGNAL_TYPES = {"want", "life_event", "preference", "health", "date"} DATE_OR_NULL = {"type": ["string", "null"]} EXTRACTION_SCHEMA = { "type": "object", "properties": { "people": {"type": "array", "items": {"type": "object", "properties": { "name": {"type": "string"}, "relationship": {"type": "string", "enum": sorted(VALID_RELATIONSHIPS)}, "summary": {"type": "string"}, }, "required": ["name", "relationship", "summary"]}}, "signals": {"type": "array", "items": {"type": "object", "properties": { "person": {"type": "string"}, "type": {"type": "string", "enum": sorted(VALID_SIGNAL_TYPES)}, "text": {"type": "string"}, "event_on": DATE_OR_NULL, "quote": {"type": "string"}, }, "required": ["person", "type", "text", "event_on", "quote"]}}, "commitments": {"type": "array", "items": {"type": "object", "properties": { "person": {"type": "string"}, "owner": {"type": "string", "enum": ["me", "them"]}, "text": {"type": "string"}, "promised_on": DATE_OR_NULL, "due": DATE_OR_NULL, "quote": {"type": "string"}, }, "required": ["person", "owner", "text", "promised_on", "due", "quote"]}}, "decisions": {"type": "array", "items": {"type": "object", "properties": { "text": {"type": "string"}, "rationale": {"type": ["string", "null"]}, "people": {"type": "array", "items": {"type": "string"}}, "decided_on": DATE_OR_NULL, }, "required": ["text", "rationale", "decided_on"]}}, }, "required": ["people", "signals", "commitments", "decisions"], } SYSTEM_PROMPT = """You extract relationship memory from a conversation transcript. The transcript is a single-microphone recording with NO speaker labels. The note-taker is "me". Infer who is speaking from context. Return STRICT JSON only — no markdown, no commentary — exactly this shape: {"people":[{"name":"","relationship":"work|personal|health|other","summary":""}], "signals":[{"person":"","type":"want|life_event|preference|health|date","text":"","event_on":null,"quote":""}], "commitments":[{"person":"","owner":"me|them","text":"","promised_on":null,"due":null,"quote":""}], "decisions":[{"text":"","rationale":null,"people":[],"decided_on":null}]} Rules: - Extract ONLY facts grounded in the transcript. Never invent details. - The note-taker holds the recorder. When someone is greeted by name ("Hey Sara."), that named person is the conversation partner — "them" — and facts they share in first person about their own life belong to THEM, not the note-taker. - The note-taker ("me", "I") is NEVER listed in people. Only the others. - ALWAYS write "text" in the third person, about the named person — NEVER copy the speaker's "I"/"my". If Sara says "I'm getting married", the signal text is "Sara is getting married", never "I'm getting married". Start each signal's text with the person's name or "Her/His/Their". - people are individual humans with personal names. Never companies, teams, places or groups. - summary: 1-2 sentences of personal nuance (what matters to them, what's going on in their life) — NOT a job description. - signals: only things a thoughtful friend would act on (a birthday, something they want, a life event, a health situation). - quote: copy the SHORT verbatim span from the transcript (one sentence, max ~20 words) that this fact came from. It must appear word-for-word in the transcript. This is how the user verifies the memory. If you cannot find a verbatim span that directly states the fact, DROP the fact entirely — do not include it. - GROUNDING: the "text" must be directly stated by its "quote". Never generalize, soften, or guess beyond the quote. If someone says "I've been coughing a lot", do NOT write "healthy with no concerns" — write what was said. No optimistic or filler summaries. If the quote doesn't clearly support the text, drop the item. - commitments: a SPECIFIC personal favour one person promised another ("I'll send you the photos", "I'll drop it off Friday"). owner "me" = the note-taker promised it, "them" = the other person promised. Do NOT record someone's job, sales pitch, or professional service as a commitment (a financial advisor "helping you reach your goals" is his job, not a personal favour) — skip those. - DATES are strict. event_on / promised_on / due is filled ONLY for an explicit calendar moment: a named date ("the 14th"), a weekday ("next Friday"), or a clear offset ("in two weeks"), resolved against today's date (given with its weekday — check it carefully), formatted YYYY-MM-DD. - NEVER turn an age, a life-stage, or a vague future into a date. "retire at 65", "when the kids grow up", "in a few years", "someday" all have event_on = null. An age is NOT a date. - Any example conversation you are shown demonstrates the FORMAT only. Never copy facts from an example into your output — only facts from the final transcript. - Omit empty lists' content but keep all four keys. Output the JSON object and nothing else.""" FEW_SHOT_USER = """Today is Monday, 2026-03-02. Transcript: Hey Marcus. Oh my God, it's been forever. How's the new place? Honestly, the move nearly killed me. But we're in. I did the whole kitchen myself. You have to come see it. I will, I promise. I'll bring that sourdough starter you wanted. I'll drop it off next weekend. You better. And isn't your birthday coming up? The 14th? Don't remind me. 40 already. Oh, and how's your mom doing after the surgery? She's recovering slowly. Walking again, which is huge. We decided to move her in with us for the spring. It just makes more sense than the drive every weekend.""" FEW_SHOT_ASSISTANT = """{"people":[{"name":"Marcus","relationship":"personal","summary":"Just moved into a new place and renovated the kitchen himself; proud of the new home. His mother is slowly recovering from surgery and is walking again."}], "signals":[{"person":"Marcus","type":"life_event","text":"Just finished moving into a new place; renovated the kitchen himself","event_on":null,"quote":"I did the whole kitchen myself."},{"person":"Marcus","type":"want","text":"Wants the sourdough starter","event_on":null,"quote":"I'll bring that sourdough starter you wanted."},{"person":"Marcus","type":"date","text":"Marcus's 40th birthday","event_on":"2026-03-14","quote":"And isn't your birthday coming up? The 14th?"},{"person":"Marcus","type":"health","text":"His mom is slowly recovering from surgery, walking again","event_on":null,"quote":"She's recovering slowly. Walking again, which is huge."}], "commitments":[{"person":"Marcus","owner":"me","text":"Drop off the sourdough starter","promised_on":"2026-03-02","due":"2026-03-08","quote":"I'll drop it off next weekend."}], "decisions":[{"text":"Marcus's mom will move in with him for the spring","rationale":"Easier than driving to her every weekend","people":["Marcus"],"decided_on":"2026-03-02"}]}""" class ExtractionError(RuntimeError): pass def _get_local_llm(): """Load the MiniCPM5-1B GGUF once. Downloads ~688 MB on first call.""" global _llm if _llm is None: from llama_cpp import Llama _llm = Llama.from_pretrained( repo_id=LOCAL_REPO, filename=LOCAL_FILE, n_ctx=8192, n_threads=os.cpu_count(), verbose=False, ) return _llm def _chat_local(messages: list[dict], temperature: float) -> str: """Local llama.cpp with JSON-schema-constrained decoding (no cloud).""" llm = _get_local_llm() out = llm.create_chat_completion( messages=messages, temperature=temperature, max_tokens=2500, response_format={"type": "json_object", "schema": EXTRACTION_SCHEMA}, ) return out["choices"][0]["message"]["content"] def _chat_hosted(messages: list[dict], temperature: float) -> str: resp = requests.post( f"{OPENBMB_BASE}/v1/chat/completions", headers={"Authorization": f"Bearer {OPENBMB_API_KEY}", "Content-Type": "application/json"}, json={ "model": OPENBMB_MODEL, "messages": messages, "temperature": temperature, "max_tokens": 2500, "chat_template_kwargs": {"enable_thinking": False}, "response_format": { "type": "json_schema", "json_schema": {"name": "kith_extraction", "schema": EXTRACTION_SCHEMA}, }, }, timeout=120, ) resp.raise_for_status() return resp.json()["choices"][0]["message"]["content"] def _chat(messages: list[dict], temperature: float) -> str: """Dispatch to the configured backend. Same (messages, temp) -> text contract.""" if KITH_BRAIN == "gpu": if _gpu_chat_fn is None: raise ExtractionError("KITH_BRAIN=gpu but no GPU generate fn was injected (set_gpu_chat).") return _gpu_chat_fn(messages, temperature) if KITH_BRAIN == "hosted": return _chat_hosted(messages, temperature) return _chat_local(messages, temperature) def _parse_json(raw: str) -> dict: text = re.sub(r".*?", "", raw, flags=re.DOTALL) text = text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip() start, end = text.find("{"), text.rfind("}") if start == -1 or end <= start: raise ExtractionError(f"No JSON object in model output: {raw[:200]!r}") body = text[start : end + 1] try: return json.loads(body) except json.JSONDecodeError: # Common failure: a verbatim quote contains an inner " that wasn't # escaped. Re-escape stray quotes that sit mid-string (not adjacent to # JSON structural characters) and retry once. repaired = re.sub(r'(?<=[^\s:,{[])"(?=[^\s:,}\]])', r"\\\"", body) return json.loads(repaired) def _verify_quote(quote, transcript_norm: str) -> str | None: """Keep the quote only if it really appears in the transcript (grounding).""" if not quote: return None q = " ".join(str(quote).split()).strip() if len(q) < 4: return None return q if q.lower() in transcript_norm else None def _clean(extraction: dict, transcript: str = "") -> dict: """Coerce the model output into the schema; drop what doesn't fit.""" transcript_norm = " ".join(transcript.split()).lower() out = {"people": [], "signals": [], "commitments": [], "decisions": []} for p in extraction.get("people") or []: if not isinstance(p, dict) or not p.get("name"): continue rel = (p.get("relationship") or "other").lower() out["people"].append({ "name": str(p["name"]).strip(), "relationship": rel if rel in VALID_RELATIONSHIPS else "other", "summary": str(p.get("summary") or "").strip(), }) known = {p["name"].lower() for p in out["people"]} for s in extraction.get("signals") or []: if not isinstance(s, dict) or not s.get("text") or not s.get("person"): continue if str(s["person"]).lower() not in known: continue # Grounding rule: no verbatim quote => we can't trust it => drop it. quote = _verify_quote(s.get("quote"), transcript_norm) if not quote: continue stype = (s.get("type") or "preference").lower() out["signals"].append({ "person": str(s["person"]), "type": stype if stype in VALID_SIGNAL_TYPES else "preference", "text": str(s["text"]).strip(), "event_on": s.get("event_on") or None, "quote": quote, }) for c in extraction.get("commitments") or []: if not isinstance(c, dict) or not c.get("text") or not c.get("person"): continue if str(c["person"]).lower() not in known: continue quote = _verify_quote(c.get("quote"), transcript_norm) if not quote: continue # same grounding rule: no quote, no commitment out["commitments"].append({ "person": str(c["person"]), "owner": "them" if c.get("owner") == "them" else "me", "text": str(c["text"]).strip(), "promised_on": c.get("promised_on") or None, "due": c.get("due") or None, "quote": quote, }) for d in extraction.get("decisions") or []: if not isinstance(d, dict) or not d.get("text"): continue out["decisions"].append({ "text": str(d["text"]).strip(), "rationale": d.get("rationale") or None, "decided_on": d.get("decided_on") or None, }) return out def extract(transcript: str, today_date: str) -> dict: """Transcript -> structured relationship memory. Retries once on bad JSON.""" from datetime import date as _date weekday = _date.fromisoformat(today_date).strftime("%A") messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": FEW_SHOT_USER}, {"role": "assistant", "content": FEW_SHOT_ASSISTANT}, {"role": "user", "content": f"Today is {weekday}, {today_date}.\n\nTranscript:\n{transcript}"}, ] # The verbatim `quote` field can carry inner quotes/punctuation that # occasionally break a decode. Retry with rising temperature so a transient # bad roll gets a fresh attempt (temp 0 would just repeat the same break). last_err: Exception | None = None for temperature in (0.0, 0.4, 0.7): try: return _clean(_parse_json(_chat(messages, temperature=temperature)), transcript) except (ExtractionError, json.JSONDecodeError) as err: last_err = err raise ExtractionError(f"Extraction failed after retries: {last_err}")