# Псевдокод Agentic RAG цикла — концептуальная модель
class AgenticRAG:
    def __init__(self):
        self.query_analyzer = QueryAnalyzer()     # Анализ запроса
        self.router = QueryRouter()               # Выбор источников
        self.retrievers = {                        # Инструменты поиска
            "vector": VectorRetriever(),
            "sql": SQLRetriever(),
            "web": WebSearchRetriever(),
            "graph": GraphRetriever()
        }
        self.reflector = SelfReflection()         # Оценка качества
        self.generator = ContextGenerator()       # Генерация ответа
        self.max_iterations = 5                   # Защита от зацикливания

    def run(self, query: str) -> Answer:
        analysis = self.query_analyzer.analyze(query)
        context_pool = []
        iteration = 0

        while iteration < self.max_iterations:
            routes = self.router.route(query, analysis, context_pool)
            for source in routes:
                docs = self.retrievers[source].retrieve(query)
                context_pool.extend(docs)

            assessment = self.reflector.assess(query, context_pool)
            if assessment.is_sufficient:
                break
            query = self.reflector.reformulate(query, assessment)
            iteration += 1

        return self.generator.generate(query, context_pool)

from llama_index.core.tools import FunctionTool
from llama_index.core.agent import ReActAgent
from pydantic import BaseModel, Field
import enum

class QueryCategory(str, enum.Enum):
    FACTUAL = "factual"
    ANALYTICAL = "analytical"
    CURRENT_EVENTS = "current_events"
    RELATIONAL = "relational"

class RoutingDecision(BaseModel):
    """Структурированное решение роутера"""
    category: QueryCategory = Field(description="Категория запроса")
    reasoning: str = Field(description="Обоснование выбора")
    sub_queries: list[str] = Field(
        default_factory=list,
        description="Подзапросы для multi-hop сценария"
    )
    needs_multi_source: bool = Field(
        default=False,
        description="Нужны ли несколько источников"
    )

async def route_query(
    query: str,
    llm,
    conversation_history: list | None = None
) -> RoutingDecision:
    """Интеллектуальный роутер: LLM решает, куда направить запрос."""
    system_prompt = """Ты — маршрутизатор запросов. Классифицируй запрос:

    - factual: фактические вопросы по документам/статьям
      Пример: "Какая ставка рефинансирования ЦБ?"
    - analytical: требует агрегации данных, расчётов
      Пример: "Средняя выручка по кварталам за 2025 год"
    - current_events: требует свежей информации из интернета
      Пример: "Курс доллара сегодня"
    - relational: связи между сущностями
      Пример: "Какие компании инвестировали в OpenAI?"

    Для сложных вопросов разбей на подзапросы.
    Если нужны данные из нескольких источников — укажи needs_multi_source."""

    response = await llm.astructured_predict(
        RoutingDecision,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Запрос: {query}"}
        ]
    )
    return response

# Сопоставление категорий с retrieval-инструментами
RETRIEVAL_MAP = {
    QueryCategory.FACTUAL: ["vector_store"],
    QueryCategory.ANALYTICAL: ["sql_database"],
    QueryCategory.CURRENT_EVENTS: ["web_search"],
    QueryCategory.RELATIONAL: ["knowledge_graph"],
}

async def execute_retrieval(
    decision: RoutingDecision,
    query: str,
    tools: dict
) -> list[Document]:
    """Выполняет поиск согласно решению роутера."""
    sources = RETRIEVAL_MAP[decision.category]
    if decision.needs_multi_source:
        sources = ["vector_store", "web_search"]

    # Параллельный поиск по выбранным источникам
    queries = decision.sub_queries or [query]
    tasks = []
    for src in sources:
        for q in queries:
            tasks.append(tools[src].aretrieve(q))

    results = await asyncio.gather(*tasks)
    # Склеиваем результаты, убираем дубликаты
    all_docs = []
    for batch in results:
        all_docs.extend(batch)
    return deduplicate(all_docs)

from pydantic import BaseModel, Field

