Multi-step retrieval: последовательные поиски с уточнением

import json
import asyncio
from openai import OpenAI

client = OpenAI()

REACT_SYSTEM = """You are a research assistant using the ReAct pattern.

For each step:
1. THINK: reason about what information you need next and why
2. SEARCH: call search_docs with a specific, focused query
3. OBSERVE: analyze what you found and what's still missing

Rules:
- Use one focused search per step (not a broad query)
- Refine your query based on what you found previously
- Stop when you have complete information to answer the question
- Maximum 4 search steps

Answer in the same language as the question."""

SEARCH_TOOL = {
    "type": "function",
    "function": {
        "name": "search_docs",
        "description": "Search the knowledge base. Use focused, specific queries. Refine based on previous results.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Specific search query. Be precise, not broad.",
                },
                "reason": {
                    "type": "string",
                    "description": "Why you're searching this — what info you expect to find",
                },
            },
            "required": ["query", "reason"],
        },
    },
}


async def react_retrieve(question: str, retriever, max_steps: int = 4) -> tuple[str, list]:
    """
    ReAct loop: возвращает (ответ, список всех найденных чанков).
    """
    messages = [
        {"role": "system", "content": REACT_SYSTEM},
        {"role": "user", "content": question},
    ]

    all_chunks = []
    step = 0

    while step < max_steps:
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            tools=[SEARCH_TOOL],
            tool_choice="auto",
            temperature=0.1,
        )
        msg = response.choices[0].message

        # Если нет вызовов инструментов — агент решил, что данных достаточно
        if not msg.tool_calls:
            return msg.content, all_chunks

        # Выполняем поиск
        messages.append(msg.model_dump(exclude_unset=True))

        for tc in msg.tool_calls:
            args = json.loads(tc.function.arguments)
            query = args["query"]

            # Поиск с дедупликацией
            new_chunks = retriever.search(query, top_k=4)
            existing_ids = {c.metadata.get("chunk_id") for c in all_chunks}
            novel_chunks = [c for c in new_chunks if c.metadata.get("chunk_id") not in existing_ids]
            all_chunks.extend(novel_chunks)

            # Формируем observation для агента
            observation = _format_observation(novel_chunks, step + 1)
            messages.append({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": observation,
            })

        step += 1

    # Достигли лимита — генерируем ответ с накопленным контекстом
    context = "\n\n---\n\n".join(c.page_content for c in all_chunks)
    messages.append({
        "role": "user",
        "content": f"Maximum steps reached. Generate final answer using context:\n{context}"
    })
    final = client.chat.completions.create(
        model="gpt-4o", messages=messages, temperature=0.1
    )
    return final.choices[0].message.content, all_chunks


def _format_observation(chunks: list, step: int) -> str:
    if not chunks:
        return f"Step {step}: No new relevant chunks found."
    parts = [f"Step {step}: Found {len(chunks)} new chunks:\n"]
    for i, chunk in enumerate(chunks[:3], 1):
        preview = chunk.page_content[:200].replace("\n", " ")
        source = chunk.metadata.get("source", "unknown")
        parts.append(f"[{i}] ({source}) {preview}...")
    return "\n".join(parts)

from dataclasses import dataclass

DECOMPOSE_PROMPT = """Break this question into sequential sub-queries for a knowledge base.

Rules:
- Each sub-query should focus on ONE specific aspect
- If sub-query N depends on result of sub-query M, note it
- Use concise search phrases, not full sentences
- Maximum 4 sub-queries
- Order them so dependencies come first

Question: {question}

Return JSON array:
[
  {{"step": 1, "query": "...", "purpose": "what info this gets", "depends_on": null}},
  {{"step": 2, "query": "...", "purpose": "...", "depends_on": 1}}
]"""


@dataclass
class SubQuery:
    step: int
    query: str
    purpose: str
    depends_on: int | None = None


def decompose_question(question: str) -> list[SubQuery]:
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": DECOMPOSE_PROMPT.format(question=question)}],
        response_format={"type": "json_object"},
        temperature=0,
    )
    data = json.loads(resp.choices[0].message.content)
    items = data if isinstance(data, list) else list(data.values())[0]
    return [SubQuery(**item) for item in items]


ENRICH_PROMPT = """Enrich this search query using the context from previous search steps.
Replace placeholders like {{answer_from_step_N}} with actual data.

Previous results:
{prev_results}

Query to enrich: {query}

Return the enriched search query (just the query string, nothing else):"""


def enrich_query(query: str, prev_results: list[str]) -> str:
    """Подставляет результаты предыдущих шагов в запрос следующего."""
    if "{{" not in query or not prev_results:
        return query  # Нет плейсхолдеров — возвращаем как есть

    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": ENRICH_PROMPT.format(
            prev_results="\n\n".join(f"Step {i+1}: {r}" for i, r in enumerate(prev_results)),
            query=query,
        )}],
        temperature=0,
        max_tokens=60,
    )
    return resp.choices[0].message.content.strip()


