Generation Pipeline: системный промпт + контекст + вопрос

from dataclasses import dataclass


@dataclass
class RetrievedChunk:
    text:   str
    source: str   # имя файла или URL
    score:  float


def format_context(chunks: list[RetrievedChunk]) -> str:
    """
    Оборачивает чанки в XML-блок с нумерацией и источниками.
    Пример вывода:
        
        [1] Источник: httpx-guide.md
        httpx.Timeout(connect=5.0, read=10.0) ...

        [2] Источник: httpx-async.md
        AsyncClient принимает тот же объект Timeout ...
        
    """
    if not chunks:
        return "\nКонтекст не найден.\n"

    parts = []
    for i, chunk in enumerate(chunks, start=1):
        source_name = chunk.source.split("/")[-1]  # только имя файла
        parts.append(f"[{i}] Источник: {source_name}\n{chunk.text.strip()}")

    inner = "\n\n".join(parts)
    return f"\n{inner}\n"


# ── Пример ─────────────────────────────────────────────────────────
chunks = [
    RetrievedChunk(
        text="httpx.Timeout(connect=5.0, read=10.0) настраивает раздельные таймауты.",
        source="docs/httpx-guide.md",
        score=0.92,
    ),
    RetrievedChunk(
        text="AsyncClient: передайте timeout=httpx.Timeout(None) чтобы отключить таймаут.",
        source="docs/httpx-async.md",
        score=0.87,
    ),
]

print(format_context(chunks))

NO_CONTEXT_REPLY = (
    "Я не нашёл информации по этому вопросу в предоставленной документации. "
    "Попробуйте переформулировать запрос или уточните тему."
)

def check_context_quality(
    chunks: list[RetrievedChunk],
    min_score: float = 0.45,
    min_chunks: int = 1,
) -> bool:
    """
    Возвращает True, если контекст достаточно релевантен для генерации.
    Два условия: хотя бы один чанк выше порога и не пустой список.
    """
    if not chunks or len(chunks) < min_chunks:
        return False
    best_score = max(ch.score for ch in chunks)
    return best_score >= min_score


def generate_answer(
    query: str,
    chunks: list[RetrievedChunk],
    system_prompt: str,
    client,           # anthropic.Anthropic
    model: str = "claude-sonnet-4-6",
    min_score: float  = 0.45,
) -> str:
    # Проверяем качество контекста до LLM-вызова
    if not check_context_quality(chunks, min_score=min_score):
        return NO_CONTEXT_REPLY

    context = format_context(chunks)
    user_message = f"{context}\n\nВопрос: {query}"

    response = client.messages.create(
        model=model,
        max_tokens=1024,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}],
    )
    return response.content[0].text

SYSTEM_PROMPT = """\
Ты — ассистент по технической документации.

## Правила ответа
1. Отвечай ТОЛЬКО на основе информации из блока .
2. Если ответа нет в контексте — напиши: «Эта информация не содержится в документации».
3. Не добавляй факты, примеры кода или числа, которых нет в источниках.

## Цитирование
- После каждого утверждения из источника ставь номер в квадратных скобках: [1], [2], ...
- В конце ответа добавь раздел «Источники:» со списком использованных документов.
- Если утверждение следует из нескольких источников — пиши [1][2].

## Формат
- Markdown для структуры (заголовки, списки, блоки кода)
- Блоки кода оборачивай в ``` с указанием языка
- Ответ должен быть конкретным и по существу — без «Конечно!» и прочих вводных
"""

from dataclasses import dataclass, field
from anthropic import Anthropic
from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, MatchValue
from sentence_transformers import SentenceTransformer


@dataclass
class GenerationConfig:
    collection_name: str   = "documents"
    model_name:      str   = "BAAI/bge-small-en-v1.5"
    llm_model:       str   = "claude-sonnet-4-6"
    top_k:           int   = 5
    min_score:       float = 0.45
    max_tokens:      int   = 1024
    temperature:     float = 0.0
    system_prompt:   str   = SYSTEM_PROMPT


@dataclass
class RAGResponse:
    answer:  str
    chunks:  list[RetrievedChunk]
    query:   str
    skipped: bool = False   # True если контекст не прошёл проверку качества


class RAGGenerationPipeline:
    def __init__(self, config: GenerationConfig,
                 qdrant_url: str = "http://localhost:6333"):
        self.cfg    = config
        self.client = Anthropic()
        self.qdrant = QdrantClient(url=qdrant_url)
        self.embed_model = SentenceTransformer(config.model_name)

    # ── Retrieval ───────────────────────────────────────────────────
    def retrieve(
        self,
        query: str,
        filter_by: dict | None = None,
    ) -> list[RetrievedChunk]:
        """
        Выполняет векторный поиск по query.
        filter_by: {"source": "httpx-guide.md"} — метadata-фильтрация
        """
        query_vec = self.embed_model.encode(
            [query], normalize_embeddings=True
        )[0].tolist()

        qdrant_filter = None
        if filter_by:
            qdrant_filter = Filter(must=[
                FieldCondition(key=k, match=MatchValue(value=v))
                for k, v in filter_by.items()
            ])

        hits = self.qdrant.search(
            collection_name=self.cfg.collection_name,
            query_vector=query_vec,
            limit=self.cfg.top_k,
            query_filter=qdrant_filter,
            with_payload=True,
        )

        return [
            RetrievedChunk(
                text=hit.payload["text"],
                source=hit.payload.get("source", "unknown"),
                score=hit.score,
            )
            for hit in hits
        ]

    # ── Generation ──────────────────────────────────────────────────
    def generate(
        self,
        query: str,
        filter_by: dict | None = None,
    ) -> RAGResponse:
        # Шаг 1: Retrieval
        chunks = self.retrieve(query, filter_by=filter_by)

        # Шаг 2: Проверка качества контекста
        if not check_context_quality(chunks, min_score=self.cfg.min_score):
            return RAGResponse(
                answer=NO_CONTEXT_REPLY,
                chunks=chunks,
                query=query,
                skipped=True,
            )

        # Шаг 3: Форматирование контекста
        context = format_context(chunks)
        user_message = f"{context}\n\nВопрос: {query}"

        # Шаг 4: LLM generation
        response = self.client.messages.create(
            model=self.cfg.llm_model,
            max_tokens=self.cfg.max_tokens,
            temperature=self.cfg.temperature,
            system=self.cfg.system_prompt,
            messages=[{"role": "user", "content": user_message}],
        )

        return RAGResponse(
            answer=response.content[0].text,
            chunks=chunks,
            query=query,
        )


