| """core/knowledge.py — it should remember, not re-ask. |
| |
| Phase 6. Two things live in one store: |
| |
| * **turns** — every conversation turn, so "what did I say yesterday" is a |
| query and not a scroll-back; |
| * **chunks** — documents the user (or JARVIS) chose to keep, embedded for |
| semantic recall. |
| |
| Embeddings are Gemini's, stored as float32 blobs in plain SQLite, and searched |
| with a cosine over numpy. PLAN.md originally said sqlite-vec; that is a Rust |
| extension with no wheel for this platform, and adding a compiler to the Space |
| image to save 20ms of search on a few thousand chunks is the wrong trade. The |
| search is O(n·d) in numpy and stays instant to ~50k chunks. |
| |
| The important design choice is the **lexical fallback**: if there is no API key, |
| or the embedding call fails, search still works on token overlap. Recall that |
| degrades honestly is far better than a vector store that returns nothing and |
| looks broken. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| import re |
| import sqlite3 |
| import threading |
| import time |
| import urllib.request |
| from pathlib import Path |
| from typing import Any, Optional |
|
|
| from core.data_paths import data_root |
|
|
| CHUNK_CHARS = 900 |
| CHUNK_OVERLAP = 120 |
| |
| |
| EMBED_MODEL = "gemini-embedding-001" |
| DIM = 768 |
|
|
| _lock = threading.RLock() |
| _db: Optional[sqlite3.Connection] = None |
|
|
|
|
| |
|
|
| def _path() -> Path: |
| return data_root() / "knowledge.sqlite" |
|
|
|
|
| def db() -> sqlite3.Connection: |
| global _db |
| with _lock: |
| if _db is not None: |
| return _db |
| p = _path() |
| p.parent.mkdir(parents=True, exist_ok=True) |
| |
| |
| |
| |
| conn = sqlite3.connect(str(p), timeout=15, check_same_thread=False) |
| conn.row_factory = sqlite3.Row |
| conn.execute("PRAGMA journal_mode=WAL") |
| conn.execute("""CREATE TABLE IF NOT EXISTS docs ( |
| id INTEGER PRIMARY KEY, title TEXT, source TEXT, uri TEXT, |
| created REAL, meta TEXT, chunks INTEGER DEFAULT 0)""") |
| conn.execute("""CREATE TABLE IF NOT EXISTS chunks ( |
| id INTEGER PRIMARY KEY, doc_id INTEGER, ord INTEGER, text TEXT, |
| vec BLOB, dim INTEGER, ts REAL)""") |
| conn.execute("""CREATE TABLE IF NOT EXISTS turns ( |
| id INTEGER PRIMARY KEY, ts REAL, role TEXT, text TEXT, session TEXT)""") |
| conn.execute("CREATE INDEX IF NOT EXISTS chunks_doc ON chunks(doc_id)") |
| conn.execute("CREATE INDEX IF NOT EXISTS turns_ts ON turns(ts)") |
| conn.commit() |
| _db = conn |
| return _db |
|
|
|
|
| def reset() -> None: |
| """Test/deploy hook.""" |
| global _db |
| with _lock: |
| if _db is not None: |
| try: |
| _db.close() |
| except Exception: |
| pass |
| _db = None |
|
|
|
|
| |
|
|
| def _api_key() -> str: |
| try: |
| from core import gemini |
| return (gemini.api_key() or "").strip() |
| except Exception: |
| return "" |
|
|
|
|
| def embed(texts: list[str], *, timeout: float = 25.0) -> Optional[list[list[float]]]: |
| """Gemini embeddings, or None. Never raises: callers fall back to lexical.""" |
| if not texts: |
| return [] |
| key = _api_key() |
| if not key: |
| return None |
| try: |
| from google import genai |
| from google.genai import types as gtypes |
| client = genai.Client(api_key=key) |
| out: list[list[float]] = [] |
| |
| for i in range(0, len(texts), 20): |
| batch = [t[:6000] for t in texts[i:i + 20]] |
| resp = client.models.embed_content( |
| model=EMBED_MODEL, contents=batch, |
| config=gtypes.EmbedContentConfig(output_dimensionality=DIM)) |
| for e in (resp.embeddings or []): |
| out.append(list((e.values or []))) |
| return out or None |
| except Exception: |
| return None |
|
|
|
|
| def _pack(vec: list[float]) -> bytes: |
| import struct |
| return struct.pack(f"<{len(vec)}f", *[float(v) for v in vec]) |
|
|
|
|
| def _unpack(blob: bytes) -> list[float]: |
| import struct |
| n = len(blob) // 4 |
| return list(struct.unpack(f"<{n}f", blob)) |
|
|
|
|
| |
|
|
| _WORD = re.compile(r"[a-z0-9_]{2,}") |
|
|
| _STOP = {"the", "and", "for", "you", "that", "this", "with", "was", "are", "but", |
| "not", "have", "has", "had", "from", "they", "them", "his", "her", |
| "she", "him", "its", "our", "your", "what", "when", "where", "who", |
| "how", "why", "can", "could", "would", "should", "will", "did", "does", |
| "been", "there", "here", "about", "just", "like", "get", "got"} |
|
|
|
|
| def _tokens(text: str) -> list[str]: |
| return [w for w in _WORD.findall((text or "").lower()) if w not in _STOP] |
|
|
|
|
| def _lexical_score(query: str, text: str) -> float: |
| """Token overlap with a small idf-ish twist: rarer words count more.""" |
| q = set(_tokens(query)) |
| if not q: |
| return 0.0 |
| t = _tokens(text) |
| if not t: |
| return 0.0 |
| overlap = sum(1 for w in set(t) if w in q) |
| if not overlap: |
| return 0.0 |
| |
| density = overlap / (len(set(t)) ** 0.5) |
| return overlap * 0.6 + density |
|
|
|
|
| |
|
|
| def record_turn(role: str, text: str, *, session: str = "") -> bool: |
| """One conversation turn. Cheap enough to always call.""" |
| text = str(text or "").strip() |
| if not text: |
| return False |
| try: |
| with _lock: |
| db().execute( |
| "INSERT INTO turns (ts, role, text, session) VALUES (?,?,?,?)", |
| (time.time(), str(role or "?"), text[:4000], str(session or ""))) |
| db().commit() |
| return True |
| except Exception: |
| return False |
|
|
|
|
| def recent_turns(limit: int = 20, *, role: str = "", query: str = "") -> list[dict]: |
| with _lock: |
| c = db() |
| if query: |
| rows = c.execute( |
| "SELECT * FROM turns WHERE text LIKE ? ORDER BY ts DESC LIMIT ?", |
| (f"%{query}%", max(1, min(limit, 500)))).fetchall() |
| else: |
| rows = c.execute( |
| "SELECT * FROM turns WHERE (? = '' OR role = ?) " |
| "ORDER BY ts DESC LIMIT ?", |
| (role, role, max(1, min(limit, 500)))).fetchall() |
| return [{"ts": r["ts"], "role": r["role"], "text": r["text"], |
| "session": r["session"]} for r in rows] |
|
|
|
|
| def turn_count() -> int: |
| with _lock: |
| return int(db().execute("SELECT COUNT(*) c FROM turns").fetchone()["c"]) |
|
|
|
|
| |
|
|
| def _chunk_text(text: str) -> list[str]: |
| """Split on paragraph boundaries, then hard-wrap anything oversized.""" |
| text = re.sub(r"\n{3,}", "\n\n", (text or "").strip()) |
| if not text: |
| return [] |
| paras = [p.strip() for p in text.split("\n\n") if p.strip()] |
| out: list[str] = [] |
| buf = "" |
| for para in paras: |
| while len(para) > CHUNK_CHARS: |
| head, para = para[:CHUNK_CHARS], para[CHUNK_CHARS - CHUNK_OVERLAP:] |
| out.append(head) |
| if len(buf) + len(para) + 2 <= CHUNK_CHARS: |
| buf = f"{buf}\n\n{para}" if buf else para |
| else: |
| if buf: |
| out.append(buf) |
| buf = para |
| if buf: |
| out.append(buf) |
| return [c for c in out if len(c) > 20] |
|
|
|
|
| def ingest(title: str, text: str, *, source: str = "note", uri: str = "", |
| meta: Optional[dict] = None) -> dict: |
| """Store a document and make it recallable. Returns the doc record.""" |
| parts = _chunk_text(text) |
| if not parts: |
| raise ValueError("nothing worth storing — that text is empty or too short") |
| vecs = embed(parts) |
| with _lock: |
| c = db() |
| cur = c.execute( |
| "INSERT INTO docs (title, source, uri, created, meta, chunks) " |
| "VALUES (?,?,?,?,?,?)", |
| (str(title or "untitled")[:160], str(source or "note")[:40], |
| str(uri or "")[:300], time.time(), |
| json.dumps(meta or {}, ensure_ascii=False), len(parts))) |
| doc_id = int(cur.lastrowid) |
| now = time.time() |
| for i, part in enumerate(parts): |
| vec = vecs[i] if vecs and i < len(vecs) else None |
| c.execute( |
| "INSERT INTO chunks (doc_id, ord, text, vec, dim, ts) " |
| "VALUES (?,?,?,?,?,?)", |
| (doc_id, i, part, |
| _pack(vec) if vec else None, len(vec) if vec else 0, now)) |
| c.commit() |
| return {"id": doc_id, "title": str(title or "untitled")[:160], |
| "source": str(source or "note")[:40], "uri": str(uri or "")[:300], |
| "chunks": len(parts), "embedded": bool(vecs)} |
|
|
|
|
| def fetch_url(url: str, *, limit: int = 40000) -> tuple[str, str]: |
| """Title + readable text from a URL. Deliberately plain: strip tags, keep |
| paragraph breaks, no readability library to keep in the image.""" |
| req = urllib.request.Request(url, headers={"User-Agent": "Jaarvis/1.0"}) |
| with urllib.request.urlopen(req, timeout=25) as r: |
| raw = r.read().decode("utf-8", "replace") |
| ctype = r.headers.get("content-type", "") |
| m = re.search(r"<title[^>]*>(.*?)</title>", raw, re.I | re.S) |
| title = re.sub(r"\s+", " ", m.group(1)).strip() if m else (url.split("/")[-1] or url) |
| if "html" not in ctype.lower() and "html" not in raw[:200].lower(): |
| return title, raw[:limit] |
| raw = re.sub(r"(?is)<(script|style|noscript|svg|head)[^>]*>.*?</\1>", " ", raw) |
| raw = re.sub(r"(?i)<br\s*/?>|</p>|</div>|</li>|</h[1-6]>", "\n", raw) |
| text = re.sub(r"<[^>]+>", " ", raw) |
| text = (text.replace(" ", " ").replace("&", "&") |
| .replace("<", "<").replace(">", ">").replace(""", '"')) |
| text = re.sub(r"[ \t]+", " ", text) |
| text = re.sub(r"\n{3,}", "\n\n", text) |
| return title, text.strip()[:limit] |
|
|
|
|
| def forget(doc_id: int) -> bool: |
| with _lock: |
| c = db() |
| row = c.execute("SELECT id FROM docs WHERE id = ?", (int(doc_id),)).fetchone() |
| if not row: |
| return False |
| c.execute("DELETE FROM chunks WHERE doc_id = ?", (int(doc_id),)) |
| c.execute("DELETE FROM docs WHERE id = ?", (int(doc_id),)) |
| c.commit() |
| return True |
|
|
|
|
| def docs(limit: int = 100) -> list[dict]: |
| with _lock: |
| rows = db().execute( |
| "SELECT * FROM docs ORDER BY created DESC LIMIT ?", |
| (max(1, min(int(limit or 100), 500)),)).fetchall() |
| return [{"id": r["id"], "title": r["title"], "source": r["source"], |
| "uri": r["uri"], "created": r["created"], "chunks": r["chunks"], |
| "meta": json.loads(r["meta"] or "{}")} for r in rows] |
|
|
|
|
| |
|
|
| def search(query: str, *, k: int = 5, include_turns: bool = True) -> list[dict]: |
| """Semantic when embeddings exist, lexical otherwise. Always returns |
| something if any text overlaps, and says which mode it used.""" |
| query = str(query or "").strip() |
| if not query: |
| return [] |
| k = max(1, min(int(k or 5), 25)) |
| qv = embed([query]) |
| use_vectors = bool(qv and qv[0]) |
| scored: list[tuple[float, dict]] = [] |
|
|
| with _lock: |
| c = db() |
| if use_vectors: |
| rows = c.execute( |
| "SELECT c.id, c.doc_id, c.text, c.vec, d.title, d.source, d.uri " |
| "FROM chunks c LEFT JOIN docs d ON d.id = c.doc_id").fetchall() |
| import numpy as np |
| qa = np.asarray(qv[0], dtype="float32") |
| qn = float(np.linalg.norm(qa)) or 1.0 |
| for r in rows: |
| if not r["vec"]: |
| continue |
| va = np.asarray(_unpack(r["vec"]), dtype="float32") |
| if va.shape[0] != qa.shape[0]: |
| continue |
| vn = float(np.linalg.norm(va)) or 1.0 |
| score = float(np.dot(qa, va) / (qn * vn)) |
| scored.append((score, _row(r))) |
| |
| |
| if include_turns: |
| for r in c.execute( |
| "SELECT id, ts, role, text FROM turns " |
| "ORDER BY ts DESC LIMIT 2000").fetchall(): |
| s = _lexical_score(query, r["text"]) |
| if s > 0: |
| scored.append((s * 0.8, { |
| "kind": "turn", "text": r["text"][:600], |
| "title": f"{r['role']} said this", |
| "source": "conversation", "uri": "", |
| "ts": r["ts"], "doc_id": None})) |
|
|
| if not use_vectors: |
| with _lock: |
| for r in db().execute( |
| "SELECT c.id, c.doc_id, c.text, d.title, d.source, d.uri " |
| "FROM chunks c LEFT JOIN docs d ON d.id = c.doc_id").fetchall(): |
| s = _lexical_score(query, r["text"]) |
| if s > 0: |
| scored.append((s * 0.7, _row(r))) |
|
|
| scored.sort(key=lambda x: -x[0]) |
| out = [] |
| for score, row in scored[:k]: |
| row["score"] = round(float(score), 4) |
| out.append(row) |
| return out |
|
|
|
|
| def _row(r) -> dict: |
| return {"kind": "doc", "text": r["text"][:800], "title": r["title"] or "untitled", |
| "source": r["source"] or "", "uri": r["uri"] or "", "doc_id": r["doc_id"], |
| "ts": None} |
|
|
|
|
| def recall_block(query: str, *, k: int = 4, budget: int = 1200) -> str: |
| """A prompt-sized block of what we already know. Returns '' when we don't, |
| so the model is never handed an empty 'here is what you remember' header.""" |
| hits = search(query, k=k) |
| if not hits: |
| return "" |
| lines, used = [], 0 |
| for h in hits: |
| snippet = " ".join(str(h.get("text") or "").split()) |
| room = budget - used |
| if room <= 60: |
| break |
| snippet = snippet[:room] |
| used += len(snippet) + 40 |
| where = h.get("title") or h.get("source") or "memory" |
| lines.append(f"- [{where}] {snippet}") |
| if not lines: |
| return "" |
| return "Relevant things you already know:\n" + "\n".join(lines) |
|
|
|
|
| def stats() -> dict: |
| with _lock: |
| c = db() |
| d = int(c.execute("SELECT COUNT(*) c FROM docs").fetchone()["c"]) |
| ch = int(c.execute("SELECT COUNT(*) c FROM chunks").fetchone()["c"]) |
| vecs = int(c.execute( |
| "SELECT COUNT(*) c FROM chunks WHERE vec IS NOT NULL").fetchone()["c"]) |
| t = int(c.execute("SELECT COUNT(*) c FROM turns").fetchone()["c"]) |
| newest = c.execute("SELECT MAX(ts) t FROM turns").fetchone()["t"] |
| return {"docs": d, "chunks": ch, "embedded_chunks": vecs, "turns": t, |
| "db": str(_path()), "bytes": _path().stat().st_size |
| if _path().exists() else 0, |
| "newest_turn": newest, "embed_model": EMBED_MODEL, |
| "mode": "semantic" if vecs else "lexical", |
| "api_key": bool(_api_key())} |
|
|
|
|
| def describe() -> str: |
| s = stats() |
| return (f"knowledge: {s['docs']} docs / {s['chunks']} chunks " |
| f"({s['embedded_chunks']} embedded, {s['mode']}) · {s['turns']} turns") |
|
|
|
|
| |
|
|
| _pending: dict[str, str] = {} |
|
|
|
|
| def note_query(text: str) -> None: |
| """Called with the user's last words so the NEXT system prompt can carry |
| what we already know about the thing they are talking about. |
| |
| Best-effort and off the hot path: an embedding call per turn is a latency |
| risk, so a failure just means the prompt carries the index instead. |
| """ |
| text = str(text or "").strip() |
| if len(text) < 8: |
| return |
| try: |
| block = recall_block(text, k=3, budget=700) |
| _pending["recall"] = block |
| except Exception: |
| pass |
|
|
|
|
| def index_line(limit: int = 12) -> str: |
| """Titles only. Listing documents is cheap; pasting their contents into every |
| prompt is not, and the model can search when it needs the detail.""" |
| try: |
| ds = docs(limit) |
| except Exception: |
| return "" |
| if not ds: |
| return "" |
| names = ", ".join(f"{d['title']}" + (f" ({d['source']})" if d.get("source") else "") |
| for d in ds[:limit]) |
| more = f" and {len(ds) - limit} more" if len(ds) > limit else "" |
| return f"You have {len(ds)} stored document(s): {names}{more}. " \ |
| f"Use the knowledge tool to read the detail of any of them." |
|
|
|
|
| def prompt_block(budget: int = 900) -> str: |
| """The knowledge section of the system prompt. Empty when there is nothing, |
| so the model is never handed an empty 'here is what you remember'.""" |
| bits = [] |
| try: |
| rec = _pending.pop("recall", "") or "" |
| except Exception: |
| rec = "" |
| if rec: |
| bits.append(rec) |
| try: |
| idx = index_line() |
| except Exception: |
| idx = "" |
| if idx: |
| bits.append(idx) |
| if not bits: |
| return "" |
| return "[WHAT YOU KNOW]\n" + "\n".join(bits) + "\n" |
|
|
