Contextual Compression: сжатие чанков до релевантных фрагментов

from anthropic import Anthropic
from sentence_transformers import CrossEncoder

client  = Anthropic()
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

EXTRACT_PROMPT = """\
Из текста ниже извлеки только фрагменты, которые отвечают на вопрос пользователя.
Извлекай дословно из текста — не перефразируй и не добавляй своих слов.
Оставь достаточно контекста, чтобы фрагмент был понятен без исходного текста.
Если ничего не подходит — ответь одним словом: IRRELEVANT

Вопрос: {query}

Текст:
{chunk}"""


def llm_extract(query: str, chunk: str, min_logit: float = 0.0) -> str | None:
    """
    Извлекает из чанка релевантный фрагмент с помощью LLM.

    Предварительно проверяет релевантность через cross-encoder:
    если logit < min_logit, LLM не вызывается → экономия токенов.
    """
    # Быстрая проверка: стоит ли вообще вызывать LLM
    logit = float(reranker.predict([[query, chunk]])[0])
    if logit < min_logit:
        return None  # чанк нерелевантен — отбрасываем без LLM-вызова

    response = client.messages.create(
        model="claude-haiku-4-5-20251001",  # дешёвая модель для preprocessing
        max_tokens=512,
        messages=[{
            "role": "user",
            "content": EXTRACT_PROMPT.format(query=query, chunk=chunk),
        }],
    )

    result = response.content[0].text.strip()
    return None if result == "IRRELEVANT" else result


def compress_chunks(
    query: str,
    chunks: list[str],
) -> list[dict]:
    """
    Применяет LLM Extractor ко всем чанкам параллельно.
    Возвращает только те, из которых удалось извлечь фрагмент.
    """
    results = []
    for i, chunk in enumerate(chunks):
        extracted = llm_extract(query, chunk)
        if extracted:
            results.append({
                "index":    i,
                "original": chunk,
                "compressed": extracted,
                "ratio":    round(len(extracted) / len(chunk), 2),
            })
    return results


# ── Пример использования ───────────────────────────────────────────
query = "как настроить timeout в httpx"

chunks = [
    "httpx — современный HTTP-клиент для Python. Поддерживает HTTP/1.1 и HTTP/2. "
    "Совместим с asyncio. Для настройки таймаутов используйте объект Timeout: "
    "httpx.Timeout(connect=5.0, read=10.0). Можно указать разные значения для "
    "connect, read, write и pool. Значение None отключает таймаут.",

    "requests — популярная библиотека для HTTP-запросов. Имеет простой API. "
    "Поддерживает сессии через requests.Session(). Connection pooling включён "
    "автоматически. Позволяет задавать заголовки, куки и параметры запроса.",

    "httpx.AsyncClient поддерживает тот же интерфейс таймаутов: "
    "async with httpx.AsyncClient(timeout=httpx.Timeout(5.0)) as client: ... "
    "По умолчанию connect timeout = 5 секунд, read timeout = 5 секунд.",
]

compressed = compress_chunks(query, chunks)

for item in compressed:
    original_len = len(item["original"].split())
    compressed_len = len(item["compressed"].split())
    print(f"Чанк {item['index']}: {original_len} слов → {compressed_len} слов")
    print(f"  Фрагмент: {item['compressed'][:120]}...")
    print()

from anthropic import Anthropic

client = Anthropic()

FILTER_PROMPT = """\
Тебе дан текст и вопрос пользователя.
Ответь YES, если текст содержит информацию, полезную для ответа на вопрос.
Ответь NO, если текст не содержит ничего полезного.
Отвечай только одним словом: YES или NO.

Вопрос: {query}

Текст:
{chunk}"""


def llm_filter(query: str, chunks: list[str]) -> list[str]:
    """
    Фильтрует чанки: оставляет только те, где LLM ответил YES.
    Чанки передаются целиком (без извлечения фрагментов).
    """
    relevant = []
    for chunk in chunks:
        response = client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=5,   # нам нужен только YES/NO
            messages=[{
                "role": "user",
                "content": FILTER_PROMPT.format(query=query, chunk=chunk),
            }],
        )
        verdict = response.content[0].text.strip().upper()
        if verdict.startswith("YES"):
            relevant.append(chunk)

    return relevant


# ── Пример ─────────────────────────────────────────────────────────
query = "как работает retry в httpx"

chunks = [
    "httpx поддерживает transport-уровень retry через httpx.HTTPTransport(retries=3). "
    "Это автоматически повторяет запрос при сетевых ошибках (ConnectError, ReadError).",

    "FastAPI — фреймворк для построения REST API. Использует Pydantic для валидации "
    "данных. Поддерживает async/await из коробки.",

    "Для сложной логики retry (экспоненциальный backoff, условия) используют tenacity: "
    "@retry(wait=wait_exponential(), stop=stop_after_attempt(3)) над функцией с httpx.",
]

result = llm_filter(query, chunks)
print(f"Осталось {len(result)} из {len(chunks)} чанков")
# → Осталось 2 из 3 чанков (FastAPI-чанк отброшен)

import re
import numpy as np
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")


def split_sentences(text: str) -> list[str]:
    """Простой сплит на предложения с сохранением нормальной длины."""
    # Разбиваем по точке/восклицательному/вопросительному знаку + пробел или перенос
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    # Убираем совсем короткие (< 10 символов) — обычно это артефакты
    return [s for s in sentences if len(s) > 10]


