Compressão contextual e filtragem por relevância

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_openai import ChatOpenAI

# LLMChainExtractor uses an LLM to extract the relevant portion
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
compressor = LLMChainExtractor.from_llm(llm)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever,
)

results = compression_retriever.invoke('how does HNSW indexing work?')
for doc in results:
    print(len(doc.page_content), 'chars:', doc.page_content[:150])

from langchain.retrievers.document_compressors import LLMChainFilter

filter_compressor = LLMChainFilter.from_llm(
    llm=ChatOpenAI(model='gpt-4o-mini', temperature=0)
)

filtering_retriever = ContextualCompressionRetriever(
    base_compressor=filter_compressor,
    base_retriever=base_retriever,
)

# Fetch 10 docs, filter drops irrelevant ones
results = filtering_retriever.invoke('what are the pgvector distance operators?')
print(f'{len(results)} chunks passed the relevance filter')

from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings()
embeddings_filter = EmbeddingsFilter(
    embeddings=embeddings,
    similarity_threshold=0.76,  # cosine similarity threshold
)

embedding_retriever = ContextualCompressionRetriever(
    base_compressor=embeddings_filter,
    base_retriever=base_retriever,
)

results = embedding_retriever.invoke('BM25 hyperparameter tuning')
print(f'Filtered to {len(results)} relevant chunks')

from langchain.retrievers.document_compressors import DocumentCompressorPipeline
from langchain_community.document_transformers import EmbeddingsRedundantFilter
from langchain_text_splitters import CharacterTextSplitter

# Step 1: split chunks into individual sentences
sentence_splitter = CharacterTextSplitter(
    chunk_size=200,
    chunk_overlap=0,
    separator='. ',
)

# Step 2: remove redundant sentences via embeddings
redundant_filter = EmbeddingsRedundantFilter(embeddings=OpenAIEmbeddings())

# Step 3: keep only relevant sentences
relevance_filter = EmbeddingsFilter(embeddings=OpenAIEmbeddings(), similarity_threshold=0.76)

pipeline = DocumentCompressorPipeline(
    transformers=[sentence_splitter, redundant_filter, relevance_filter]
)

piped_retriever = ContextualCompressionRetriever(
    base_compressor=pipeline,
    base_retriever=base_retriever,
)

from langchain_community.document_transformers import EmbeddingsRedundantFilter
from langchain_core.documents import Document

redundant_filter = EmbeddingsRedundantFilter(
    embeddings=OpenAIEmbeddings(),
    similarity_threshold=0.95,  # treat docs with >95% cosine sim as duplicates
)

# Simulate near-duplicate documents
docs = [
    Document(page_content='pgvector is a PostgreSQL extension for vector similarity search'),
    Document(page_content='pgvector extends PostgreSQL to support vector similarity search'),  # near duplicate
    Document(page_content='HNSW and IVFFlat are the two index types in pgvector'),
]

filtered = redundant_filter.transform_documents(docs, query='')
print(f'Kept {len(filtered)} of {len(docs)} documents after deduplication')

from sentence_transformers import CrossEncoder
import re

sentence_reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def compress_with_cross_encoder(
    query: str, chunk: str, min_score: float = 0.0
) -> str:
    sentences = re.split(r'(?<=[.!?]) +', chunk)
    if len(sentences) <= 1:
        return chunk

    pairs = [[query, s] for s in sentences]
    scores = sentence_reranker.predict(pairs)

    relevant = [
        sentence
        for sentence, score in zip(sentences, scores)
        if float(score) >= min_score
    ]
    return ' '.join(relevant) if relevant else chunk

def compress_to_budget(
    query: str,
    docs: list[str],
    total_token_budget: int = 2000,
    tokens_per_char: float = 0.25,
) -> list[str]:
    compressed = []
    used_tokens = 0

    for doc in docs:
        compressed_doc = compress_with_cross_encoder(query, doc)
        doc_tokens = int(len(compressed_doc) * tokens_per_char)

        if used_tokens + doc_tokens <= total_token_budget:
            compressed.append(compressed_doc)
            used_tokens += doc_tokens
        else:
            break  # stop when budget is exhausted

    print(f'Using {used_tokens} tokens across {len(compressed)} docs')
    return compressed

def measure_compression_faithfulness(query, original_docs, compressed_docs, llm):
    # Use LLM to check if answer from compressed context matches
    # answer from full context
    def generate_answer(docs, q):
        context = '\n'.join(docs)
        resp = llm.invoke(f'Answer from context:\n{context}\nQ: {q}')
        return resp.content

    full_answer = generate_answer([d.page_content for d in original_docs], query)
    comp_answer = generate_answer(compressed_docs, query)

    # Check semantic similarity between answers
    vecs = [embed(full_answer), embed(comp_answer)]
    sim = cosine_similarity(vecs[0], vecs[1])
    print(f'Answer similarity after compression: {sim:.3f}')
    return sim

class AdaptiveCompressor:
    def __init__(self, embeddings, cross_encoder=None, threshold=0.76):
        self.emb_filter = EmbeddingsFilter(embeddings=embeddings,
                                           similarity_threshold=threshold)
        self.dedup = EmbeddingsRedundantFilter(embeddings=embeddings)
        self.cross_encoder = cross_encoder

    def compress(self, query: str, docs, use_cross_encoder: bool = False):
        # Stage 1: embedding filter
        docs = self.emb_filter.compress_documents(docs, query=query)
        # Stage 2: deduplication
        docs = self.dedup.transform_documents(docs, query=query)
        # Stage 3: optional sentence-level extraction
        if use_cross_encoder and self.cross_encoder:
            docs = [
                type(d)(page_content=compress_with_cross_encoder(
                    query, d.page_content
                ), metadata=d.metadata)
                for d in docs
            ]
        return docs