Rethinking with Retrieval (RwR)

from openai import OpenAI
from dataclasses import dataclass

client = OpenAI()


@dataclass
class VerificationResult:
    """Результат верификации одного шага."""
    original_step: str
    evidence: list[str]
    verdict: str  # "confirmed", "corrected", "unverified"
    corrected_step: str


COT_PROMPT = """Разбей ответ на вопрос на отдельные шаги рассуждения.
Каждый шаг должен содержать одно конкретное, проверяемое утверждение.

Вопрос: {question}

Формат:
Step 1: ...
Step 2: ...
Step N: ..."""

QUERY_PROMPT = """Сформулируй поисковый запрос для проверки этого утверждения.
Запрос должен быть коротким, конкретным и содержать ключевые сущности.

Утверждение: {step}

Верни только поисковый запрос, ничего больше."""

VERIFY_PROMPT = """Сопоставь утверждение с найденным evidence.

Утверждение: {step}

Evidence:
{evidence}

Предыдущие верифицированные шаги:
{previous_steps}

Определи:
1. Утверждение ПОДТВЕРЖДЕНО, ОПРОВЕРГНУТО или НЕЛЬЗЯ ПРОВЕРИТЬ?
2. Если опровергнуто — дай исправленную версию.

Формат ответа:
VERDICT: confirmed | corrected | unverified
CORRECTED: <исправленное утверждение или "N/A">"""

FINAL_PROMPT = """Сформулируй финальный ответ на основе верифицированных шагов.

Вопрос: {question}

Верифицированные шаги:
{steps}

Дай точный, фактологически верный ответ."""


def generate_cot_steps(question: str) -> list[str]:
    """Шаг 1: генерация CoT."""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": COT_PROMPT.format(
            question=question
        )}],
        temperature=0,
    )
    lines = response.choices[0].message.content.strip().split("\n")
    steps = []
    for line in lines:
        line = line.strip()
        if line.lower().startswith("step"):
            step = line.split(":", 1)[1].strip() if ":" in line else line
            steps.append(step)
    return steps


def formulate_query(step: str) -> str:
    """Шаг 2: формирование поискового запроса."""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": QUERY_PROMPT.format(
            step=step
        )}],
        temperature=0,
    )
    return response.choices[0].message.content.strip()


def retrieve_evidence(query: str, retriever) -> list[str]:
    """Шаг 3: retrieval из базы знаний.

    retriever — любая функция str -> list[str]:
    - BM25 по Wikipedia
    - Vector search (Qdrant, Chroma)
    - Web search API (Tavily, Exa)
    """
    return retriever(query)


def verify_step(
    step: str,
    evidence: list[str],
    previous_steps: list[str]
) -> VerificationResult:
    """Шаг 4: верификация с учётом evidence."""
    evidence_text = "\n---\n".join(evidence) if evidence else "(нет evidence)"
    prev_text = "\n".join(
        f"  {i+1}. {s}" for i, s in enumerate(previous_steps)
    ) or "(нет)"

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": VERIFY_PROMPT.format(
            step=step,
            evidence=evidence_text,
            previous_steps=prev_text,
        )}],
        temperature=0,
    )

    content = response.choices[0].message.content.strip()

    # Парсинг ответа
    verdict = "unverified"
    corrected = step
    for line in content.split("\n"):
        line = line.strip()
        if line.startswith("VERDICT:"):
            v = line.split(":", 1)[1].strip().lower()
            if v in ("confirmed", "corrected", "unverified"):
                verdict = v
        elif line.startswith("CORRECTED:"):
            c = line.split(":", 1)[1].strip()
            if c != "N/A":
                corrected = c

    return VerificationResult(
        original_step=step,
        evidence=evidence,
        verdict=verdict,
        corrected_step=corrected if verdict == "corrected" else step,
    )


def rethinking_with_retrieval(
    question: str,
    retriever,  # Callable[[str], list[str]]
    top_k: int = 3,
) -> dict:
    """Полный пайплайн Rethinking with Retrieval."""

    # 1. Генерация CoT
    steps = generate_cot_steps(question)

    # 2-4. Для каждого шага: query -> retrieve -> verify
    verified_steps: list[str] = []
    verification_log: list[VerificationResult] = []

    for step in steps:
        query = formulate_query(step)
        evidence = retrieve_evidence(query, retriever)[:top_k]
        result = verify_step(step, evidence, verified_steps)
        verification_log.append(result)
        verified_steps.append(result.corrected_step)

    # 5. Финальный ответ
    steps_text = "\n".join(
        f"  {i+1}. [{v.verdict}] {v.corrected_step}"
        for i, v in enumerate(verification_log)
    )
    final = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": FINAL_PROMPT.format(
            question=question,
            steps=steps_text,
        )}],
        temperature=0,
    )

    return {
        "answer": final.choices[0].message.content.strip(),
        "steps": [
            {
                "original": v.original_step,
                "corrected": v.corrected_step,
                "verdict": v.verdict,
                "evidence_count": len(v.evidence),
            }
            for v in verification_log
        ],
        "total_steps": len(steps),
        "corrected_count": sum(
            1 for v in verification_log if v.verdict == "corrected"
        ),
    }


# — Пример с BM25 retriever ---
# pip install rank-bm25

from rank_bm25 import BM25Okapi


def build_bm25_retriever(corpus: list[str]):
    """Простой BM25 retriever по корпусу текстов."""
    tokenized = [doc.lower().split() for doc in corpus]
    bm25 = BM25Okapi(tokenized)

    def search(query: str, k: int = 3) -> list[str]:
        scores = bm25.get_scores(query.lower().split())
        top_indices = sorted(
            range(len(scores)), key=lambda i: scores[i], reverse=True
        )[:k]
        return [corpus[i] for i in top_indices if scores[i] > 0]

    return search


# Пример использования
knowledge_base = [
    "Санкт-Петербург основан 27 мая 1703 года Петром I.",
    "Санкт-Петербург был столицей Российской империи с 1712 по 1918 год.",
    "Население Санкт-Петербурга — около 5.6 млн человек (2024).",
    "Эрмитаж — один из крупнейших музеев мира, основан в 1764 году.",
    "Московский Кремль построен в XV веке при Иване III.",
]

retriever = build_bm25_retriever(knowledge_base)
result = rethinking_with_retrieval(
    question="Когда был основан Санкт-Петербург и кем?",
    retriever=retriever,
)
print(result["answer"])
print(f"Исправлено шагов: {result['corrected_count']}/{result['total_steps']}")

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

qdrant = QdrantClient(":memory:")  # или url="http://localhost:6333"


def build_vector_retriever(
    documents: list[str],
    collection_name: str = "knowledge",
):
    """Vector retriever через Qdrant + OpenAI embeddings."""

    # Создание эмбеддингов
    embed_response = client.embeddings.create(
        model="text-embedding-3-small",
        input=documents,
    )
    vectors = [e.embedding for e in embed_response.data]
    dim = len(vectors[0])

    # Создание коллекции
    qdrant.create_collection(
        collection_name=collection_name,
        vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
    )

    # Загрузка документов
    points = [
        PointStruct(id=i, vector=v, payload={"text": doc})
        for i, (v, doc) in enumerate(zip(vectors, documents))
    ]
    qdrant.upsert(collection_name=collection_name, points=points)

    def search(query: str, k: int = 3) -> list[str]:
        q_embed = client.embeddings.create(
            model="text-embedding-3-small",
            input=[query],
        ).data[0].embedding
        results = qdrant.search(
            collection_name=collection_name,
            query_vector=q_embed,
            limit=k,
        )
        return [hit.payload["text"] for hit in results]

    return search