gart/rag/search.py

"""Busca semântica RAG do projeto G-ART.

Consulta `gart.code_chunks` no Postgres (via túnel SSH) combinando dois
sinais e devolvendo faixas de linha, para o agente ler só o trecho relevante
em vez do arquivo inteiro.

Uso:
    rag/search_gart.sh "como funciona o export FCPXML"
    rag/search_gart.sh "FCPXMLModifier" --snippet
    rag/search_gart.sh "voice_actions" --module fcpxml --json
    rag/search_gart.sh --map fcpxml          # inventário do módulo

Como módulo:
    from search import rag_search
    rag_search("consulta", top_k=5)

Por que busca híbrida
---------------------
A busca puramente densa erra nomes exatos: procurar `FCPXMLModifier` pode não
trazer `modifier.py` no topo. Por isso rodamos duas listas em paralelo —
densa (embedding) e lexical (pg_trgm sobre os símbolos declarados) — e
fundimos com RRF, que soma 1/(k+posição) de cada lista e portanto não exige
normalizar escalas diferentes. Adaptado do rag/search.py dos projetos irmãos
(Doza, Tigre, Jhonny) — mesmos parâmetros, calibrados lá via rag/bench.py.
"""

import argparse
import json
import os
import sys

import psycopg2
from dotenv import load_dotenv

RAG_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(RAG_DIR)
# Por padrão reaproveita as mesmas credenciais do indexador
# (admin/gart-rag.env), para não duplicar a senha em dois arquivos.
_load_path = os.environ.get("RAG_ENV_FILE", os.path.join(ROOT, "admin", "gart-rag.env"))
load_dotenv(_load_path)

if RAG_DIR not in sys.path:
    sys.path.insert(0, RAG_DIR)

from embed import QUERY_PREFIX, _embed  # noqa: E402

# Constante do Reciprocal Rank Fusion e pesos por lista — valores herdados
# dos projetos irmãos (calibrados lá via varredura sobre consultas douradas,
# ver rag/bench.py no projeto Doza/Tigre). Ao mexer nesses números, monte um
# conjunto de consultas de referência para o G-ART antes.
RRF_K = 8
W_DENSE = 1.0
W_LEX = 1.0       # multiplicado pela similaridade bruta do casamento
W_SUMMARY = 1.5   # o resumo é curto e preciso: um acerto ali vale mais

# Piso do casamento lexical: abaixo disso o lado lexical se cala em vez de
# afogar a lista densa com ruído.
LEX_MIN = 0.5

# Quantos chunks do mesmo arquivo podem ocupar o top-k, para um arquivo
# grande não tomar todos os lugares.
MAX_PER_FILE = 2

# Quanto código o modo --snippet mostra por resultado.
SNIPPET_CHARS = 300

_EXTRA_COLS = ("start_line", "end_line", "symbols", "module", "kind")
_cols_cache = {}


def _available_cols(cur, schema):
    if schema not in _cols_cache:
        cur.execute(
            """SELECT column_name FROM information_schema.columns
               WHERE table_schema = %s AND table_name = 'code_chunks'""",
            (schema,),
        )
        _cols_cache[schema] = {r[0] for r in cur.fetchall()}
    return _cols_cache[schema]


def _select_cols(cur, schema):
    have = _available_cols(cur, schema)
    extra = ", ".join(f"c.{c}" if c in have else f"NULL AS {c}"
                      for c in _EXTRA_COLS)
    return f"c.id, c.file_path, c.content, {extra}"


def _db_connect():
    return psycopg2.connect(
        host=os.environ.get("RAG_DB_HOST", "127.0.0.1"),
        port=os.environ.get("RAG_DB_PORT", "55435"),
        dbname=os.environ.get("RAG_DB_NAME", "rag_gart"),
        user=os.environ.get("RAG_DB_USER", "gart_rag_indexer"),
        password=os.environ["RAG_DB_PASSWORD"],
        connect_timeout=5,
    )


def _schema():
    return os.environ.get("RAG_DB_SCHEMA", "gart")


def _qid(schema):
    """Schema como identificador SQL seguro (aspas duplas)."""
    return '"' + schema.replace('"', '""') + '"'