async def decomposed_retrieve(question: str, retriever) -> tuple[str, list]:
    """Декомпозирует вопрос и выполняет поиски последовательно."""
    subqueries = decompose_question(question)
    all_chunks = []
    step_results: list[str] = []  # Краткие summary каждого шага для enrich_query

    for sq in sorted(subqueries, key=lambda x: x.step):
        # Обогащаем запрос результатами предыдущих шагов
        final_query = enrich_query(sq.query, step_results)

        # Поиск
        chunks = retriever.search(final_query, top_k=4)
        new_chunks = _deduplicate(chunks, all_chunks)
        all_chunks.extend(new_chunks)

        # Краткий summary для следующих шагов
        if new_chunks:
            summary = " | ".join(c.page_content[:100] for c in new_chunks[:2])
            step_results.append(summary)

    # Генерируем финальный ответ
    context = "\n\n---\n\n".join(c.page_content for c in all_chunks)
    resp = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "Answer using the provided context."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"},
        ],
        temperature=0.1,
    )
    return resp.choices[0].message.content, all_chunks


def _deduplicate(new_chunks: list, existing: list) -> list:
    existing_ids = {c.metadata.get("chunk_id", c.page_content[:50]) for c in existing}
    return [c for c in new_chunks if c.metadata.get("chunk_id", c.page_content[:50]) not in existing_ids]

GAP_ANALYSIS_PROMPT = """Analyze if the context is sufficient to fully answer the question.

Question: {question}

Current context ({n_chunks} chunks):
{context_preview}

Evaluate:
1. What aspects of the question ARE covered?
2. What specific information is MISSING?
3. What search query would best find the missing info?
4. Sufficiency score 0-100 (0=useless, 50=partial, 100=complete)

Return JSON:
{{
  "score": int,
  "covered": ["aspect1", "aspect2"],
  "missing": ["missing1", "missing2"],
  "followup_query": "specific search query",
  "sufficient": bool  // true if score >= 75
}}"""


def analyze_gaps(question: str, chunks: list) -> dict:
    """Анализирует что нашли и что ещё нужно."""
    context_preview = "\n\n".join(c.page_content[:300] for c in chunks[:5])

    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": GAP_ANALYSIS_PROMPT.format(
            question=question,
            n_chunks=len(chunks),
            context_preview=context_preview,
        )}],
        response_format={"type": "json_object"},
        temperature=0,
    )
    return json.loads(resp.choices[0].message.content)


async def gap_retrieve(question: str, retriever, max_steps: int = 4) -> tuple[str, list]:
    """Итеративный поиск с gap analysis после каждого шага."""
    all_chunks = []
    current_query = question
    history: list[dict] = []  # История шагов для отладки

    for step in range(max_steps):
        # Поиск
        new_chunks = retriever.search(current_query, top_k=4)
        all_chunks.extend(_deduplicate(new_chunks, all_chunks))

        # Gap analysis
        analysis = analyze_gaps(question, all_chunks)
        history.append({
            "step": step + 1,
            "query": current_query,
            "found": len(new_chunks),
            "score": analysis["score"],
            "missing": analysis["missing"],
        })

        # Достаточно?
        if analysis["sufficient"] or step == max_steps - 1:
            break

        # Уточняем запрос
        current_query = analysis["followup_query"]

    # Генерируем ответ
    context = "\n\n---\n\n".join(c.page_content for c in all_chunks)
    resp = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "Answer the question using the provided context. Be specific."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"},
        ],
        temperature=0.1,
    )
    return resp.choices[0].message.content, all_chunks, history

from dataclasses import dataclass, field


@dataclass
class StopCondition:
    max_steps: int = 4
    min_score: float = 0.75    # Gap analysis score threshold
    min_novel_chunks: int = 1  # Минимум новых чанков за шаг
    max_overlap: float = 0.8   # Максимальный overlap с предыдущими результатами


def check_overlap(new_chunks: list, existing: list) -> float:
    """Доля новых чанков, которые уже есть в existing."""
    if not new_chunks:
        return 1.0
    existing_ids = {c.metadata.get("chunk_id", c.page_content[:50]) for c in existing}
    duplicates = sum(
        1 for c in new_chunks
        if c.metadata.get("chunk_id", c.page_content[:50]) in existing_ids
    )
    return duplicates / len(new_chunks)


