João Henrique
b541f502ba
- Adicionado estrutura completa do projeto - Configurado MCP server para Premiere Pro - Adicionado documentação e skills - Configurado Gitignore para o projeto
491 lines
19 KiB
Python
491 lines
19 KiB
Python
"""Busca semântica RAG do projeto.
|
||
|
||
Consulta `<schema>.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_tigre.sh "como funciona o export FCPXML"
|
||
rag/search_tigre.sh "FCPXMLValidator" --snippet
|
||
rag/search_tigre.sh "wizard" --module TigreAppUI --json
|
||
rag/search_tigre.sh --map TigreAI # 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 `FCPXMLValidator` não
|
||
trazia `FCPXMLValidator.swift`, e `SilenceCutPipeline.swift` não aparecia em
|
||
nenhum top-5. 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. Ver rag/bench.py para os números antes/depois.
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
|
||
import psycopg2
|
||
from dotenv import load_dotenv
|
||
|
||
RAG_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
# Mesma mecânica do indexador: respeita RAG_ENV_FILE para permitir wrapper
|
||
# por sistema (ex: search_tigre.sh com RAG_ENV_FILE=tigre.env).
|
||
_load_path = os.environ.get("RAG_ENV_FILE", os.path.join(RAG_DIR, ".env"))
|
||
load_dotenv(_load_path)
|
||
|
||
if RAG_DIR not in sys.path:
|
||
sys.path.insert(0, RAG_DIR)
|
||
|
||
from index_code import QUERY_PREFIX, _embed # noqa: E402
|
||
|
||
# Constante do Reciprocal Rank Fusion e pesos por lista. Os valores saíram de
|
||
# uma varredura sobre o conjunto dourado (rag/bench.py), não da literatura: o
|
||
# k=60 clássico achata demais a curva para um corpus deste tamanho e deixava
|
||
# um 1º lugar isolado perder para um medíocre presente em duas listas.
|
||
# k=8 e peso 1.5 no resumo vêm de uma varredura (k × pesos) sobre as 18
|
||
# consultas de rag/bench.py. Não é o pico da grade — é o meio de um platô
|
||
# estável (k de 5 a 8 dá o mesmo recall, e em k=8 o resultado não muda com o
|
||
# peso do resumo entre 1.0 e 2.0), escolhido justamente para não sobreajustar
|
||
# um conjunto dourado pequeno. Ao mexer nesses números, rode rag/bench.sh.
|
||
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. Calibrado por medição: com os probes abaixo, um
|
||
# nome de tipo procurado diretamente pontua 1.00 e o segundo colocado no
|
||
# máximo 0.95, enquanto NENHUMA consulta conceitual em português passa de
|
||
# 0.35. Ou seja, 0.5 faz o lado lexical se calar quando não tem nada a dizer,
|
||
# em vez de afogar a lista densa com ruído.
|
||
LEX_MIN = 0.5
|
||
|
||
# Quantos chunks do mesmo arquivo podem ocupar o top-k. Sem esse limite um
|
||
# arquivo grande toma todos os lugares: numa medição o WizardModel.swift
|
||
# ficou com 3 dos 5 resultados, gastando token sem trazer nada novo.
|
||
MAX_PER_FILE = 2
|
||
|
||
# Quanto código o modo --snippet mostra por resultado.
|
||
SNIPPET_CHARS = 300
|
||
|
||
# Colunas extras introduzidas pela migration 002. O banco legado `doza` não
|
||
# as tem — este mesmo código atende os dois, então descobrimos quais existem
|
||
# e substituímos as ausentes por NULL, em vez de quebrar a busca do legado.
|
||
_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", "5433"),
|
||
dbname=os.environ.get("RAG_DB_NAME", "rag_doza"),
|
||
user=os.environ.get("RAG_DB_USER", "rag_doza_indexer"),
|
||
password=os.environ["RAG_DB_PASSWORD"],
|
||
connect_timeout=5,
|
||
)
|
||
|
||
|
||
def _schema():
|
||
return os.environ.get("RAG_DB_SCHEMA", "doza")
|
||
|
||
|
||
def _qid(schema):
|
||
"""Schema como identificador SQL seguro (aspas duplas) — aceita guífen
|
||
(ex: `jhonny-rag`) sem quebrar a sintaxe. Retroativo: produtos sem guífen
|
||
voltam idênticos (`doza` -> `"doza"`)."""
|
||
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.
|
||
|
||
Só dois: a consulta inteira e a versão sem espaços — é esta que faz
|
||
"silence cut pipeline" casar 1.00 com `SilenceCutPipeline`.
|
||
|
||
Palavra a palavra NÃO funciona, por mais tentador que pareça: termos
|
||
portugueses comuns casam com qualquer lista de símbolos ("fora" pontuou
|
||
1.00 contra `WizardStep`, "forma" 0.83 contra meia base) e o ruído
|
||
expulsava do top-5 acertos legítimos do lado denso — o `FacePerceiver`
|
||
era o 1º da lista densa para "detectar rostos" e sumia do resultado.
|
||
"""
|
||
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):
|
||
# O HNSW só explora `ef_search` candidatos; 64 dá folga sobre um top-k
|
||
# pequeno sem custar latência perceptível nesta escala.
|
||
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):
|
||
# Casa contra `symbols` (nome do arquivo + tipo + membros declarados),
|
||
# nunca contra o caminho inteiro: nomes de diretório como `Pipeline/`
|
||
# casariam com meia base e afogariam o sinal.
|
||
have = _available_cols(cur, schema)
|
||
stem = "regexp_replace(c.file_path, '^.*/|\\.[^.]*$', '', 'g')"
|
||
target = f"coalesce(c.symbols, {stem})" if "symbols" in have else stem
|
||
# Desempate estável: pela linha inicial onde ela existe, senão pelo id.
|
||
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, não sobre o código.
|
||
|
||
Devolve o melhor trecho de cada arquivo cujo resumo casa com a consulta.
|
||
Existe porque a lista por trecho é dominada por arquivos grandes — eles
|
||
casam morno com tudo e empurram para baixo arquivos pequenos e exatos.
|
||
"""
|
||
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}
|
||
|
||
# Para cada arquivo achado pelo resumo, o trecho mais próximo da consulta.
|
||
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.
|
||
|
||
RRF puro só olha a posição, e isso empatava o 1º lugar lexical com o 1º
|
||
dense — deixando de fora casos em que só o lado lexical acertava
|
||
(`AtomicFileIO` para "gravar arquivo de forma atomica"). Por isso cada
|
||
lista traz uma função de peso: a lexical escala com a similaridade
|
||
bruta, de modo que um casamento quase exato de nome vence o empate.
|
||
"""
|
||
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()
|
||
# Ampliamos o pool antes de fundir e deduplicar: o item certo pode estar
|
||
# em 8º numa lista e em 1º na outra, e a fusão é que o traz para cima.
|
||
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
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Linhas que não descrevem nada: delimitadores de docstring, cercas e
|
||
# separadores. Aparecem no topo de arquivos Python do índice legado.
|
||
_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` — caso do banco
|
||
legado, que não passou pela migration 002.
|
||
"""
|
||
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
|
||
nenhum corpo de código. É o modo mais barato e quase sempre suficiente —
|
||
com a faixa em mãos o agente lê exatamente a janela que precisa.
|
||
"""
|
||
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):
|
||
# O banco legado não tem faixa de linhas; sem ela, mostramos só o
|
||
# caminho em vez de um ":None-None" que não serve para nada.
|
||
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"])
|
||
# A regra "uma classe por arquivo" faz o tipo principal repetir o
|
||
# nome do arquivo quase sempre; imprimir os dois é token jogado fora.
|
||
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":
|
||
# Corta em fronteira de LINHA: um trecho de código partido no
|
||
# meio de um identificador não ajuda ninguém a decidir se vale
|
||
# abrir o arquivo.
|
||
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)")
|
||
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: TigreAI)")
|
||
ap.add_argument("--path", help="restringe a caminhos contendo este texto")
|
||
ap.add_argument("--ext", help="restringe a uma extensao (ex: .swift)")
|
||
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()
|