def _filters(module, path, ext, have=()):
    """Cláusulas de escopo aplicadas antes do ranqueamento."""
    clauses, params = [], []
    if module and "module" in have:
        clauses.append("c.module = %s")
        params.append(module)
    if path:
        clauses.append("c.file_path ILIKE %s")
        params.append(f"%{path}%")
    if ext:
        clauses.append("c.file_path LIKE %s")
        params.append(f"%{ext}")
    return (" AND " + " AND ".join(clauses) if clauses else ""), params


def _probe_terms(query):
    """Termos que o lado lexical tenta casar contra os símbolos: a consulta
    inteira e a versão sem espaços (faz "voice timeline" casar com
    `VoiceTimeline`)."""
    terms = [query, query.replace(" ", "")]
    return list(dict.fromkeys(t for t in terms if t))


def _row_to_dict(row):
    return {
        "id": row[0], "file_path": row[1], "content": row[2],
        "start_line": row[3], "end_line": row[4],
        "symbols": row[5], "module": row[6], "kind": row[7],
    }


def _dense(cur, schema, q_emb, limit, where, params):
    cur.execute("SET LOCAL hnsw.ef_search = 64")
    cur.execute(
        f"""
        SELECT {_select_cols(cur, schema)}, 1 - (c.embedding <=> %s::vector) AS s
        FROM {_qid(schema)}.code_chunks c
        WHERE c.embedding IS NOT NULL {where}
        ORDER BY c.embedding <=> %s::vector
        LIMIT %s
        """,
        [q_emb] + params + [q_emb, limit],
    )
    return [(_row_to_dict(r), float(r[8])) for r in cur.fetchall()]


def _lexical(cur, schema, terms, limit, where, params):
    have = _available_cols(cur, schema)
    stem = "regexp_replace(c.file_path, '^.*/|\\.[^.]*$', '', 'g')"
    target = f"coalesce(c.symbols, {stem})" if "symbols" in have else stem
    tiebreak = "c.start_line" if "start_line" in have else "c.id"
    score = f"(SELECT max(word_similarity(t, {target})) FROM unnest(%s::text[]) t)"
    cur.execute(
        f"""
        SELECT {_select_cols(cur, schema)}, {score} AS s
        FROM {_qid(schema)}.code_chunks c
        WHERE {score} >= %s {where}
        ORDER BY s DESC, {tiebreak} ASC
        LIMIT %s
        """,
        [terms] + [terms, LEX_MIN] + params + [limit],
    )
    return [(_row_to_dict(r), float(r[8])) for r in cur.fetchall()]


def _summary_dense(cur, schema, q_emb, limit, where, params):
    """Terceira lista: busca sobre o resumo do arquivo (file_index), não
    sobre o código — resgata arquivos pequenos e precisos que a lista por
    trecho, dominada por arquivos grandes, deixa passar."""
    cur.execute("SELECT to_regclass(%s)", (f"{_qid(schema)}.file_index",))
    if cur.fetchone()[0] is None:
        return []
    cur.execute("SET LOCAL hnsw.ef_search = 64")
    cur.execute(
        f"""
        SELECT file_path, 1 - (summary_embedding <=> %s::vector) AS s
        FROM {_qid(schema)}.file_index
        WHERE summary_embedding IS NOT NULL
        ORDER BY summary_embedding <=> %s::vector
        LIMIT %s
        """,
        (q_emb, q_emb, limit),
    )
    hits = cur.fetchall()
    if not hits:
        return []
    order = {fp: i for i, (fp, _) in enumerate(hits)}
    raw = {fp: s for fp, s in hits}

    cur.execute(
        f"""
        SELECT DISTINCT ON (c.file_path) {_select_cols(cur, schema)}
        FROM {_qid(schema)}.code_chunks c
        WHERE c.file_path = ANY(%s) AND c.embedding IS NOT NULL {where}
        ORDER BY c.file_path, c.embedding <=> %s::vector
        """,
        [list(order)] + params + [q_emb],
    )
    rows = [_row_to_dict(r) for r in cur.fetchall()]
    rows.sort(key=lambda r: order[r["file_path"]])
    return [(r, raw[r["file_path"]]) for r in rows]


