Advanced mode: chat-driven form filling (fill_form tool), 10 official WA forms auto-grabbed + auto-mapped, doc vault, web search, per-user doc RAG

userdocs.py

@@ -0,0 +1,117 @@
+# -*- coding: utf-8 -*-
+"""Per-user document index — ingest uploaded files into a vector array so the
+attorney can reference the user's actual documents alongside the case law.
+Reuses rag.py's embedder + chunker + cosine. stdlib + optional pymupdf (PDF).
+"""
+import json
+import os
+import re
+
+from rag import embed, _split_long, _cosine, MAX_CHUNK
+
+INDEX_DIR = "/opt/astraea/indexes"
+
+
+def _path(uid):
+    return os.path.join(INDEX_DIR, f"user_{uid}.json")
+
+
+def _docx_text(path):
+    import zipfile
+    try:
+        with zipfile.ZipFile(path) as z:
+            xml = z.read("word/document.xml").decode("utf-8", "ignore")
+        out = []
+        for p in re.split(r"</w:p>", xml):
+            ts = re.findall(r"<w:t[^>]*>(.*?)</w:t>", p)
+            if ts:
+                out.append("".join(ts))
+        return "\n".join(out)
+    except Exception:
+        return ""
+
+
+def _pdf_text(path):
+    try:
+        import fitz  # pymupdf
+        doc = fitz.open(path)
+        parts = []
+        for page in doc:
+            parts.append(page.get_text())
+        doc.close()
+        return "\n".join(parts)
+    except Exception:
+        return ""
+
+
+def extract_text(orig_name, path):
+    ext = os.path.splitext(orig_name)[1].lower()
+    if ext == ".docx":
+        return _docx_text(path)
+    if ext == ".pdf":
+        return _pdf_text(path)
+    if ext in (".txt", ".md", ".csv", ".log", ".json", ".html", ".xml", ".rtf"):
+        try:
+            with open(path, "r", encoding="utf-8", errors="replace") as f:
+                return f.read()
+        except Exception:
+            return ""
+    return ""
+
+
+def build(uid, files, upload_dir):
+    """Embed every file's text into a per-user vector index. Returns chunk count."""
+    chunks = []
+    for rec in files:
+        path = os.path.join(upload_dir, rec["filename"])
+        if not os.path.exists(path):
+            continue
+        text = extract_text(rec.get("orig_name", rec["filename"]), path)
+        if not text or not text.strip():
+            continue
+        for piece in _split_long(text, MAX_CHUNK):
+            chunks.append({"source": rec.get("orig_name", rec["filename"]), "text": piece, "vector": None})
+    for c in chunks:
+        try:
+            c["vector"] = embed(c["text"])
+        except Exception:
+            c["vector"] = None
+    chunks = [c for c in chunks if c["vector"]]
+    os.makedirs(INDEX_DIR, exist_ok=True)
+    try:
+        with open(_path(uid), "w") as f:
+            json.dump({"chunks": chunks}, f)
+    except Exception:
+        pass
+    return len(chunks)
+
+
+def retrieve(uid, query, top_k=4):
+    if not os.path.exists(_path(uid)):
+        return []
+    try:
+        with open(_path(uid)) as f:
+            data = json.load(f)
+    except Exception:
+        return []
+    qvec = embed(query)
+    scored = []
+    for c in data.get("chunks", []):
+        v = c.get("vector")
+        if not v:
+            continue
+        scored.append((_cosine(qvec, v), c))
+    scored.sort(key=lambda x: x[0], reverse=True)
+    return [{"source": c["source"], "text": c["text"], "score": round(s, 3)}
+            for s, c in scored[:top_k] if s > 0.15]
+
+
+def status(uid):
+    if not os.path.exists(_path(uid)):
+        return {"indexed": False, "chunks": 0}
+    try:
+        with open(_path(uid)) as f:
+            data = json.load(f)
+        return {"indexed": True, "chunks": len(data.get("chunks", []))}
+    except Exception:
+        return {"indexed": False, "chunks": 0}