v3.2: fp32 matrix at load (was recasting 186MB/query); add scope to split personal vs public layer

main.py

@@ -1,8 +1,13 @@
-"""OmniBrain API v3.1 — real routes over the exported brain.
+"""OmniBrain API v3.2 — real routes over the exported brain.
 
-Boot: pull parquet shards from the BRAIN_REPO dataset repo, load fp16 matrix
+Boot: pull parquet shards from the BRAIN_REPO dataset repo, load an fp32 matrix
 in RAM, serve brute-force cosine. Space disk is ephemeral, so the dataset repo
 is the durable store; restarting the Space just re-pulls.
+
+v3.2: (1) matrix kept fp32 at load — v3.1 re-cast the whole fp16 matrix to
+fp32 on EVERY query (~186MB alloc per request). (2) `scope` splits the personal
+layer from bulk public reference rows, which otherwise outnumber it ~7:1 and
+crowd real answers out of the top-k.
 """
 import os, glob
 from collections import Counter
@@ -17,6 +22,16 @@ from sentence_transformers import SentenceTransformer
 os.environ.setdefault("HF_HOME", "/tmp/hf")
 REPO_ID = os.environ.get("BRAIN_REPO", "Fonedo/omnibrain-caleb-brain")
 
+# Bulk public reference sets. Everything else (memories, vault, project docs,
+# agents, skills) is Caleb's own knowledge and is what search should hit first.
+PUBLIC_SOURCES = {
+    "teknium/OpenHermes-2.5", "openbmb/UltraFeedback",
+    "nvidia/OpenCodeReasoning", "nvidia/OpenCodeReasoning-2",
+    "ise-uiuc/Magicoder-OSS-Instruct-75K", "princeton-nlp/SWE-bench",
+    "zai-org/AgentInstruct", "markov-ai/gaming-500-hours",
+    "HuggingFaceFW/fineweb-edu",
+}
+
 app = FastAPI(title="OmniBrain")
 S = {"ready": False, "error": None}
 
@@ -29,39 +44,45 @@ def load():
         for f in sorted(glob.glob(f"{d}/*.parquet")):
             t = pq.read_table(f)
             vecs.append(np.asarray(t["vector"].combine_chunks().flatten(),
-                                   dtype=np.float16).reshape(-1, 768))
+                                   dtype=np.float32).reshape(-1, 768))
             meta.extend(zip(t["project"].to_pylist(), t["source"].to_pylist(),
                             t["type"].to_pylist(), t["text"].to_pylist()))
-        S["V"] = np.vstack(vecs) if vecs else np.zeros((0, 768), np.float16)
+        # fp32 once, not per query
+        S["V"] = np.vstack(vecs) if vecs else np.zeros((0, 768), np.float32)
         S["M"] = meta
+        S["IS_PUB"] = np.array([m[1] in PUBLIC_SOURCES for m in meta], dtype=bool)
         S["model"] = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5",
                                          trust_remote_code=True, device="cpu")
         S["model"].max_seq_length = 512
         S["ready"] = True
-        print(f"Loaded {len(meta):,} vectors from {REPO_ID}")
+        n_pub = int(S["IS_PUB"].sum())
+        print(f"Loaded {len(meta):,} vectors "
+              f"({len(meta)-n_pub:,} personal / {n_pub:,} public) from {REPO_ID}")
     except Exception as e:
         S["error"] = str(e)[:500]
         print("STARTUP FAILED:", S["error"])
 
 @app.get("/health")
 def health():
-    return {"status": "online", "version": "3.1", "ready": S["ready"],
+    return {"status": "online", "version": "3.2", "ready": S["ready"],
             "vectors": len(S.get("M", [])), "error": S["error"]}
 
 @app.get("/stats")
 def stats():
     c = Counter(m[0] for m in S.get("M", []))
     t = Counter(m[2] for m in S.get("M", []))
-    return {"vectors": len(S.get("M", [])), "by_project": dict(c),
-            "by_type": dict(t), "dim": 768,
-            "model": "nomic-ai/nomic-embed-text-v1.5",
-            "brain_repo": REPO_ID}
+    pub = int(S["IS_PUB"].sum()) if S.get("ready") else 0
+    return {"vectors": len(S.get("M", [])),
+            "personal": len(S.get("M", [])) - pub, "public": pub,
+            "by_project": dict(c), "by_type": dict(t), "dim": 768,
+            "model": "nomic-ai/nomic-embed-text-v1.5", "brain_repo": REPO_ID}
 
 class Q(BaseModel):
     query: str
     k: int = 8
     project: str | None = None
     agent: str | None = None          # kimi | claude | codex — biases type
+    scope: str = "personal"           # personal | all | public
 
 # each AI gets a different slice of the brain, per the OmniBrain design
 AGENT_BIAS = {
@@ -70,10 +91,19 @@ AGENT_BIAS = {
     "codex":  {"code": 1.2, "doc": 1.05},                                 # code
 }
 
-def _search(q: str, k: int, project: str | None, agent: str | None):
+def _search(q: str, k: int, project: str | None, agent: str | None,
+            scope: str = "personal"):
+    if not len(S["M"]):
+        return []
     # asymmetric model: docs were embedded with "search_document: "
     v = S["model"].encode(["search_query: " + q], normalize_embeddings=True)[0]
-    sims = S["V"].astype(np.float32) @ v.astype(np.float32)
+    sims = S["V"] @ v.astype(np.float32)
+
+    if scope == "personal":
+        sims = np.where(S["IS_PUB"], -1e9, sims)
+    elif scope == "public":
+        sims = np.where(S["IS_PUB"], sims, -1e9)
+
     if project:
         mask = np.array([m[0] == project or m[0] == "meta" for m in S["M"]])
         sims = np.where(mask, sims, -1e9)
@@ -81,7 +111,9 @@ def _search(q: str, k: int, project: str | None, agent: str | None):
         b = AGENT_BIAS[agent]
         boost = np.array([b.get(m[2], 1.0) for m in S["M"]], dtype=np.float32)
         sims = np.where(sims > -1e8, sims * boost, sims)
-    idx = np.argpartition(-sims, min(k, len(sims) - 1))[:k]
+
+    k = max(1, min(k, len(sims)))
+    idx = np.argpartition(-sims, k - 1)[:k]
     idx = idx[np.argsort(-sims[idx])]
     return [{"score": round(float(sims[i]), 4), "project": S["M"][i][0],
              "source": S["M"][i][1], "type": S["M"][i][2],
@@ -90,12 +122,12 @@ def _search(q: str, k: int, project: str | None, agent: str | None):
 @app.post("/search")
 def search(q: Q):
     if not S["ready"]: return {"error": S["error"] or "still loading"}
-    return {"results": _search(q.query, q.k, q.project, q.agent)}
+    return {"results": _search(q.query, q.k, q.project, q.agent, q.scope)}
 
 @app.post("/generate_context")
 def generate_context(q: Q):
     if not S["ready"]: return {"error": S["error"] or "still loading"}
-    hits = _search(q.query, q.k, q.project, q.agent)
+    hits = _search(q.query, q.k, q.project, q.agent, q.scope)
     ctx = "\n\n---\n".join(
         f"[{h['project']}/{h['type']}] {h['text']}" for h in hits)
     return {"context": ctx, "n": len(hits)}
@@ -103,4 +135,5 @@ def generate_context(q: Q):
 @app.get("/")
 def root():
     return {"message": "OmniBrain API",
-            "routes": ["/health", "/stats", "/search", "/generate_context"]}
+            "routes": ["/health", "/stats", "/search", "/generate_context"],
+            "scopes": ["personal (default)", "all", "public"]}