def _fuse(ranked_lists):
    """Reciprocal Rank Fusion ponderada sobre listas já ordenadas."""
    scores, best = {}, {}
    for results, weight_of in ranked_lists:
        for rank, (row, raw) in enumerate(results, 1):
            contrib = weight_of(raw) / (RRF_K + rank)
            scores[row["id"]] = scores.get(row["id"], 0.0) + contrib
            best.setdefault(row["id"], row)
    ordered = sorted(scores.items(), key=lambda kv: -kv[1])
    return [dict(best[i], score=s) for i, s in ordered]


def _dedupe(rows, top_k, max_per_file=MAX_PER_FILE):
    """Limita chunks por arquivo; o excedente vira uma nota de localização."""
    kept, counts, extras = [], {}, {}
    for row in rows:
        fp = row["file_path"]
        if counts.get(fp, 0) < max_per_file:
            counts[fp] = counts.get(fp, 0) + 1
            row["also_at"] = []
            kept.append(row)
        else:
            extras.setdefault(fp, []).append((row["start_line"], row["end_line"]))
    for row in kept:
        row["also_at"] = extras.get(row["file_path"], [])[:3]
    return kept[:top_k]


def _attach_map(cur, schema, rows):
    """Anexa tipo principal e resumo de `file_index`, quando existir."""
    cur.execute("SELECT to_regclass(%s)", (f"{_qid(schema)}.file_index",))
    if cur.fetchone()[0] is None or not rows:
        return rows
    paths = list({r["file_path"] for r in rows})
    cur.execute(
        f"SELECT file_path, main_type, summary FROM {_qid(schema)}.file_index "
        f"WHERE file_path = ANY(%s)",
        (paths,),
    )
    info = {p: (t, s) for p, t, s in cur.fetchall()}
    for row in rows:
        main_type, summary = info.get(row["file_path"], (None, None))
        row["main_type"] = main_type
        row["summary"] = summary
    return rows


def rag_search(query, top_k=5, module=None, path=None, ext=None, dense_only=False):
    """Busca híbrida. Devolve dicts com file_path, faixa de linhas e score."""
    schema = _schema()
    pool = max(top_k * 4, 20)

    conn = _db_connect()
    try:
        with conn.cursor() as cur:
            where, params = _filters(module, path, ext,
                                     _available_cols(cur, schema))
            q_emb = _embed(query, prefix=QUERY_PREFIX)
            lists = [(_dense(cur, schema, q_emb, pool, where, params),
                      lambda raw: W_DENSE)]
            if not dense_only:
                lists.append((_lexical(cur, schema, _probe_terms(query),
                                       pool, where, params),
                              lambda raw: W_LEX * (0.5 + raw)))
                lists.append((_summary_dense(cur, schema, q_emb, pool,
                                             where, params),
                              lambda raw: W_SUMMARY))
            rows = _dedupe(_fuse(lists), top_k)
            rows = _attach_map(cur, schema, rows)
        conn.commit()
    finally:
        conn.close()
    return rows


def map_files(term=None, module=None, limit=60):
    """Inventário de arquivos — responde "onde fica X" sem corpo de código."""
    schema = _schema()
    conn = _db_connect()
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT to_regclass(%s)", (f"{_qid(schema)}.file_index",))
            if cur.fetchone()[0] is None:
                return []
            clauses, params = [], []
            if module:
                clauses.append("module = %s")
                params.append(module)
            if term:
                clauses.append("(file_path ILIKE %s OR main_type ILIKE %s "
                               "OR summary ILIKE %s)")
                params += [f"%{term}%"] * 3
            where = "WHERE " + " AND ".join(clauses) if clauses else ""
            cur.execute(
                f"""SELECT file_path, module, main_type, summary, n_lines
                    FROM {_qid(schema)}.file_index {where}
                    ORDER BY module, file_path LIMIT %s""",
                params + [limit],
            )
            return [
                {"file_path": r[0], "module": r[1], "main_type": r[2],
                 "summary": r[3], "n_lines": r[4]}
                for r in cur.fetchall()
            ]
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# Formatação
# ---------------------------------------------------------------------------

_NOISE = {'"""', "'''", "# ---", "---", "/*", "*/", "{", "}", "*"}


