transcription/scripts/smoke_rag.py

68 lines
2.3 KiB
Python
Raw Normal View History

"""Ручной smoke-test движка через CLI: ingest → query → exit."""
import asyncio
import sys
import tempfile
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
from src.rag.engine import get_or_create_engine # noqa: E402
SAMPLES = {
"plan.md": (
"# План 3-го этажа\n\n"
"План 3-го этажа жилого дома. Оси: А, Б, В, Г. Размеры между осями А и Б: 5400 мм.\n"
"Квартиры: 301, 302, 303.\n"
),
"auth.md": (
"# Авторизация\n\n"
"Авторизация работает через JWT-токены с TTL 24 часа.\n"
"Refresh-токен живёт 30 дней.\n"
),
"schedule.md": (
"# График работ\n\n"
"Строительство начинается 1 июня 2026. Срок сдачи — 30 ноября 2027.\n"
),
}
async def main() -> int:
with tempfile.TemporaryDirectory() as tmp:
coll = Path(tmp) / "demo" / "qmd_collections" / "demo"
coll.mkdir(parents=True)
eng = get_or_create_engine(coll)
eng.warmup()
# Ingest
for name, text in SAMPLES.items():
result = eng.index_text(text, source_path=name)
print(f" + {name}: chunks={result.chunks_indexed} vectors={result.vectors_indexed} skipped={result.skipped}")
# Status
print("\nStatus:", eng.status())
# Search BM25
print('\n--- BM25 "авторизация JWT" ---')
for h in eng.search("авторизация JWT"):
print(f" {h.score:.3f} {h.file_path}: {h.snippet(60)}")
# Search vector
print('\n--- Vector "как устроен вход в систему" ---')
for h in eng.vsearch("как устроен вход в систему"):
print(f" {h.score:.3f} {h.file_path}: {h.snippet(60)}")
# Hybrid query
print('\n--- Hybrid "когда сдача объекта" ---')
for h in eng.query("когда сдача объекта", limit=3, use_rerank=False):
print(f" {h.score:.3f} {h.source_annotation()}: {h.snippet(60)}")
eng.close()
print("\nAll good [OK]")
return 0
if __name__ == "__main__":
raise SystemExit(asyncio.run(main()))