Self-RAG: агент оценивает релевантность и при необходимости ищет ещё

from openai import AsyncOpenAI
from pydantic import BaseModel
from enum import StrEnum

client = AsyncOpenAI()

RETRIEVE_SYSTEM = """Ты классификатор: нужен ли поиск по базе знаний для ответа на вопрос.

Retrieval НЕ нужен если:
- вопрос требует только общих знаний LLM (объяснение концепций, написание кода)
- вопрос математический или логический
- вопрос субъективный (советы, мнения)

Retrieval НУЖЕН если:
- вопрос про конкретные данные из корпуса (версии, лимиты, конфигурации)
- вопрос про события, документы, изменения в системе
- вопрос про специфику продукта/компании/проекта

Отвечай строго: "yes" или "no"."""


class RetrieveDecision(StrEnum):
    YES      = "yes"
    NO       = "no"
    CONTINUE = "continue"   # для посегментной генерации


class RetrieveResult(BaseModel):
    decision: RetrieveDecision
    reason: str


async def check_retrieve_needed(question: str) -> RetrieveResult:
    """[Retrieve] token: определяем нужен ли retrieval."""
    resp = await client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": RETRIEVE_SYSTEM},
            {"role": "user", "content": f"Вопрос: {question}"},
        ],
        response_format=RetrieveResult,
        temperature=0,
    )
    return resp.choices[0].message.parsed


# Примеры использования
async def demo_retrieve_check():
    cases = [
        "Что такое JWT?",                          # no
        "Какой лимит на запросы к нашему API?",    # yes
        "Как работает TCP handshake?",             # no
        "Что изменилось в версии 2.4.1?",          # yes
    ]
    for q in cases:
        result = await check_retrieve_needed(q)
        print(f"[{result.decision.upper():3}] {q}")
        # [NO ] Что такое JWT?
        # [YES] Какой лимит на запросы к нашему API?
        # [NO ] Как работает TCP handshake?
        # [YES] Что изменилось в версии 2.4.1?

from langchain_core.documents import Document

ISREL_SYSTEM = """Ты оцениваешь, является ли документ релевантным для ответа на вопрос.

Документ РЕЛЕВАНТЕН если:
- содержит прямую информацию для ответа на вопрос
- содержит связанный контекст, который поможет ответить

Документ НЕРЕЛЕВАНТЕН если:
- тема документа другая, хотя встречаются похожие термины
- документ затрагивает тему косвенно, но не помогает ответить на конкретный вопрос

Отвечай строго: "relevant" или "irrelevant"."""


class IsRelResult(BaseModel):
    verdict: str   # "relevant" | "irrelevant"
    reason: str


async def check_relevance(question: str, chunk: Document) -> IsRelResult:
    """[ISREL] token: проверяем релевантность чанка вопросу."""
    resp = await client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": ISREL_SYSTEM},
            {"role": "user", "content": (
                f"Вопрос: {question}\n\n"
                f"Документ:\n{chunk.page_content}"
            )},
        ],
        response_format=IsRelResult,
        temperature=0,
    )
    return resp.choices[0].message.parsed


async def filter_relevant_chunks(
    question: str,
    chunks: list[Document],
    concurrency: int = 4,
) -> list[Document]:
    """Параллельная фильтрация: оставляем только релевантные чанки."""
    import asyncio

    semaphore = asyncio.Semaphore(concurrency)

    async def check_one(chunk: Document) -> tuple[Document, IsRelResult]:
        async with semaphore:
            result = await check_relevance(question, chunk)
            return chunk, result

    results = await asyncio.gather(*[check_one(c) for c in chunks])

    relevant = [chunk for chunk, r in results if r.verdict == "relevant"]
    print(f"ISREL: {len(relevant)}/{len(chunks)} чанков релевантны")
    return relevant

from pydantic import BaseModel, Field

class ChunkVerdict(BaseModel):
    chunk_index: int
    verdict: str    # "relevant" | "irrelevant"
    reason: str