def _first_doc_line(content):
    """Primeira linha que serve de descrição do trecho, usada quando o
    arquivo não tem resumo em `file_index`."""
    for line in content.splitlines():
        s = line.strip()
        if s.startswith(("///", "//", "#", "*")):
            s = s.lstrip("/#* ").strip()
            if s and s not in _NOISE:
                return s
    for line in content.splitlines():
        s = line.strip().strip("\"'").strip()
        if s and s not in _NOISE and not s.startswith("import"):
            return s
    return ""


def format_results(results, mode="map"):
    """Renderiza o resultado no modo pedido. O padrão é `map`: caminho +
    faixa de linhas + uma linha de descrição, sem corpo de código."""
    if not results:
        return "[RAG vazio, buscando local]\n"

    if mode == "json":
        return json.dumps(results, ensure_ascii=False, indent=2) + "\n"

    out = []
    for i, r in enumerate(results, 1):
        loc = r["file_path"]
        if r.get("start_line") and r.get("end_line"):
            loc += f":{r['start_line']}-{r['end_line']}"
        out.append(f"{i}  {r['score']:.2f}  {loc}")

        desc = r.get("summary") or _first_doc_line(r["content"])
        stem = os.path.splitext(os.path.basename(r["file_path"]))[0]
        label = r.get("main_type") or ""
        if label == stem:
            label = ""
        if label and desc:
            out.append(f"          {label} · {desc[:78]}")
        elif desc:
            out.append(f"          {desc[:88]}")

        spans = [f"{a}-{b}" for a, b in r.get("also_at", []) if a and b]
        if spans:
            out.append(f"          (+ tambem em {', '.join(spans)})")

        if mode == "snippet":
            body, size = [], 0
            for ln in r["content"].splitlines():
                if body and size + len(ln) > SNIPPET_CHARS:
                    body.append("…")
                    break
                body.append(ln)
                size += len(ln) + 1
            out.append("".join(f"          | {ln}\n" for ln in body))
        elif mode == "full":
            out.append("".join(f"          | {ln}\n" for ln in r["content"].splitlines()))
    return "\n".join(out) + "\n"


def format_map(rows):
    if not rows:
        return "[RAG vazio, buscando local]\n"
    out = []
    current = None
    for r in rows:
        if r["module"] != current:
            current = r["module"]
            out.append(f"\n{current or '(sem modulo)'}")
        name = os.path.basename(r["file_path"])
        desc = (r["summary"] or "")[:78]
        out.append(f"  {name:<38} {r['n_lines']:>5}L  {desc}")
    return "\n".join(out) + "\n"


def main():
    ap = argparse.ArgumentParser(
        description="Busca RAG hibrida (densa + lexical, fundidas com RRF) do G-ART")
    ap.add_argument("query", nargs="?", help="consulta")
    ap.add_argument("top_k", nargs="?", type=int, default=5)
    ap.add_argument("--snippet", action="store_true", help="mostra 300 chars do trecho")
    ap.add_argument("--full", action="store_true", help="mostra o trecho inteiro")
    ap.add_argument("--json", action="store_true", help="saida estruturada")
    ap.add_argument("--module", help="restringe a um modulo (ex: fcpxml)")
    ap.add_argument("--path", help="restringe a caminhos contendo este texto")
    ap.add_argument("--ext", help="restringe a uma extensao (ex: .py)")
    ap.add_argument("--map", dest="map_term", nargs="?", const="",
                    help="inventario de arquivos em vez de busca por trecho")
    ap.add_argument("--dense-only", action="store_true",
                    help="desliga o lado lexical (para comparacao)")
    args = ap.parse_args()

    if args.map_term is not None:
        print(format_map(map_files(term=args.map_term or None, module=args.module)))
        return

    if not args.query:
        ap.error("informe a consulta, ou use --map")

    mode = ("json" if args.json else "full" if args.full
            else "snippet" if args.snippet else "map")
    results = rag_search(args.query, args.top_k, module=args.module,
                         path=args.path, ext=args.ext, dense_only=args.dense_only)
    print(format_results(results, mode=mode))


if __name__ == "__main__":
    main()