# ── Использование ───────────────────────────────────────────────────
pipeline = RAGGenerationPipeline(
    config=GenerationConfig(collection_name="my_docs"),
    qdrant_url="http://localhost:6333",
)

result = pipeline.generate("как настроить timeout в httpx?")
print(result.answer)
print(f"\nИспользовано {len(result.chunks)} чанков")

from anthropic import Anthropic

client = Anthropic()


def generate_streaming(
    query: str,
    chunks: list[RetrievedChunk],
    system_prompt: str,
    model: str = "claude-sonnet-4-6",
):
    """
    Generator: yield токены по одному по мере генерации.
    Используй в FastAPI с StreamingResponse или в CLI.
    """
    if not check_context_quality(chunks):
        yield NO_CONTEXT_REPLY
        return

    context = format_context(chunks)
    user_message = f"{context}\n\nВопрос: {query}"

    with client.messages.stream(
        model=model,
        max_tokens=1024,
        temperature=0,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}],
    ) as stream:
        for text in stream.text_stream:
            yield text


# ── FastAPI StreamingResponse ───────────────────────────────────────
from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()


@app.get("/ask")
async def ask(query: str):
    # Retrieval — отдельно (не стримим)
    chunks = pipeline.retrieve(query)

    # Generation — стрим
    def token_generator():
        for token in generate_streaming(query, chunks, SYSTEM_PROMPT):
            yield token

    return StreamingResponse(token_generator(), media_type="text/plain")


# ── CLI вывод ───────────────────────────────────────────────────────
if __name__ == "__main__":
    query  = "как настроить таймаут в httpx?"
    chunks = pipeline.retrieve(query)
    print("Ответ: ", end="", flush=True)
    for token in generate_streaming(query, chunks, SYSTEM_PROMPT):
        print(token, end="", flush=True)
    print()

from anthropic import Anthropic

client = Anthropic()

REWRITE_PROMPT = """\
Тебе дана история разговора и последний вопрос пользователя.
Переформулируй последний вопрос в самодостаточный, сохранив его смысл.
Ответь только переформулированным вопросом, без пояснений.

История разговора:
{history}

Последний вопрос: {question}"""


def rewrite_question(
    question: str,
    history: list[dict],
) -> str:
    """
    Переформулирует вопрос с учётом истории.
    «А что насчёт AsyncClient?» → «Как настроить таймауты в httpx AsyncClient?»
    """
    if not history:
        return question  # первый вопрос — переформулировка не нужна

    history_text = "\n".join(
        f"{'Пользователь' if m['role']=='user' else 'Ассистент'}: {m['content'][:200]}"
        for m in history[-4:]  # берём последние 2 обмена
    )

    response = client.messages.create(
        model="claude-haiku-4-5-20251001",  # дешёвая модель для rewriting
        max_tokens=128,
        messages=[{
            "role": "user",
            "content": REWRITE_PROMPT.format(
                history=history_text,
                question=question,
            ),
        }],
    )
    return response.content[0].text.strip()


class MultiTurnRAG:
    def __init__(self, pipeline: RAGGenerationPipeline):
        self.pipeline = pipeline
        self.history: list[dict] = []

    def chat(self, user_message: str) -> str:
        # Переформулируем вопрос для поиска
        search_query = rewrite_question(user_message, self.history)

        # Retrieval по переформулированному запросу
        chunks = self.pipeline.retrieve(search_query)

        # Сохраняем оригинальный вопрос в историю
        self.history.append({"role": "user", "content": user_message})

        # Generation с полной историей в messages
        context = format_context(chunks)
        messages = [
            *self.history[:-1],      # история без последнего (уже добавили выше)
            {
                "role": "user",
                "content": f"{context}\n\nВопрос: {user_message}",
            },
        ]

        response = self.pipeline.client.messages.create(
            model=self.pipeline.cfg.llm_model,
            max_tokens=1024,
            temperature=0,
            system=self.pipeline.cfg.system_prompt,
            messages=messages,
        )

        answer = response.content[0].text
        self.history.append({"role": "assistant", "content": answer})
        return answer


# ── Пример диалога ──────────────────────────────────────────────────
rag = MultiTurnRAG(pipeline)

print(rag.chat("как настроить connect timeout в httpx?"))
# → Используйте httpx.Timeout(connect=5.0) ... [1]

print(rag.chat("а как отключить таймаут полностью?"))
# rewrite → «как отключить таймаут в httpx полностью?»
# → Передайте httpx.Timeout(None) ... [2]