class BatchIsRelResult(BaseModel):
    verdicts: list[ChunkVerdict]


async def filter_relevant_batch(
    question: str,
    chunks: list[Document],
) -> list[Document]:
    """Один LLM-вызов для оценки всех чанков сразу."""
    chunks_text = "\n\n".join(
        f"[Чанк {i}]\n{c.page_content}" for i, c in enumerate(chunks)
    )

    resp = await client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": (
                "Оцени каждый из предоставленных чанков: "
                "релевантен ли он для ответа на вопрос. "
                "Верни verdict для каждого чанка по его индексу."
            )},
            {"role": "user", "content": (
                f"Вопрос: {question}\n\n{chunks_text}"
            )},
        ],
        response_format=BatchIsRelResult,
        temperature=0,
    )
    result = resp.choices[0].message.parsed

    relevant_indices = {
        v.chunk_index for v in result.verdicts
        if v.verdict == "relevant"
    }
    return [c for i, c in enumerate(chunks) if i in relevant_indices]

from enum import StrEnum

class GroundingVerdict(StrEnum):
    FULLY    = "fully_supported"
    PARTIAL  = "partially_supported"
    NONE     = "no_support"

    @property
    def weight(self) -> float:
        return {
            GroundingVerdict.FULLY:    1.0,
            GroundingVerdict.PARTIAL:  0.5,
            GroundingVerdict.NONE:     0.0,
        }[self]


ISGRT_SYSTEM = """Ты оцениваешь, насколько сгенерированный ответ основан на предоставленных документах.

Критерии:
- "fully_supported": каждое утверждение в ответе подтверждается документами
- "partially_supported": часть утверждений подтверждается, часть добавлена из знаний модели
- "no_support": ответ не подкреплён документами или противоречит им"""


class IsGrtResult(BaseModel):
    verdict: GroundingVerdict
    unsupported_claims: list[str] = Field(
        default_factory=list,
        description="Список утверждений без подтверждения в документах"
    )
    reason: str


async def check_grounding(
    question: str,
    answer: str,
    context_chunks: list[Document],
) -> IsGrtResult:
    """[ISGRT] token: насколько ответ заземлён в документах?"""
    context = "\n\n".join(
        f"[Документ {i+1}]\n{c.page_content}"
        for i, c in enumerate(context_chunks)
    )

    resp = await client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": ISGRT_SYSTEM},
            {"role": "user", "content": (
                f"Вопрос: {question}\n\n"
                f"Документы:\n{context}\n\n"
                f"Ответ:\n{answer}"
            )},
        ],
        response_format=IsGrtResult,
        temperature=0,
    )
    return resp.choices[0].message.parsed

ISUSE_SYSTEM = """Оцени насколько ответ полезен для пользователя, задавшего вопрос.

Шкала:
5 — Полный, точный, конкретный ответ. Решает задачу пользователя.
4 — Хороший ответ, возможно с незначительными пропусками.
3 — Частично помогает, но неполный или нечёткий.
2 — Поверхностный. Касается темы, но не решает задачу.
1 — Нерелевантный, неверный или бессодержательный ответ.

Возвращай только число от 1 до 5."""


class IsUseResult(BaseModel):
    score: int        # 1-5
    reason: str


async def check_utility(question: str, answer: str) -> IsUseResult:
    """[ISUSE] token: насколько ответ полезен?"""
    resp = await client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": ISUSE_SYSTEM},
            {"role": "user", "content": (
                f"Вопрос: {question}\n\nОтвет:\n{answer}"
            )},
        ],
        response_format=IsUseResult,
        temperature=0,
    )
    return resp.choices[0].message.parsed

from dataclasses import dataclass, field

GENERATE_SYSTEM = """Ты — ассистент, отвечающий на вопросы строго на основе предоставленных документов.
Если документы не содержат необходимой информации, скажи об этом прямо.
Не добавляй информацию из своих знаний."""