class ReflectionResult(BaseModel):
    """Результат самооценки агента."""
    relevance_score: float = Field(
        ge=0, le=1,
        description="Насколько контекст релевантен вопросу"
    )
    completeness_score: float = Field(
        ge=0, le=1,
        description="Достаточно ли контекста для полного ответа"
    )
    has_contradictions: bool = Field(
        description="Есть ли противоречия между источниками"
    )
    contradiction_detail: str = Field(
        default="",
        description="Детали противоречий, если есть"
    )
    needs_refinement: bool = Field(
        description="Нужен ли дополнительный поиск"
    )
    refined_query: str = Field(
        default="",
        description="Переформулированный запрос для уточнения"
    )
    missing_information: list[str] = Field(
        default_factory=list,
        description="Какой информации не хватает"
    )

async def self_reflect(
    original_query: str,
    retrieved_context: list[Document],
    llm
) -> ReflectionResult:
    """Агент оценивает качество найденного контекста."""
    context_text = "\n---\n".join(
        f"[{i}] {doc.text[:500]}"
        for i, doc in enumerate(retrieved_context, 1)
    )

    reflection_prompt = """Ты — критический оценщик качества поиска. Проанализируй:

    1. RELEVANCE (0-1): насколько найденный контекст отвечает на вопрос?
       - 0.8-1.0: прямой ответ, высокая релевантность
       - 0.4-0.7: частично релевантно, много шума
       - 0.0-0.3: нерелевантно

    2. COMPLETENESS (0-1): можно ли дать исчерпывающий ответ?
       - 1.0: вся информация есть
       - 0.5-0.9: основной ответ есть, не хватает деталей
       - 0.0-0.4: информации недостаточно

    3. CONTRADICTIONS: есть ли противоречия между источниками?

    Если оценка ниже 0.7 по любой метрике — укажи:
    - needs_refinement = True
    - refined_query — переформулируй вопрос для лучшего поиска
    - missing_information — что именно нужно найти"""

    result = await llm.astructured_predict(
        ReflectionResult,
        messages=[
            {"role": "system", "content": reflection_prompt},
            {"role": "user",
             "content": f"Вопрос: {original_query}\n\nКонтекст:\n{context_text}"}
        ]
    )
    return result


# Основной цикл с self-reflection
async def agentic_rag_cycle(query: str, llm, tools: dict, max_rounds=3):
    """Полный цикл Agentic RAG с самооценкой."""
    decision = await route_query(query, llm)
    all_context = []
    current_query = query

    for round_num in range(1, max_rounds + 1):
        # Поиск
        docs = await execute_retrieval(decision, current_query, tools)
        all_context.extend(docs)

        # Самооценка
        reflection = await self_reflect(query, all_context, llm)

        print(f"[Round {round_num}] relevance={reflection.relevance_score:.2f}, "
              f"completeness={reflection.completeness_score:.2f}")

        if reflection.relevance_score >= 0.8 and reflection.completeness_score >= 0.8:
            print("✅ Контекст достаточен — генерируем ответ")
            break

        if reflection.needs_refinement:
            print(f"🔄 Уточняем запрос: {reflection.refined_query}")
            current_query = reflection.refined_query
            # Обновляем решение роутера под новый запрос
            decision = await route_query(current_query, llm)
        else:
            break
    else:
        print("⚠️ Достигнут лимит итераций")

    return all_context

from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import MultiStepQueryEngine

