# ОБЫЧНЫЙ RAG: одношаговая стратегия
# Запрос → Retrieve → Augment → Generate
def naive_rag(query):
    docs = vectorstore.similarity_search(query, k=3)
    context = "\n".join([d.page_content for d in docs])
    return llm.invoke(f"Ответь на основе: {context}\nВопрос: {query}")

# AGENTIC RAG: многошаговая итеративная стратегия
# Plan → Retrieve → Evaluate → Replan → Generate
def agentic_rag(query):
    plan = planner.decompose(query)  # Разбить на подзадачи
    context = []
    for sub_query in plan:
        docs = retrieve_and_grade(sub_query)
        if not is_relevant(docs):
            sub_query = rewrite_query(sub_query)  # Переформулировать
            docs = retrieve_and_grade(sub_query)
        context.extend(docs)
    return synthesize(context, query)

# Установка: pip install langchain langchain-community chromadb pypdf
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.llms import Ollama

# Загрузка PDF и создание векторного индекса
loader = PyPDFLoader("docs/report.pdf")
docs = loader.load()

splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = splitter.split_documents(docs)

vectorstore = Chroma.from_documents(
    chunks,
    OllamaEmbeddings(model="nomic-embed-text")
)

# Retrieval chain — 3 строки!
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
llm = Ollama(model="llama3.1")
result = retriever.invoke("Какие KPI в Q1 2025?")
print(result)

# Агент сам решает: достаточно ли контекста? Если нет — переформулирует запрос
from langchain.agents import create_react_agent, AgentExecutor
from langchain.tools import Tool

class IterativeRetriever:
    def __init__(self, vectorstore, llm):
        self.vs = vectorstore
        self.llm = llm
        self.max_iterations = 3

    def search(self, query):
        context = []
        current_query = query
        for i in range(self.max_iterations):
            docs = self.vs.similarity_search(current_query, k=3)
            context.extend(docs)
            # LLM оценивает: достаточно ли информации?
            sufficiency = self._grade_sufficiency(context, query)
            if sufficiency == "sufficient":
                break
            # Переформулировка: "search for: X" → "search for: X related to Y"
            current_query = self._rewrite(current_query, context)
        return self._generate_answer(context, query)

    def _grade_sufficiency(self, context, query):
        prompt = f"Достаточно ли контекста для ответа?\nВопрос: {query}\nКонтекст: {context}\nОтветь: sufficient/insufficient"
        return self.llm.invoke(prompt).strip().lower()

# Агент извлекает фильтры из вопроса: "покажи статьи 2024" → filter year=2024
from langchain.retrievers import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo

# Описываем метаданные документов
metadata_field_info = [
    AttributeInfo(name="year", description="Год публикации", type="integer"),
    AttributeInfo(name="author", description="Автор документа", type="string"),
    AttributeInfo(name="category", description="Категория: tech, finance, legal", type="string"),
]

# Self-query retriever — LLM сам построит фильтр
self_query_retriever = SelfQueryRetriever.from_llm(
    llm=Ollama(model="llama3.1"),
    vectorstore=vectorstore,
    document_contents="Научные статьи и отчёты",
    metadata_field_info=metadata_field_info,
)

# Запрос: "статьи Петрова за 2024 про нейросети"
# Агент автоматически построит фильтр: year=2024, author="Петров"
docs = self_query_retriever.invoke(
    "статьи Петрова за 2024 год про нейросети в категории tech"
)
print(f"Найдено документов: {len(docs)}")

# Агент разбивает сложный вопрос на подвопросы, ищет каждый, объединяет ответы
def multi_hop_qa(complex_query, retriever, llm):
    # Шаг 1: Декомпозиция вопроса на подвопросы
    decompose_prompt = f"""Разбей сложный вопрос на простые подвопросы (по одному на строку):
Вопрос: {complex_query}
Подвопросы:"""
    sub_questions = llm.invoke(decompose_prompt).strip().split("\n")

    # Шаг 2: Поиск ответа на каждый подвопрос
    intermediate_answers = {}
    for sq in sub_questions:
        docs = retriever.invoke(sq)
        context = "\n".join([d.page_content for d in docs])
        answer = llm.invoke(f"Контекст: {context}\nВопрос: {sq}\nКраткий ответ:")
        intermediate_answers[sq] = answer

    # Шаг 3: Синтез финального ответа из промежуточных
    synthesis_prompt = f"Объедини ответы в один связный ответ на вопрос: {complex_query}\n\n"
    for q, a in intermediate_answers.items():
        synthesis_prompt += f"Q: {q}\nA: {a}\n\n"
    return llm.invoke(synthesis_prompt)

# Пример: "Сравни цены на GPU у поставщика X и Y и скажи, у кого дешевле"
answer = multi_hop_qa(
    "Сравни цены на A100 у NVIDIA и AMD. Кто предлагает лучшую цену за терафлопс?",
    retriever, llm
)

# pip install ragas
from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy, context_precision
from datasets import Dataset

# Подготовка датасета для оценки
eval_data = {
    "question": ["Какие KPI в Q1?", "Кто CEO компании?", "Бюджет на 2025?"],
    "answer": ["KPI: выручка +15%...", "CEO: Иванов И.И.", "Бюджет: 50 млн..."],
    "contexts": [["отчёт Q1 2025"], ["страница команды"], ["финплан 2025"]],
    "ground_truth": ["Выручка выросла на 15% в Q1 2025", "CEO Иванов Иван", "50 млн рублей"],
}
dataset = Dataset.from_dict(eval_data)

# Расчёт метрик
result = evaluate(
    dataset,
    metrics=[faithfulness, answer_relevancy, context_precision],
)
print(f"Faithfulness: {result['faithfulness']:.2%}")
print(f"Answer Relevancy: {result['answer_relevancy']:.2%}")
print(f"Context Precision: {result['context_precision']:.2%}")