Retrieval-Augmented Thoughts (RAT)

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

llm = ChatOpenAI(model="gpt-4o", temperature=0)
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")


def build_retriever(documents: list[str], k: int = 3):
    """Создать retriever из списка текстовых документов."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500, chunk_overlap=50
    )
    docs = [Document(page_content=d) for d in documents]
    chunks = splitter.split_documents(docs)
    vectorstore = Chroma.from_documents(chunks, embeddings)
    return vectorstore.as_retriever(search_kwargs={"k": k})


INITIAL_COT_PROMPT = """Разбей задачу на шаги рассуждения.

Задача: {question}

Верни список шагов в формате:
Thought 1: ...
Thought 2: ...
Thought N: ...

Каждый шаг --- отдельная мысль, конкретная и проверяемая."""

REVISION_PROMPT = """Пересмотри мысль с учётом найденной информации.

Задача: {question}

Предыдущие (уже пересмотренные) мысли:
{previous_thoughts}

Текущая мысль: {thought}

Найденная информация:
{retrieved_context}

Верни ИСПРАВЛЕННУЮ мысль. Если мысль корректна --- верни как есть.
Если найдена ошибка --- исправь и укажи, что изменилось."""

FINAL_ANSWER_PROMPT = """Сформулируй финальный ответ на основе пересмотренных мыслей.

Задача: {question}

Пересмотренные мысли:
{revised_thoughts}

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


class RATPipeline:
    """Retrieval-Augmented Thoughts: retrieval на каждом шаге CoT."""

    def __init__(self, retriever, llm_instance=None):
        self.retriever = retriever
        self.llm = llm_instance or llm

    def _generate_initial_thoughts(self, question: str) -> list[str]:
        """Шаг 1: генерация начальной цепочки мыслей."""
        response = self.llm.invoke(
            INITIAL_COT_PROMPT.format(question=question)
        )
        lines = response.content.strip().split("\n")
        thoughts = []
        for line in lines:
            line = line.strip()
            if line.startswith("Thought"):
                # Убираем префикс "Thought N: "
                thought = line.split(":", 1)[1].strip() if ":" in line else line
                thoughts.append(thought)
        return thoughts

    def _retrieve(self, query: str) -> str:
        """Шаг 2: поиск по базе знаний."""
        docs = self.retriever.invoke(query)
        return "\n---\n".join(doc.page_content for doc in docs)

    def _revise_thought(
        self,
        question: str,
        thought: str,
        context: str,
        previous: list[str]
    ) -> str:
        """Шаг 3: ревизия мысли с учётом retrieval."""
        prev_text = "\n".join(
            f"  {i+1}. {t}" for i, t in enumerate(previous)
        ) or "(нет)"
        response = self.llm.invoke(
            REVISION_PROMPT.format(
                question=question,
                previous_thoughts=prev_text,
                thought=thought,
                retrieved_context=context
            )
        )
        return response.content.strip()

    def run(self, question: str) -> dict:
        """Полный цикл RAT."""
        # 1. Начальные мысли
        thoughts = self._generate_initial_thoughts(question)

        # 2-3. Для каждой мысли: retrieval + revision
        revised = []
        retrieval_log = []

        for i, thought in enumerate(thoughts):
            context = self._retrieve(thought)
            retrieval_log.append({
                "thought": thought,
                "retrieved": context[:200] + "..."
            })
            revised_thought = self._revise_thought(
                question, thought, context, revised
            )
            revised.append(revised_thought)

        # 4. Финальный ответ
        revised_text = "\n".join(
            f"  {i+1}. {t}" for i, t in enumerate(revised)
        )
        final = self.llm.invoke(
            FINAL_ANSWER_PROMPT.format(
                question=question,
                revised_thoughts=revised_text
            )
        )

        return {
            "answer": final.content.strip(),
            "initial_thoughts": thoughts,
            "revised_thoughts": revised,
            "retrieval_log": retrieval_log,
            "num_retrievals": len(thoughts)
        }


# Пример использования
knowledge_base = [
    "Python 3.12 выпущен в октябре 2023 года.",
    "В Python 3.12 добавлены type parameter syntax (PEP 695).",
    "asyncio.TaskGroup появился в Python 3.11.",
    "match/case (structural pattern matching) добавлен в Python 3.10.",
    "GIL в Python планируется сделать опциональным в Python 3.13 (PEP 703).",
]

retriever = build_retriever(knowledge_base, k=2)
rat = RATPipeline(retriever)
result = rat.run("Какие основные новшества появились в Python 3.12?")
print(result["answer"])

CONFIDENCE_PROMPT = """Оцени свою уверенность в этом утверждении от 0 до 1.

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

Верни только число (например, 0.85)."""


class SelectiveRAT(RATPipeline):
    """RAT с выборочным retrieval по confidence threshold."""

    def __init__(self, retriever, threshold: float = 0.7, **kwargs):
        super().__init__(retriever, **kwargs)
        self.threshold = threshold

    def _get_confidence(self, thought: str) -> float:
        """Оценить уверенность модели в мысли."""
        response = self.llm.invoke(
            CONFIDENCE_PROMPT.format(thought=thought)
        )
        try:
            return float(response.content.strip())
        except ValueError:
            return 0.0  # При ошибке парсинга --- всегда делаем retrieval

    def run(self, question: str) -> dict:
        thoughts = self._generate_initial_thoughts(question)
        revised = []
        retrieval_log = []
        skipped = 0

        for thought in thoughts:
            confidence = self._get_confidence(thought)

            if confidence >= self.threshold:
                # Модель уверена --- пропускаем retrieval
                revised.append(thought)
                skipped += 1
            else:
                # Низкая уверенность --- делаем retrieval
                context = self._retrieve(thought)
                retrieval_log.append({
                    "thought": thought,
                    "confidence": confidence,
                    "retrieved": context[:200]
                })
                revised_thought = self._revise_thought(
                    question, thought, context, revised
                )
                revised.append(revised_thought)

        revised_text = "\n".join(
            f"  {i+1}. {t}" for i, t in enumerate(revised)
        )
        final = self.llm.invoke(
            FINAL_ANSWER_PROMPT.format(
                question=question, revised_thoughts=revised_text
            )
        )

        return {
            "answer": final.content.strip(),
            "initial_thoughts": thoughts,
            "revised_thoughts": revised,
            "retrieval_log": retrieval_log,
            "retrievals_made": len(thoughts) - skipped,
            "retrievals_skipped": skipped
        }