deploy: acra.py

acra.py

@@ -0,0 +1,124 @@
+import os
+import google.generativeai as genai
+from db import supabase
+from classifier_inference import classify_query
+from typing import List
+
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+EMBED_MODEL = "models/text-embedding-004"
+GEN_MODEL   = "gemma-3-27b-it"
+DEPTH       = {0: 0, 1: 3, 2: 6, 3: 10}
+
+def embed_texts(texts):
+    return [genai.embed_content(model=EMBED_MODEL, content=t,
+            task_type="retrieval_document")["embedding"] for t in texts]
+
+def embed_query(q):
+    return genai.embed_content(model=EMBED_MODEL, content=q,
+           task_type="retrieval_query")["embedding"]
+
+def adaptive_chunk(text, max_tok=512):
+    paras = [p.strip() for p in text.split("\n\n") if p.strip()]
+    chunks, cur = [], ""
+    for p in paras:
+        if (len(cur.split()) + len(p.split())) / 0.75 < max_tok:
+            cur = (cur + "\n\n" + p).strip()
+        else:
+            if cur: chunks.append(cur)
+            cur = p
+    if cur: chunks.append(cur)
+    return chunks or [text]
+
+def decompose(query):
+    m = genai.GenerativeModel(GEN_MODEL)
+    r = m.generate_content(
+        f"Decompose into 2-4 simpler sub-queries. "
+        f"Return numbered list only.\n\nQuery: {query}")
+    lines = [l.strip().lstrip("1234567890.). ")
+             for l in r.text.strip().split("\n") if l.strip()]
+    return lines[:4] or [query]
+
+def compress(query, chunks):
+    m = genai.GenerativeModel(GEN_MODEL)
+    out = []
+    for c in chunks:
+        r = m.generate_content(
+            f"Extract only sentences relevant to the query. "
+            f"Return empty if none.\n\nQuery: {query}\nChunk: {c}")
+        if r.text.strip(): out.append(r.text.strip())
+    return out
+
+def vsearch(query, namespace, user_id, k):
+    return (supabase.rpc("match_documents", {
+        "query_embedding":  embed_query(query),
+        "match_count":      k,
+        "filter_namespace": namespace,
+        "filter_user_id":   user_id
+    }).execute().data or [])
+
+PROMPTS = {
+    1: "Answer using ONLY the context below. Be concise.\n\nContext:\n{ctx}\n\nQuestion: {q}\nAnswer:",
+    2: "Synthesize the context to answer. Think step by step.\n\nContext:\n{ctx}\n\nQuestion: {q}\nAnswer:",
+    3: "Use chain-of-thought to answer this complex question.\nAddress each aspect. Note any gaps.\n\nContext:\n{ctx}\n\nQuestion: {q}\nReasoning and answer:",
+}
+
+async def ingest_pipeline(texts, metadata, namespace, user_id):
+    chunks, meta = [], []
+    for i, t in enumerate(texts):
+        for j, c in enumerate(adaptive_chunk(t)):
+            chunks.append(c)
+            meta.append({**metadata[i], "source_index": i, "chunk_index": j})
+    rows = [{"content": c, "embedding": e, "metadata": m,
+             "namespace": namespace, "user_id": user_id}
+            for c, e, m in zip(chunks, embed_texts(chunks), meta)]
+    for i in range(0, len(rows), 50):
+        supabase.table("documents").insert(rows[i:i+50]).execute()
+    return len(chunks)
+
+async def query_pipeline(query, namespace, top_k, rerank, user_id):
+    cls   = classify_query(query)
+    level = cls["level"]
+    k     = DEPTH[level]
+    model = genai.GenerativeModel(GEN_MODEL)
+
+    if level == 0:
+        r = model.generate_content(
+            f"Answer concisely from your knowledge:\n\n{query}")
+        return {"answer": r.text.strip(), "sources": [], "complexity": cls}
+
+    hits = []
+    if level == 3:
+        seen = set()
+        for sq in decompose(query):
+            for h in vsearch(sq, namespace, user_id, 4):
+                if h["id"] not in seen:
+                    seen.add(h["id"]); hits.append(h)
+    else:
+        hits = vsearch(query, namespace, user_id, k)
+
+    if not hits:
+        return {"answer": "No relevant documents found. Ingest some first.",
+                "sources": [], "complexity": cls}
+
+    chunks = [h["content"] for h in hits]
+    if rerank and level >= 2:
+        chunks = [c for c in compress(query, chunks) if c.strip()]
+
+    ctx = "\n\n---\n\n".join(chunks[:k])
+    r   = model.generate_content(PROMPTS[level].format(ctx=ctx, q=query))
+
+    return {
+        "answer": r.text.strip(),
+        "sources": [{"content": h["content"][:200],
+                     "metadata": h.get("metadata", {}),
+                     "score": h.get("similarity", 0)}
+                    for h in hits[:len(chunks)]],
+        "complexity": cls
+    }
+
+async def run_acra_pipeline(mode, **kw):
+    if mode == "ingest":
+        return await ingest_pipeline(kw["texts"], kw["metadata"],
+                                     kw["namespace"], kw["user_id"])
+    return await query_pipeline(kw["query"], kw["namespace"],
+                                kw["top_k"], kw["rerank"], kw["user_id"])