class MultiHopRetriever:
    """Многошаговый retriever — каждый шаг уточняет поиск."""

    def __init__(self, index: VectorStoreIndex, llm, max_hops=4):
        self.index = index
        self.llm = llm
        self.max_hops = max_hops
        self.seen_queries = set()  # Защита от циклов

    async def decompose(self, query: str) -> list[str]:
        """Разбивает сложный запрос на цепочку подзапросов."""
        prompt = """Разбей сложный вопрос на цепочку простых подзапросов.
        Каждый последующий подзапрос должен опираться на результат предыдущего.

        Пример:
        Вопрос: "Кто автор статьи о CRISPR, опубликованной в самом цитируемом журнале 2023 года?"
        Шаг 1: "Самый цитируемый научный журнал 2023 года"
        Шаг 2: "Статья о CRISPR в {журнал из шага 1} за 2023 год"
        Шаг 3: "Автор статьи {название статьи из шага 2}"

        Выдай ТОЛЬКО список шагов, по одному на строку, без нумерации."""

        response = await self.llm.acomplete(
            f"{prompt}\n\nВопрос: {query}"
        )
        steps = [s.strip() for s in response.text.split("\n") if s.strip()]
        return steps[:self.max_hops]

    async def retrieve_chain(self, query: str) -> list[dict]:
        """Выполняет цепочку поисков, передавая результаты между шагами."""
        steps = await self.decompose(query)
        chain_results = []
        accumulated_context = {}

        for i, step in enumerate(steps):
            # Подставляем результаты предыдущих шагов
            resolved_step = step.format(**accumulated_context)

            # Проверка на цикличность
            if resolved_step in self.seen_queries:
                print(f"⚠️ Обнаружен цикл на шаге {i}, останавливаемся")
                break
            self.seen_queries.add(resolved_step)

            # Поиск и извлечение
            retriever = self.index.as_retriever(similarity_top_k=3)
            nodes = await retriever.aretrieve(resolved_step)

            # Извлекаем ключевую информацию через LLM
            extraction = await self.llm.astructured_predict(
                KeyFacts,
                messages=[{
                    "role": "user",
                    "content": f"Извлеки ключевые факты:\n{step}\n\nДокументы:\n{nodes}"
                }]
            )

            accumulated_context.update(extraction.dict())
            chain_results.append({
                "step": i + 1,
                "query": resolved_step,
                "nodes": nodes,
                "extracted_facts": extraction
            })

        return chain_results


class KeyFacts(BaseModel):
    """Факты, извлечённые на шаге поиска."""
    primary_finding: str
    names: list[str] = Field(default_factory=list)
    dates: list[str] = Field(default_factory=list)
    metrics: dict = Field(default_factory=dict)

from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver
from typing import TypedDict, Annotated, Literal
from langchain_core.messages import HumanMessage, AIMessage

# ===== Состояние агента =====
class AgenticRAGState(TypedDict):
    messages: Annotated[list, add_messages]
    original_query: str
    current_query: str
    route_decision: dict
    retrieved_docs: list[dict]
    reflection: dict
    iteration: int
    max_iterations: int
    final_answer: str

# ===== Нода 1: Анализ запроса =====
def analyze_query(state: AgenticRAGState) -> dict:
    """Анализируем запрос, определяем сложность и тип."""
    query = state["messages"][-1].content

    analysis_prompt = """Проанализируй запрос и определи:
    - Тип: simple / complex / multi-hop
    - Предметная область
    - Нужен ли актуальный веб-поиск
    Ответь кратко, в 2-3 предложения."""

    response = llm.invoke([
        SystemMessage(content=analysis_prompt),
        HumanMessage(content=query)
    ])

    return {
        "original_query": query,
        "current_query": query,
        "iteration": 0,
        "retrieved_docs": [],
        "messages": [AIMessage(content=f"[Analysis] {response.content}")]
    }

# ===== Нода 2: Маршрутизация =====
def route_node(state: AgenticRAGState) -> dict:
    """Выбираем источники для поиска."""
    # Используем route_query из секции 3
    decision = route_query_sync(state["current_query"])
    return {
        "route_decision": decision.dict(),
        "iteration": state["iteration"] + 1
    }

# ===== Нода 3: Поиск =====
def retrieve_node(state: AgenticRAGState) -> dict:
    """Выполняем поиск по выбранным источникам."""
    decision = state["route_decision"]
    query = state["current_query"]
    new_docs = []

    for source in decision.get("sources", ["vector_store"]):
        if source == "vector_store":
            docs = vector_store.similarity_search(query, k=5)
            new_docs.extend(docs)
        elif source == "web_search":
            docs = tavily_search(query, max_results=3)
            new_docs.extend(docs)

    all_docs = state["retrieved_docs"] + new_docs
    return {"retrieved_docs": all_docs}