@dataclass
class Candidate:
    """Один кандидат-ответ, связанный с конкретным чанком."""
    chunk: Document
    answer: str
    isrel: IsRelResult   | None = None
    isgrt: IsGrtResult   | None = None
    isuse: IsUseResult   | None = None
    score: float = 0.0   # итоговый балл

    @property
    def is_grounded(self) -> bool:
        return self.isgrt is not None and self.isgrt.verdict != GroundingVerdict.NONE


async def generate_candidate(
    question: str,
    chunk: Document,
    model: str = "gpt-4o-mini",
) -> str:
    """Генерируем ответ на основе одного чанка."""
    resp = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": GENERATE_SYSTEM},
            {"role": "user", "content": (
                f"Документ:\n{chunk.page_content}\n\n"
                f"Вопрос: {question}"
            )},
        ],
        temperature=0.1,
    )
    return resp.choices[0].message.content


async def evaluate_candidate(
    question: str,
    candidate: Candidate,
) -> Candidate:
    """ISGRT + ISUSE для одного кандидата, параллельно."""
    isgrt, isuse = await asyncio.gather(
        check_grounding(question, candidate.answer, [candidate.chunk]),
        check_utility(question, candidate.answer),
    )
    candidate.isgrt = isgrt
    candidate.isuse = isuse
    candidate.score = isuse.score * isgrt.verdict.weight
    return candidate

import asyncio
from typing import Callable

# Type: retriever принимает строку, возвращает список Document
Retriever = Callable[[str], list[Document]]


@dataclass
class SelfRAGResult:
    answer: str
    used_retrieval: bool
    winner_score: float
    candidates: list[Candidate] = field(default_factory=list)
    retrieve_reason: str = ""


async def self_rag(
    question: str,
    retriever: Retriever,
    top_k: int = 4,
    min_score: float = 2.0,    # порог: ниже → fallback на direct
    model: str = "gpt-4o-mini",
) -> SelfRAGResult:
    """
    Полный Self-RAG пайплайн:
    1. [Retrieve] — нужен ли поиск?
    2. [ISREL]    — фильтруем нерелевантные чанки
    3. Generate   — k кандидатов из релевантных чанков
    4. [ISGRT]    — заземлённость каждого кандидата
    5. [ISUSE]    — полезность каждого кандидата
    6. Rank       — выбираем лучший
    """

    # ── Шаг 1: [Retrieve] ──────────────────────────────────────────────
    retrieve_result = await check_retrieve_needed(question)

    if retrieve_result.decision == RetrieveDecision.NO:
        # Прямая генерация без retrieval
        direct_answer = await generate_direct(question, model)
        isuse = await check_utility(question, direct_answer)
        return SelfRAGResult(
            answer=direct_answer,
            used_retrieval=False,
            winner_score=float(isuse.score),
            retrieve_reason=retrieve_result.reason,
        )

    # ── Шаг 2: Retrieval + [ISREL] ─────────────────────────────────────
    raw_chunks = retriever(question)
    relevant_chunks = await filter_relevant_batch(question, raw_chunks[:top_k])

    if not relevant_chunks:
        # Не нашли ничего релевантного — fallback на direct
        direct_answer = await generate_direct(question, model)
        isuse = await check_utility(question, direct_answer)
        return SelfRAGResult(
            answer=direct_answer,
            used_retrieval=True,
            winner_score=float(isuse.score),
            retrieve_reason="no relevant chunks found",
        )

    # ── Шаг 3: Генерация кандидатов (параллельно) ──────────────────────
    answers = await asyncio.gather(
        *[generate_candidate(question, chunk, model) for chunk in relevant_chunks]
    )
    candidates = [
        Candidate(chunk=chunk, answer=answer)
        for chunk, answer in zip(relevant_chunks, answers)
    ]

    # ── Шаг 4+5: [ISGRT] + [ISUSE] для каждого кандидата (параллельно) ─
    evaluated = await asyncio.gather(
        *[evaluate_candidate(question, c) for c in candidates]
    )

    # ── Шаг 6: Rank — выбираем лучший кандидат ──────────────────────────
    ranked = sorted(evaluated, key=lambda c: c.score, reverse=True)
    winner = ranked[0]

    if winner.score < min_score:
        # Ни один кандидат не набрал достаточного балла → direct fallback
        direct_answer = await generate_direct(question, model)
        isuse_d = await check_utility(question, direct_answer)
        if isuse_d.score > winner.score:
            return SelfRAGResult(
                answer=direct_answer,
                used_retrieval=True,
                winner_score=float(isuse_d.score),
                candidates=list(ranked),
            )

    return SelfRAGResult(
        answer=winner.answer,
        used_retrieval=True,
        winner_score=winner.score,
        candidates=list(ranked),
    )


