simbiot-harness/harness/rag_lookup.py

"""Семантический поиск по базе кооператива для тиков харнеса (t-260929212316).

Вендор: /root/rag (fastembed bge-small-en-v1.5 + sqlite), базы:
- rag.db — /root/coop/context + agent-skills + PORTAL_API.md (12.5k чанков)
- hermes.db — 54k чанков общего корпуса (продукты, код, история)
Каждый запрос — отдельный сабпроцесс /root/rag-venv/bin/python rag.py query:
модель embedding (ONNX, ~130МБ) тяжёлая для импорта в процесс воркера,
изолированный сабпроцесс не раздувает память LLM-тика и не ломает
глобальный рантайм Strands.

Лимиты (по постановке Алекса): 1 запрос на тик, k<=5, таймаут 20с,
RAG_DISABLED=1 выключает. Трейс запроса пишется в event store
(RagQueried) — часовой уровень видит в событиях, что мин-пульс
ходил в знания кооператива.
"""
import json
import os
import subprocess

_HARNESS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
from . import paths as _paths  # t-261004113351-2858c1: RAG-индекс инстанса
_RAG_VENV_PY = "/root/rag-venv/bin/python"
_RAG_SCRIPT = "/root/rag/rag.py"

# База по умолчанию: корпус кооператива (12.5k чанков) — то, что просил Алекс
# для мин-пульса; hermes.db (54k) — опцией RAG_DB для верхних уровней.
# SIMBIOT_RAG_DB — точечное переопределение; SIMBIOT_HOME → <home>/rag/rag.db
# (t-261004113351: инстанс Юры не должен искать по нашим знаниям).
DEFAULT_DB = _paths.rag_db()

# Лимиты на тик: 1 запрос, k<=5, таймаут 20с (эмбеддинг ~2.7с + запас)
MAX_K = 5
TIMEOUT = 20


def _parse_output(out: str) -> list:
    """Разбор вывода rag.py: чанки, разделённые '═══ score | file#N ═══'."""
    chunks, cur = [], None
    for line in out.splitlines():
        if line.startswith("═══"):
            if cur is not None:
                chunks.append(cur)
            head = line.strip("═ ").strip()
            try:
                score, src = head.split("|", 1)
                score = float(score.strip())
            except ValueError:
                score, src = 0.0, head
            cur = {"score": score, "source": src.strip(), "text": ""}
        elif cur is not None:
            cur["text"] += line + "\n"
    if cur is not None:
        chunks.append(cur)
    for c in chunks:
        c["text"] = c["text"].strip()
    return chunks


def _trace(level, trace_id, payload):
    """Трейс RAG-запроса в event store — spend/observability едины с LLM-тиками."""
    if not (level and trace_id):
        return
    try:
        from . import es
        seq = 1
        rows = []  # считаем версию по aggregate_id — append_event сам валиден
        es.append_event("trace", f"{trace_id}:rag{seq}", "RagQueried", {
            **payload, "level": level,
        })
    except Exception:
        pass  # трейс — best effort, не валить тик


def rag_query(question: str, k: int = MAX_K, level: str = None,
              trace_id: str = None, db: str = None, timeout: int = TIMEOUT) -> list:
    """Семантический поиск по базе кооператива. Возвращает список чанков.

    [{score, source, text}] — отсортировано по релевантности. Пустой список —
    ничего не найдено или вендор упал; текст ошибки не поднимается наружу:
    мин-тик не должен падать из-за RAG.
    """
    if os.environ.get("RAG_DISABLED"):
        return []
    question = (question or "").strip()
    if not question:
        return []
    k = max(1, min(int(k or MAX_K), MAX_K))
    env = dict(os.environ)
    env["RAG_DB"] = db or DEFAULT_DB
    if not os.access(_RAG_VENV_PY, os.X_OK) or not os.path.exists(_RAG_SCRIPT):
        return []
    try:
        proc = subprocess.run(
            [_RAG_VENV_PY, _RAG_SCRIPT, "query", question, "-k", str(k)],
            capture_output=True, text=True, timeout=timeout or TIMEOUT,
            cwd=_HARNESS_DIR, env=env,
        )
        chunks = _parse_output(proc.stdout)
    except subprocess.TimeoutExpired:
        _trace(level, trace_id, {"query": question[:300], "k": k,
                                "error": f"timeout after {timeout}s"})
        return []
    except Exception as e:
        _trace(level, trace_id, {"query": question[:300], "k": k,
                                "error": str(e)[:200]})
        return []
    _trace(level, trace_id, {"query": question[:300], "k": k,
                             "results": len(chunks),
                             "top": [c["source"] for c in chunks[:3]]})
    return chunks


def format_chunks(chunks: list, max_chars: int = 6000) -> str:
    """Чанки → компактный текст для контекста тика."""
    out = []
    used = 0
    for c in chunks:
        body = c["text"]
        if not body:
            continue
        line = f"═══ {c['score']:.3f} | {c['source']} ═══\n{body}"
        if used + len(line) > max_chars:
            break
        out.append(line)
        used += len(line)
    return "\n\n".join(out)