# ===== Нода 4: Self-Reflection =====
def reflect_node(state: AgenticRAGState) -> dict:
    """Оценка качества найденного контекста."""
    reflection = self_reflect_sync(
        state["original_query"],
        state["retrieved_docs"]
    )
    return {
        "reflection": reflection.dict(),
        "messages": [AIMessage(content=f"[Reflection] relevance={reflection.relevance_score:.2f}")]
    }

# ===== Нода 5: Генерация ответа =====
def generate_node(state: AgenticRAGState) -> dict:
    """Генерируем финальный ответ с цитатами."""
    docs = state["retrieved_docs"]
    context = "\n".join(
        f"[Источник {i}] {doc.page_content[:400]}"
        for i, doc in enumerate(docs, 1)
    )

    response = llm.invoke([
        SystemMessage(content="""Ты — AI-ассистент с доступом к документам. Ответь на вопрос,
        используя ТОЛЬКО предоставленный контекст. Указывай номера источников.
        Если информации недостаточно — скажи об этом честно."""),
        HumanMessage(content=f"Контекст:\n{context}\n\nВопрос: {state['original_query']}")
    ])
    return {"final_answer": response.content}

# ===== Условный роутер: куда идти после reflection =====
def should_continue(state: AgenticRAGState) -> Literal["generate", "route"]:
    reflection = state.get("reflection", {})
    if state["iteration"] >= state.get("max_iterations", 3):
        return "generate"  # Лимит итераций
    if reflection.get("relevance_score", 0) >= 0.8:
        return "generate"  # Контекст достаточен
    return "route"  # Нужен ещё поиск

# ===== Сборка графа =====
builder = StateGraph(AgenticRAGState)

builder.add_node("analyze", analyze_query)
builder.add_node("route", route_node)
builder.add_node("retrieve", retrieve_node)
builder.add_node("reflect", reflect_node)
builder.add_node("generate", generate_node)

builder.set_entry_point("analyze")
builder.add_edge("analyze", "route")
builder.add_edge("route", "retrieve")
builder.add_edge("retrieve", "reflect")
builder.add_conditional_edges(
    "reflect",
    should_continue,
    {"generate": "generate", "route": "route"}
)
builder.add_edge("generate", END)

# Компиляция с поддержкой checkpoint
memory = MemorySaver()
agentic_rag_graph = builder.compile(checkpointer=memory)

# ===== Запуск =====
config = {"configurable": {"thread_id": "user-session-001"}}
result = agentic_rag_graph.invoke(
    {"messages": [HumanMessage(content="Какие последние изменения в регулировании AI в ЕС?")]},
    config=config
)
print(result["final_answer"])

pip install langchain langchain-community langchain-openai chromadb tavily-python langgraph

import os
from typing import TypedDict, Annotated, Literal
from langchain_openai import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field

# ===== Конфигурация =====
llm = ChatOpenAI(model="gpt-4o", temperature=0)
embeddings = OpenAIEmbeddings()
web_search = TavilySearchResults(max_results=3)

# Векторное хранилище с документами (замените на свои данные)
vectorstore = Chroma(
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)

# ===== Pydantic-модели =====
class RouteDecision(BaseModel):
    source: Literal["local", "web", "both"]
    reasoning: str

class ReflectionAssessment(BaseModel):
    relevance: float = Field(ge=0, le=1)
    completeness: float = Field(ge=0, le=1)
    is_sufficient: bool
    suggestion: str = ""

# ===== Состояние =====
class State(TypedDict):
    messages: Annotated[list, add_messages]
    query: str
    source: str
    local_docs: list[str]
    web_docs: list[str]
    all_docs: list[str]
    assessment: dict
    iteration: int

