118 lines
3.4 KiB
Python
118 lines
3.4 KiB
Python
# -*- 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}
|