def should_stop(
    step: int,
    new_chunks: list,
    all_chunks: list,
    analysis: dict,
    cond: StopCondition,
) -> tuple[bool, str]:
    """Возвращает (нужно_ли_остановиться, причина)."""
    # 1. Достигли лимита шагов
    if step >= cond.max_steps:
        return True, f"max_steps={cond.max_steps} reached"

    # 2. LLM считает контекст достаточным
    if analysis.get("score", 0) >= cond.min_score * 100:
        return True, f"sufficient: score={analysis['score']}"

    # 3. Ничего нового не нашли
    novel = [c for c in new_chunks
             if c.metadata.get("chunk_id", c.page_content[:50]) not in
             {e.metadata.get("chunk_id", e.page_content[:50]) for e in all_chunks}]
    if len(novel) < cond.min_novel_chunks:
        return True, "no new relevant chunks found"

    # 4. Ищем по кругу (высокий overlap)
    overlap = check_overlap(new_chunks, all_chunks[:-len(new_chunks)])
    if overlap >= cond.max_overlap:
        return True, f"circular search: overlap={overlap:.0%}"

    return False, ""


# Использование в main loop
stop_cond = StopCondition(max_steps=4, min_score=0.75, min_novel_chunks=1)

for step in range(stop_cond.max_steps):
    new_chunks = retriever.search(current_query)
    all_chunks.extend(new_chunks)
    analysis = analyze_gaps(question, all_chunks)

    stop, reason = should_stop(step + 1, new_chunks, all_chunks, analysis, stop_cond)
    if stop:
        print(f"Stopped at step {step+1}: {reason}")
        break

    current_query = analysis["followup_query"]

from enum import StrEnum
from dataclasses import dataclass, field
from typing import Protocol


class Strategy(StrEnum):
    REACT      = "react"       # Reason-Act-Observe loop
    DECOMPOSE  = "decompose"   # Предварительная декомпозиция
    GAP        = "gap"         # Реактивный gap analysis
    AUTO       = "auto"        # Выбор стратегии автоматически


class Retriever(Protocol):
    def search(self, query: str, top_k: int = 4) -> list: ...


@dataclass
class MultiStepConfig:
    strategy: Strategy = Strategy.AUTO
    max_steps: int = 4
    top_k: int = 4
    score_threshold: float = 0.75
    model: str = "gpt-4o-mini"
    reasoning_model: str = "gpt-4o"  # Для финального ответа


@dataclass
class RetrievalResult:
    answer: str
    chunks: list
    steps: int
    strategy: Strategy
    stop_reason: str = ""


class MultiStepRAG:
    def __init__(self, retriever: Retriever, config: MultiStepConfig = None):
        self.retriever = retriever
        self.config = config or MultiStepConfig()

    async def answer(self, question: str, strategy: Strategy = None) -> RetrievalResult:
        strategy = strategy or self.config.strategy

        if strategy == Strategy.AUTO:
            strategy = self._choose_strategy(question)

        if strategy == Strategy.REACT:
            text, chunks = await react_retrieve(
                question, self.retriever, max_steps=self.config.max_steps
            )
        elif strategy == Strategy.DECOMPOSE:
            text, chunks = await decomposed_retrieve(question, self.retriever)
        else:  # GAP (default)
            text, chunks, history = await gap_retrieve(
                question, self.retriever, max_steps=self.config.max_steps
            )

        return RetrievalResult(
            answer=text,
            chunks=chunks,
            steps=len(chunks) // self.config.top_k,
            strategy=strategy,
        )

    def _choose_strategy(self, question: str) -> Strategy:
        """Эвристика выбора стратегии без LLM-вызова."""
        q = question.lower()

        # Явно составной вопрос — decompose
        COMPOSITE_MARKERS = ["и ", " а также", "плюс ", "кроме того", "помимо"]
        QUESTION_WORDS = ["чем отличается", "compare", "vs ", "versus"]
        if any(m in q for m in COMPOSITE_MARKERS + QUESTION_WORDS):
            return Strategy.DECOMPOSE

        # Исследовательский / открытый вопрос — ReAct
        OPEN_MARKERS = ["расскажи подробно", "объясни как", "how does", "explain"]
        if any(m in q for m in OPEN_MARKERS) or len(question.split()) > 20:
            return Strategy.REACT

        # По умолчанию — gap analysis (хорошо для технических вопросов)
        return Strategy.GAP


# Использование
async def main():
    rag = MultiStepRAG(
        retriever=my_vector_db,
        config=MultiStepConfig(max_steps=4, score_threshold=0.75),
    )

    questions = [
        "Как настроить Redis timeout?",                           # GAP (1-2 steps)
        "Чем Enterprise отличается от Business и какие скидки?",  # DECOMPOSE
        "Расскажи подробно про архитектуру нашей системы кэширования",  # REACT
    ]

    for q in questions:
        result = await rag.answer(q)
        print(f"Q: {q}")
        print(f"  Strategy: {result.strategy}, Steps: {result.steps}")
        print(f"  Chunks: {len(result.chunks)}")
        print(f"  Answer: {result.answer[:100]}...")
        print()

all_chunks, query = [], question
for step in range(4):
    new = retriever.search(query, top_k=4)
    all_chunks.extend(_deduplicate(new, all_chunks))
    analysis = analyze_gaps(question, all_chunks)
    if analysis["sufficient"]: break
    query = analysis["followup_query"]
answer = generate(question, all_chunks)