86 lines
3.2 KiB
Python
86 lines
3.2 KiB
Python
"""Tests feature RAG : authoring corpus (rag_upsert/dump/delete) + retrieval (mockés)."""
|
|
|
|
import base64
|
|
import json as _json
|
|
|
|
import nacl.pwhash
|
|
|
|
from kazeia_central.core.adb import Adb
|
|
from kazeia_central.core.store import Store
|
|
from kazeia_central.features.medical.rag import service as rsvc
|
|
|
|
_OPS = nacl.pwhash.argon2id.OPSLIMIT_MIN
|
|
_MEM = nacl.pwhash.argon2id.MEMLIMIT_MIN
|
|
NOW = 1_700_000_000_000
|
|
|
|
|
|
def _store(tmp_path):
|
|
s = Store(tmp_path / "c.db", kdf_ops=_OPS, kdf_mem=_MEM)
|
|
s.unlock("pw")
|
|
return s
|
|
|
|
|
|
class _RagAdb(Adb):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._docs = {}
|
|
self.upserts = []
|
|
self.deleted = []
|
|
|
|
def content_query(self, path, *, serial=None, where=None, expect_rows=False, _retry=True):
|
|
if path == "rag_status":
|
|
return [{"enabled": 1, "ready": 1, "model": "e5-small", "dim": 384,
|
|
"doc_count": len(self._docs), "chunk_count": len(self._docs), "pending": 0}]
|
|
if path == "rag":
|
|
return [{"source": s, "char_count": len(t), "chunk_count": 1, "updated_at": 1}
|
|
for s, t in self._docs.items()]
|
|
if path == "rag_query":
|
|
return [{"source": "doc1", "score": 0.91, "text": "chunk pertinent pour: " + (where or "")}]
|
|
return []
|
|
|
|
def content_call(self, method, *, arg_json=None, extras=None, serial=None):
|
|
if method == "rag_upsert_json":
|
|
self._docs[arg_json["source"]] = arg_json["text"]
|
|
self.upserts.append(arg_json)
|
|
return {"ok": "true"}
|
|
if method == "rag_dump_json":
|
|
src = arg_json["source"]
|
|
if src not in self._docs:
|
|
return {"ok": "false", "error": "not_found"}
|
|
payload = {"source": src, "text": self._docs[src]}
|
|
return {"ok": "true", "result_b64": base64.b64encode(_json.dumps(payload).encode()).decode()}
|
|
return {"ok": "true"}
|
|
|
|
def shell(self, command, *, serial=None, timeout=None):
|
|
if "content delete" in command and "/rag/" in command:
|
|
src = command.split("/rag/")[1].strip()
|
|
self.deleted.append(src)
|
|
self._docs.pop(src, None)
|
|
return ""
|
|
|
|
|
|
def test_rag_upsert_get_richtext(tmp_path):
|
|
s = _store(tmp_path)
|
|
adb = _RagAdb()
|
|
rsvc.upsert_doc(adb, s, "D1", "anxiete", "Anxiété : signes.\nGestion : respiration lente.", now=NOW)
|
|
assert adb.upserts[0]["source"] == "anxiete" and ":" in adb.upserts[0]["text"]
|
|
d = rsvc.get_doc(adb, "D1", "anxiete") # rag_dump → texte riche intact
|
|
assert d["text"].startswith("Anxiété") and ":" in d["text"] and "\n" in d["text"]
|
|
assert "anxiete" in {x["source"] for x in rsvc.list_docs(adb, "D1")}
|
|
|
|
|
|
def test_rag_delete(tmp_path):
|
|
s = _store(tmp_path)
|
|
adb = _RagAdb()
|
|
rsvc.upsert_doc(adb, s, "D1", "temp", "x", now=NOW)
|
|
rsvc.delete_doc(adb, s, "D1", "temp", now=NOW)
|
|
assert adb.deleted == ["temp"] and "temp" not in adb._docs
|
|
|
|
|
|
def test_rag_query_and_status():
|
|
adb = _RagAdb()
|
|
hits = rsvc.query(adb, "D1", "je me sens anxieux")
|
|
assert hits and hits[0]["source"] == "doc1" and hits[0]["score"] > 0.5
|
|
st = rsvc.status(adb, "D1")
|
|
assert st["ready"] is True and st["model"] == "e5-small"
|