async def generate_direct(question: str, model: str) -> str:
    """Генерация без retrieval-контекста."""
    resp = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "Ты полезный ассистент. Отвечай точно и лаконично."},
            {"role": "user", "content": question},
        ],
        temperature=0.1,
    )
    return resp.choices[0].message.content


# Пример использования
async def main():
    result = await self_rag(
        question="Как настроить rate limiting в Nginx?",
        retriever=lambda q: vector_db.similarity_search(q, k=4),
        top_k=4,
        min_score=2.0,
    )
    print(f"Used retrieval: {result.used_retrieval}")
    print(f"Winner score: {result.winner_score:.1f}")
    print(f"Candidates evaluated: {len(result.candidates)}")
    for c in result.candidates:
        print(f"  score={c.score:.1f} isgrt={c.isgrt.verdict} isuse={c.isuse.score}")
    print(f"\nAnswer:\n{result.answer}")

import json, logging
from datetime import datetime

logger = logging.getLogger("self_rag")


def log_self_rag_trace(question: str, result: SelfRAGResult) -> dict:
    """Структурированный лог трассировки Self-RAG."""
    trace = {
        "ts": datetime.utcnow().isoformat(),
        "question": question,
        "retrieve": result.used_retrieval,
        "retrieve_reason": result.retrieve_reason,
        "candidates": [],
        "winner_score": result.winner_score,
    }

    for c in result.candidates:
        trace["candidates"].append({
            "chunk_id": c.chunk.metadata.get("chunk_id", "?"),
            "isrel": "relevant",                            # прошёл фильтр
            "isgrt": c.isgrt.verdict if c.isgrt else None,
            "isgrt_unsupported": c.isgrt.unsupported_claims if c.isgrt else [],
            "isuse": c.isuse.score if c.isuse else None,
            "score": c.score,
        })

    logger.info(json.dumps(trace, ensure_ascii=False))
    return trace


# Мониторинг: считаем процент no_support по chunk_id
# Если конкретный чанк систематически даёт no_support → плохой контент
from collections import Counter

def analyze_grounding_issues(traces: list[dict]) -> None:
    no_support_chunks = Counter()
    for t in traces:
        for c in t["candidates"]:
            if c["isgrt"] == "no_support":
                no_support_chunks[c["chunk_id"]] += 1

    print("Чанки с частыми галлюцинациями (no_support):")
    for chunk_id, count in no_support_chunks.most_common(5):
        print(f"  {chunk_id}: {count} раз")

score = isuse.score * isgrt.verdict.weight
# fully_supported  → weight=1.0  → score=ISUSE
# partially_supported → weight=0.5 → score=ISUSE*0.5
# no_support       → weight=0.0  → score=0 (галлюцинация, не возвращается)

chunks  = retriever(question)
relevant = await filter_relevant_batch(question, chunks)
if not relevant:
    return await generate_direct(question)

# k кандидатов из релевантных чанков
answers = await asyncio.gather(
    *[generate_candidate(question, c) for c in relevant]
)
# выбираем заземлённый ответ
for answer, chunk in zip(answers, relevant):
    grounding = await check_grounding(question, answer, [chunk])
    if grounding.verdict == GroundingVerdict.FULLY:
        return answer         # первый полностью заземлённый → возвращаем

return answers[0]             # fallback: первый кандидат