# ===== Ноды графа =====
def route(state: State) -> dict:
    """Агент решает: искать в локальной базе, в интернете или везде."""
    decision = llm.with_structured_output(RouteDecision).invoke([
        SystemMessage(content="""Определи, где искать ответ:
        - local: если вопрос про внутренние документы, политики, продукты компании
        - web: если нужна актуальная информация из интернета
        - both: если вопрос требует и того, и другого"""),
        HumanMessage(content=state["query"])
    ])
    print(f"🎯 Router: source={decision.source} — {decision.reasoning}")
    return {"source": decision.source, "iteration": state.get("iteration", 0) + 1}

def retrieve(state: State) -> dict:
    """Поиск по выбранным источникам."""
    source = state["source"]
    local_results = []
    web_results = []

    local_query = state["query"]
    if state.get("assessment") and state["assessment"].get("suggestion"):
        local_query = state["assessment"]["suggestion"]

    if source in ("local", "both"):
        docs = vectorstore.similarity_search(local_query, k=5)
        local_results = [doc.page_content for doc in docs]
        print(f"📚 Local: найдено {len(local_results)} документов")

    if source in ("web", "both"):
        web_results_raw = web_search.invoke({"query": state["query"]})
        web_results = [r["content"] for r in web_results_raw]
        print(f"🌐 Web: найдено {len(web_results)} результатов")

    all_docs = local_results + web_results
    return {"local_docs": local_results, "web_docs": web_results, "all_docs": all_docs}

def reflect(state: State) -> dict:
    """Self-reflection: оцениваем, достаточно ли контекста."""
    if not state["all_docs"]:
        return {"assessment": {"relevance": 0, "completeness": 0, "is_sufficient": False, "suggestion": state["query"]}}

    context = "\n".join(state["all_docs"][:5])
    assessment = llm.with_structured_output(ReflectionAssessment).invoke([
        SystemMessage(content="""Оцени качество контекста для ответа на вопрос.
        relevance: насколько контекст релевантен (0-1)
        completeness: насколько контекст полон (0-1)
        is_sufficient: True если обе метрики >= 0.7
        suggestion: если недостаточно — как переформулировать запрос"""),
        HumanMessage(content=f"Вопрос: {state['query']}\n\nКонтекст:\n{context}")
    ])
    print(f"🪞 Reflection: relevance={assessment.relevance:.2f} completeness={assessment.completeness:.2f}")
    return {"assessment": {
        "relevance": assessment.relevance,
        "completeness": assessment.completeness,
        "is_sufficient": assessment.is_sufficient,
        "suggestion": assessment.suggestion
    }}

def generate(state: State) -> dict:
    """Генерация финального ответа."""
    context = "\n\n".join(
        f"[{i}] {doc[:300]}"
        for i, doc in enumerate(state["all_docs"], 1)
    )
    response = llm.invoke([
        SystemMessage(content="""Ты — AI-аналитик. Ответь на вопрос, используя ТОЛЬКО предоставленный контекст.
        Указывай [Источник N] для каждой цитаты. Если данных не хватает — честно скажи."""),
        HumanMessage(content=f"Контекст:\n{context}\n\nВопрос: {state['query']}")
    ])
    print("✅ Ответ сгенерирован")
    return {"messages": [AIMessage(content=response.content)]}

def should_retry(state: State) -> Literal["route", "generate"]:
    assessment = state.get("assessment", {})
    if state["iteration"] >= 3:
        return "generate"
    if assessment.get("is_sufficient"):
        return "generate"
    return "route"

# ===== Сборка и запуск =====
builder = StateGraph(State)
builder.add_node("route", route)
builder.add_node("retrieve", retrieve)
builder.add_node("reflect", reflect)
builder.add_node("generate", generate)

builder.set_entry_point("route")
builder.add_edge("route", "retrieve")
builder.add_edge("retrieve", "reflect")
builder.add_conditional_edges("reflect", should_retry, {
    "route": "route",
    "generate": "generate"
})
builder.add_edge("generate", END)

graph = builder.compile(checkpointer=MemorySaver())

# Запуск
config = {"configurable": {"thread_id": "demo-1"}}
result = graph.invoke(
    {"query": "Какие изменения внёс AI Act 2025 в регулирование систем с высоким риском?"},
    config=config
)
print("\n" + "="*70)
print(result["messages"][-1].content)