def embeddings_filter(
    query: str,
    chunks: list[str],
    threshold: float = 0.55,
    min_sentences: int = 1,
) -> list[str]:
    """
    Фильтрует предложения в каждом чанке по cosine similarity с запросом.

    threshold: минимальная cosine similarity (0.55 — хорошая точка отсчёта)
    min_sentences: если после фильтрации осталось меньше — чанк отбрасывается целиком
    """
    # Один раз векторизуем запрос
    query_vec = model.encode([query], normalize_embeddings=True)[0]

    result = []
    for chunk in chunks:
        sentences = split_sentences(chunk)
        if not sentences:
            continue

        # Векторизуем все предложения чанка батчем
        sent_vecs = model.encode(sentences, normalize_embeddings=True)

        # Cosine similarity = dot product для нормализованных векторов
        scores = sent_vecs @ query_vec

        # Оставляем предложения выше порога, сохраняя исходный порядок
        kept = [s for s, score in zip(sentences, scores) if score >= threshold]

        if len(kept) >= min_sentences:
            result.append(" ".join(kept))

    return result


# ── Пример ─────────────────────────────────────────────────────────
query = "как настроить timeout в httpx"

chunk = (
    "httpx — современный HTTP-клиент для Python. "
    "Поддерживает HTTP/1.1 и HTTP/2. "
    "Для настройки таймаутов используйте httpx.Timeout(connect=5.0, read=10.0). "
    "Можно задать разные значения для connect, read, write и pool. "
    "Библиотека совместима с asyncio и trio. "
    "Connection pooling включён автоматически."
)

compressed = embeddings_filter(query, [chunk], threshold=0.55)
print("До:", len(chunk.split()), "слов")
print("После:", len(compressed[0].split()) if compressed else 0, "слов")
print("Результат:", compressed[0] if compressed else "—")

from langchain_anthropic import ChatAnthropic
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import (
    LLMChainExtractor,   # LLM Extractor
    LLMChainFilter,      # LLM Filter
    EmbeddingsFilter,    # Embeddings Filter
)

# ── Базовый ретривер ────────────────────────────────────────────────
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
vectorstore = Chroma(embedding_function=embeddings, persist_directory="./chroma_db")
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# ── Вариант 1: LLM Extractor ────────────────────────────────────────
llm = ChatAnthropic(model="claude-haiku-4-5-20251001", temperature=0)

extractor = LLMChainExtractor.from_llm(llm)
compression_retriever_extract = ContextualCompressionRetriever(
    base_compressor=extractor,
    base_retriever=base_retriever,
)

# ── Вариант 2: LLM Filter ───────────────────────────────────────────
llm_filter = LLMChainFilter.from_llm(llm)
compression_retriever_filter = ContextualCompressionRetriever(
    base_compressor=llm_filter,
    base_retriever=base_retriever,
)

# ── Вариант 3: Embeddings Filter (без LLM) ──────────────────────────
emb_filter = EmbeddingsFilter(
    embeddings=embeddings,
    similarity_threshold=0.55,
)
compression_retriever_emb = ContextualCompressionRetriever(
    base_compressor=emb_filter,
    base_retriever=base_retriever,
)

# ── Использование в chain ───────────────────────────────────────────
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

prompt = ChatPromptTemplate.from_template(
    "Ответь на вопрос, используя только предоставленный контекст.\n\n"
    "Контекст:\n{context}\n\n"
    "Вопрос: {question}"
)

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# Цепочка с LLM Extractor (лучшее качество)
chain = (
    {
        "context":  compression_retriever_extract | format_docs,
        "question": RunnablePassthrough(),
    }
    | prompt
    | ChatAnthropic(model="claude-sonnet-4-6", temperature=0)
    | StrOutputParser()
)

answer = chain.invoke("как настроить timeout в httpx?")
print(answer)

import asyncio
from anthropic import AsyncAnthropic

client = AsyncAnthropic()

EXTRACT_PROMPT = """\
Извлеки из текста только фрагменты, релевантные вопросу.
Ответь дословными цитатами или словом IRRELEVANT.

Вопрос: {query}
Текст:
{chunk}"""


async def extract_one(query: str, chunk: str) -> str | None:
    """Один асинхронный вызов для одного чанка."""
    response = await client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=512,
        messages=[{
            "role": "user",
            "content": EXTRACT_PROMPT.format(query=query, chunk=chunk),
        }],
    )
    result = response.content[0].text.strip()
    return None if result == "IRRELEVANT" else result


async def compress_parallel(query: str, chunks: list[str]) -> list[str]:
    """
    Параллельное сжатие всех чанков.
    Latency = max(один вызов), а не sum(все вызовы).
    """
    tasks = [extract_one(query, chunk) for chunk in chunks]
    results = await asyncio.gather(*tasks)
    # Фильтруем None (нерелевантные чанки)
    return [r for r in results if r is not None]


# ── Запуск ──────────────────────────────────────────────────────────
async def main():
    query = "как настроить timeout в httpx"
    chunks = [...]  # список из 5 чанков

    import time
    t0 = time.perf_counter()
    compressed = await compress_parallel(query, chunks)
    elapsed = time.perf_counter() - t0

    print(f"Сжато {len(compressed)} чанков за {elapsed:.2f}с")

asyncio.run(main())