RAG e Agentic Search com Claude: chunking, embeddings e BM25

# RAG agêntico com Claude: retrieval como tool use

import anthropic, json
from typing import Callable

client = anthropic.Anthropic()

# Definição das tools de retrieval
retrieval_tools = [
    {
        "name": "buscar_documentos",
        "description": """Busca documentos relevantes no corpus usando busca híbrida.
                         Use quando precisar de informação específica não presente no contexto.
                         Reformule a query para maximizar relevância — queries específicas
                         retornam melhores resultados que queries genéricas.
                         Retorna top-5 chunks com score de relevância.""",
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Query de busca otimizada para recuperar a informação necessária"
                },
                "filtros": {
                    "type": "object",
                    "description": "Filtros opcionais: {'tipo': 'contrato', 'ano': 2024}",
                    "additionalProperties": True
                }
            },
            "required": ["query"]
        }
    }
]

# Implementação de busca híbrida com Reciprocal Rank Fusion (RRF):

from sentence_transformers import SentenceTransformer
from rank_bm25 import BM25Okapi
import numpy as np

class HybridRetriever:
    def __init__(self, documents: list[dict]):
        self.documents = documents
        self.texts = [d["content"] for d in documents]

        # BM25 para keyword matching
        tokenized = [text.lower().split() for text in self.texts]
        self.bm25 = BM25Okapi(tokenized)

        # Embedding para busca semântica
        self.model = SentenceTransformer("intfloat/multilingual-e5-base")
        self.embeddings = self.model.encode(self.texts, normalize_embeddings=True)

    def search(self, query: str, top_k: int = 5, alpha: float = 0.5) -> list[dict]:
        """
        Busca híbrida com RRF.
        alpha: peso do embedding vs BM25 (0.5 = igual, 1.0 = só embedding)
        """
        # BM25 scores
        bm25_scores = self.bm25.get_scores(query.lower().split())
        bm25_ranked = np.argsort(bm25_scores)[::-1]

        # Embedding scores
        query_emb = self.model.encode([query], normalize_embeddings=True)
        emb_scores = (query_emb @ self.embeddings.T)[0]
        emb_ranked = np.argsort(emb_scores)[::-1]

        # Reciprocal Rank Fusion
        k = 60  # constante RRF padrão
        rrf_scores = {}
        for rank, idx in enumerate(bm25_ranked):
            rrf_scores[idx] = rrf_scores.get(idx, 0) + (1 - alpha) / (k + rank + 1)
        for rank, idx in enumerate(emb_ranked):
            rrf_scores[idx] = rrf_scores.get(idx, 0) + alpha / (k + rank + 1)

        # Top-k resultados
        top_indices = sorted(rrf_scores, key=rrf_scores.get, reverse=True)[:top_k]
        return [
            {**self.documents[i], "relevance_score": rrf_scores[i]}
            for i in top_indices
        ]

# Por que RRF funciona:
# BM25: "art. 15" encontra documentos com exatamente "art. 15" (keyword exact)
# Embedding: "rescisão" encontra "encerramento do contrato" (semântica)
# RRF combina os rankings sem precisar normalizar scores em escala comum

# Loop completo de RAG agêntico:

retriever = HybridRetriever(carregar_documentos())  # seu corpus

def executar_busca(query: str, filtros: dict = None) -> str:
    """Função que o loop chama quando Claude usa a tool."""
    resultados = retriever.search(query, top_k=5)
    if filtros:
        resultados = [r for r in resultados
                     if all(r.get(k) == v for k, v in filtros.items())]

    return json.dumps([{
        "id": r["id"],
        "titulo": r["titulo"],
        "conteudo": r["content"][:1000],  # limita por chunk
        "relevancia": round(r["relevance_score"], 3)
    } for r in resultados], ensure_ascii=False)

# Loop agêntico:
def rag_agentico(pergunta: str, max_buscas: int = 10) -> str:
    messages = [{"role": "user", "content": pergunta}]
    buscas_realizadas = 0

    while True:
        response = client.messages.create(
            model="claude-opus-4-6",
            max_tokens=2048,
            tools=retrieval_tools,
            system="""Você é um assistente especializado em consultar documentos.
                     Use a ferramenta buscar_documentos quando precisar de informação.
                     Reformule queries se os resultados não forem suficientes.
                     Responda apenas baseado nos documentos recuperados.""",
            messages=messages
        )

        if response.stop_reason == "end_turn":
            return response.content[0].text

        if response.stop_reason == "tool_use" and buscas_realizadas < max_buscas:
            messages.append({"role": "assistant", "content": response.content})
            tool_results = []

            for block in response.content:
                if block.type == "tool_use":
                    buscas_realizadas += 1
                    resultado = executar_busca(
                        block.input["query"],
                        block.input.get("filtros")
                    )
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": resultado
                    })

            messages.append({"role": "user", "content": tool_results})

        else:
            # Limite de buscas atingido — Claude responde com o que tem
            messages.append({
                "role": "user",
                "content": "Limite de buscas atingido. Responda com as informações obtidas."
            })

# Uso:
resposta = rag_agentico(
    "Quais são as condições de rescisão nos contratos de 2023 "
    "e como elas diferem dos contratos de 2022?"
)
# Claude vai buscar: "rescisão contratos 2023" → "rescisão contratos 2022"
# → "diferenças cláusulas rescisão" e combinar os resultados

# Cross-encoder reranking: segunda etapa após busca inicial
# Retrieve broad (top-20) → rerank precision (top-5)

from sentence_transformers import CrossEncoder

class RerankedRetriever(HybridRetriever):
    def __init__(self, documents: list[dict]):
        super().__init__(documents)
        # Cross-encoder: mais pesado mas muito mais preciso
        self.reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

    def search_with_reranking(self, query: str, top_k: int = 5) -> list[dict]:
        # Fase 1: retrieval amplo (candidatos)
        candidates = self.search(query, top_k=20)

        # Fase 2: reranking preciso
        pairs = [(query, doc["content"][:512]) for doc in candidates]
        rerank_scores = self.reranker.predict(pairs)

        # Ordena por score do cross-encoder
        reranked = sorted(
            zip(candidates, rerank_scores),
            key=lambda x: x[1],
            reverse=True
        )

        return [
            {**doc, "rerank_score": float(score)}
            for doc, score in reranked[:top_k]
        ]

# Por que o reranking importa:
# Bi-encoder (embedding): rápido, pode errar em nuance semântica
# Cross-encoder: lento, olha query E documento juntos — mais preciso
# Padrão: retrieve 20 com bi-encoder → rerank top-5 com cross-encoder