Продвинутый RAG: переформулировка, гибридный поиск, re-ranking | Урок 9

"Расскажи про это" — что "это"? Retriever ищет буквально.

Пользователь: "Как работает asyncio?"
Пользователь: "А как его тестировать?" — retriever не знает, что "его" = asyncio

"Сравни FastAPI и Django" — нужны документы по обоим, но один запрос может найти только один

Нашли нужный документ, но в нём 80% не по теме вопроса

"asyncio.gather" — векторный поиск может не найти точное имя функции

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

load_dotenv()
model = ChatOpenAI(model=os.getenv("MODEL_NAME", "gpt-4o"), temperature=0)

# Цепочка переформулировки
rewrite_prompt = ChatPromptTemplate.from_messages([
    ("system", """Ты эксперт по информационному поиску.
Перефразируй вопрос пользователя так, чтобы он лучше подходил для поиска
в технической документации. Сделай его конкретнее и убери лишние слова.
Верни только переформулированный вопрос, без пояснений."""),
    ("human", "{question}"),
])

rewriter = rewrite_prompt | model | StrOutputParser()

# Пример
original = "Расскажи, как вообще сделать так, чтобы несколько штук работало одновременно"
rewritten = rewriter.invoke({"question": original})
print(f"Оригинал: {original}")
print(f"Переформулировка: {rewritten}")
# "параллельное выполнение задач asyncio Python coroutines"

# Встраиваем в RAG цепочку
rag_with_rewrite = (
    {"question": rewriter, "original_question": RunnablePassthrough()}
    # | ...  далее стандартная цепочка RAG
)

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document


load_dotenv()
model = ChatOpenAI(model=os.getenv("MODEL_NAME", "gpt-4o"), temperature=0)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# --- Создаем VectorStore ---
texts = [
    "FastAPI поддерживает async/await и работает на Starlette.",
    "Django — полнофункциональный фреймворк с ORM, админкой и шаблонами.",
    "FastAPI автоматически генерирует документацию Swagger.",
    "Django REST Framework добавляет поддержку REST API в Django.",
    "FastAPI использует Pydantic для валидации входных данных.",
    "Django имеет встроенную систему аутентификации пользователей.",
]
vectorstore = FAISS.from_texts(texts, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

# --- Multi-query: генерируем 3 варианта запроса ---

multi_query_prompt = ChatPromptTemplate.from_messages([
    ("system", """Сгенерируй 3 разных варианта поискового запроса для нахождения
информации по вопросу пользователя. Варианты должны охватывать разные аспекты вопроса.
Верни только запросы, по одному на строке, без нумерации и пояснений."""),
    ("human", "{question}"),
])

query_generator = multi_query_prompt | model | StrOutputParser()


def generate_queries(question: str) -> list[str]:
    result = query_generator.invoke({"question": question})
    queries = [q.strip() for q in result.strip().split("\n") if q.strip()]
    queries.append(question)  # добавляем оригинальный вопрос
    return queries


def multi_query_search(question: str) -> list[Document]:
    queries = generate_queries(question)
    print(f"Запросы: {queries}")

    # Собираем все результаты, дедуплицируем по содержимому
    seen = set()
    unique_docs = []
    for query in queries:
        docs = retriever.invoke(query)
        for doc in docs:
            if doc.page_content not in seen:
                seen.add(doc.page_content)
                unique_docs.append(doc)
    return unique_docs


# Тест
docs = multi_query_search("Чем FastAPI лучше Django для REST API?")
print(f"\nНайдено уникальных документов: {len(docs)}")
for doc in docs:
    print(f"  - {doc.page_content[:80]}")

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_classic.retrievers import ContextualCompressionRetriever
from langchain_classic.retrievers.document_compressors import LLMChainExtractor

load_dotenv()
model = ChatOpenAI(model=os.getenv("MODEL_NAME", "gpt-4o"), temperature=0)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Большой текст с лишней информацией
texts = [
    """asyncio — это стандартная библиотека Python для написания конкурентного кода.
Она основана на event loop, который управляет выполнением корутин.
История библиотеки началась в Python 3.4 как экспериментальная возможность.
asyncio.gather() позволяет запускать несколько корутин параллельно и дождаться
всех результатов. Для работы с файлами используй aiofiles.
Библиотека была значительно улучшена в Python 3.7 с введением asyncio.run().""",
]

vectorstore = FAISS.from_texts(texts, embeddings)
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 1})

# Компрессор: извлекает только релевантные части чанка
compressor = LLMChainExtractor.from_llm(model)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever,
)

# Без компрессии будет весь чанк
docs = base_retriever.invoke("как запустить несколько корутин параллельно?")
print("БЕЗ компрессии:")
print(docs[0].page_content)

# С компрессией, только нужная часть
compressed_docs = compression_retriever.invoke("как запустить несколько корутин параллельно?")
print("\nС компрессией:")
print(compressed_docs[0].page_content if compressed_docs else "ничего не найдено")
# "asyncio.gather() позволяет запускать несколько корутин параллельно и дождаться всех результатов."

from dotenv import load_dotenv
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.retrievers import BM25Retriever
from langchain_classic.retrievers import EnsembleRetriever
from langchain_core.documents import Document

load_dotenv()
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

texts = [
    "asyncio.gather() запускает несколько корутин параллельно.",
    "asyncio.sleep() приостанавливает выполнение корутины на заданное время.",
    "FastAPI использует async def для асинхронных эндпоинтов.",
    "Pydantic BaseModel используется для валидации данных в FastAPI.",
    "SQLAlchemy AsyncSession нужен для асинхронной работы с базой данных.",
    "HTTPException в FastAPI вызывается для возврата HTTP-ошибок.",
]

docs = [Document(page_content=t) for t in texts]

# Векторный retriever
vectorstore = FAISS.from_documents(docs, embeddings)
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# BM25 retriever (точный поиск по ключевым словам)
bm25_retriever = BM25Retriever.from_documents(docs)
bm25_retriever.k = 3

# Ensemble: объединяем с весами
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6],  # 40% BM25, 60% векторный
)

# Тест: точный термин — BM25 справится лучше
results = ensemble_retriever.invoke("asyncio.gather")
print("Гибридный поиск по 'asyncio.gather':")
for doc in results:
    print(f"  - {doc.page_content}")

Гибридный поиск по 'asyncio.gather':
  - asyncio.gather() запускает несколько корутин параллельно.
  - asyncio.sleep() приостанавливает выполнение корутины на заданное время.
  - FastAPI использует async def для асинхронных эндпоинтов.
  - HTTPException в FastAPI вызывается для возврата HTTP-ошибок.
  - SQLAlchemy AsyncSession нужен для асинхронной работы с базой данных.
  - Pydantic BaseModel используется для валидации данных в FastAPI.

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.retrievers import BM25Retriever
from langchain_classic.retrievers import EnsembleRetriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document

load_dotenv()
model = ChatOpenAI(model=os.getenv("MODEL_NAME", "gpt-4o"), temperature=0)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Пример базы знаний для re-ranking
texts = [
    "asyncio.gather() запускает несколько корутин параллельно.",
    "asyncio.sleep() приостанавливает выполнение корутины.",
    "FastAPI использует async def для асинхронных эндпоинтов.",
    "aiohttp — библиотека для асинхронных HTTP запросов.",
    "HTTPException в FastAPI вызывается для возврата HTTP ошибок.",
    "SQLAlchemy AsyncSession нужен для асинхронной работы с БД.",
]
docs_for_search = [Document(page_content=t) for t in texts]
vectorstore = FAISS.from_documents(docs_for_search, embeddings)
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
bm25_retriever = BM25Retriever.from_documents(docs_for_search)
bm25_retriever.k = 3

ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6],
)

rerank_prompt = ChatPromptTemplate.from_messages([
    ("system", """Оцени релевантность фрагмента документа для ответа на вопрос.
Верни только число от 0 до 10, где:
0 — фрагмент совершенно не относится к вопросу
10 — фрагмент напрямую отвечает на вопрос
Только число, без пояснений."""),
    ("human", "Вопрос: {question}\n\nФрагмент: {document}"),
])

scorer = rerank_prompt | model | StrOutputParser()


def rerank_documents(question: str, docs: list[Document], top_k: int = 3) -> list[Document]:
    """Переранжирует документы по релевантности к вопросу."""
    scored = []
    for doc in docs:
        try:
            score_str = scorer.invoke({
                "question": question,
                "document": doc.page_content,
            })
            score = float(score_str.strip())
        except (ValueError, Exception):
            score = 0.0
        scored.append((score, doc))

    # Сортируем по убыванию оценки
    scored.sort(key=lambda x: x[0], reverse=True)

    print("Оценки re-ranking:")
    for score, doc in scored[:top_k]:
        print(f"  {score:.1f}: {doc.page_content[:70]}...")

    return [doc for _, doc in scored[:top_k]]


# Пример использования с retriever
question = "Как сделать асинхронный HTTP запрос?"
candidates = ensemble_retriever.invoke(question)  # 6 кандидатов
reranked = rerank_documents(question, candidates, top_k=3)

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory

load_dotenv()

model = ChatOpenAI(model=os.getenv("MODEL_NAME", "gpt-4o"), temperature=0)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# VectorStore
texts = [
    "asyncio.gather() запускает корутины параллельно и возвращает список результатов.",
    "asyncio.create_task() создаёт задачу из корутины без ожидания её завершения.",
    "asyncio.run() — точка входа для запуска асинхронного кода в синхронном контексте.",
    "pytest-asyncio позволяет тестировать async def функции с помощью @pytest.mark.asyncio.",
    "aiohttp — библиотека для асинхронных HTTP запросов, альтернатива requests.",
]
vectorstore = FAISS.from_texts(texts, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# Шаг 1: переформулируем вопрос с учетом истории
contextualize_prompt = ChatPromptTemplate.from_messages([
    ("system", """Используя историю диалога и последний вопрос пользователя,
сформулируй самодостаточный поисковый запрос. Если вопрос уже самодостаточен, верни его без изменений.
Верни только запрос, без пояснений."""),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}"),
])

contextualize_chain = contextualize_prompt | model | StrOutputParser()

# Шаг 2: Промпт RAG с историей
rag_prompt = ChatPromptTemplate.from_messages([
    ("system", """Отвечай на вопрос только на основе контекста.
Если ответа нет, скажи об этом.

Контекст:
{context}"""),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}"),
])


def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


# Сборка: contextualize → retriever → generate
def make_rag_chain(retriever):
    return (
        RunnablePassthrough.assign(
            # Переформулируем вопрос с учетом истории, используем для поиска
            context=contextualize_chain | retriever | format_docs,
        )
        | rag_prompt
        | model
        | StrOutputParser()
    )

conversational_rag = make_rag_chain(retriever)

# Оборачиваем с историей
store = {}


def get_session_history(session_id: str) -> ChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


chain_with_history = RunnableWithMessageHistory(
    conversational_rag,
    get_session_history,
    input_messages_key="question",
    history_messages_key="history",
)

config = {"configurable": {"session_id": "dev_session"}}


def chat(question: str) -> str:
    return chain_with_history.invoke({"question": question}, config=config)


# Тест, контекстный диалог
print(chat("Расскажи про asyncio.gather"))
print()
print(chat("А как его тестировать?"))
# Второй вопрос переформулируется в "как тестировать asyncio.gather?"
print()
print(chat("Есть ли альтернативы для HTTP запросов?"))
# Переформулируется с учетом контекста про asyncio

# Дорого: 20 вызовов LLM на каждый запрос
candidates = retriever.invoke(question)  # k=20
reranked = rerank_documents(question, candidates)

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.retrievers import BM25Retriever
from langchain_classic.retrievers import EnsembleRetriever
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

load_dotenv()

model = ChatOpenAI(model=os.getenv("MODEL_NAME", "gpt-4o"), temperature=0)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# --- База знаний ---

raw_texts = [
    "asyncio.gather() принимает корутины и запускает их параллельно. Возвращает список результатов в том же порядке.",
    "asyncio.create_task() создаёт Task из корутины. Задача начинает выполняться немедленно в фоне.",
    "asyncio.sleep(n) приостанавливает корутину на n секунд без блокировки event loop.",
    "asyncio.wait() похож на gather, но возвращает множества done и pending задач.",
    "aiohttp.ClientSession используется для асинхронных HTTP-запросов. Сессию нужно закрывать явно.",
    "httpx.AsyncClient — современная альтернатива aiohttp с похожим API на requests.",
    "pytest-asyncio добавляет поддержку async def тестов. Используй декоратор @pytest.mark.asyncio.",
    "asyncio.timeout() (Python 3.11+) устанавливает таймаут на выполнение корутины.",
    "FastAPI автоматически запускает async def эндпоинты в event loop без дополнительных настроек.",
    "SQLAlchemy AsyncSession требует async with для управления транзакциями.",
]

docs = [Document(page_content=t, metadata={"id": i}) for i, t in enumerate(raw_texts)]

# --- Гибридный retriever ---

vectorstore = FAISS.from_documents(docs, embeddings)
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

bm25_retriever = BM25Retriever.from_documents(docs)
bm25_retriever.k = 4

hybrid_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6],
)

# --- Query rewriting ---

rewrite_prompt = ChatPromptTemplate.from_messages([
    ("system", "Перефразируй вопрос для технического поиска по документации. Только запрос, без пояснений."),
    ("human", "{question}"),
])
rewriter = rewrite_prompt | model | StrOutputParser()

# --- Re-ranking ---

rerank_prompt = ChatPromptTemplate.from_messages([
    ("system", "Оцени релевантность фрагмента к вопросу от 0 до 10. Только число."),
    ("human", "Вопрос: {question}\n\nФрагмент: {document}"),
])
scorer = rerank_prompt | model | StrOutputParser()

def rerank(question: str, docs: list[Document], top_k: int = 3) -> list[Document]:
    scored = []
    for doc in docs:
        try:
            s = float(scorer.invoke({"question": question, "document": doc.page_content}).strip())
        except Exception:
            s = 0.0
        scored.append((s, doc))
    scored.sort(key=lambda x: x[0], reverse=True)
    return [d for _, d in scored[:top_k]]

# --- Промпт RAG ---

rag_prompt = ChatPromptTemplate.from_messages([
    ("system", """Отвечай только на основе контекста. Если ответа нет, скажи об этом.

Контекст:
{context}"""),
    ("human", "{question}"),
])

def format_docs(docs):
    return "\n\n".join(f"[{i+1}] {d.page_content}" for i, d in enumerate(docs))

# --- Продвинутая RAG цепочка ---

chain = (
    RunnablePassthrough.assign(
        rewritten=lambda x: rewriter.invoke({"question": x["question"]}),
    )
    | RunnablePassthrough.assign(
        candidates=lambda x: hybrid_retriever.invoke(x["rewritten"]),
    )
    | RunnablePassthrough.assign(
        top_docs=lambda x: rerank(x["rewritten"], x["candidates"], top_k=3),
    )
    | RunnablePassthrough.assign(
        context=lambda x: format_docs(x["top_docs"]),
    )
    | rag_prompt
    | model
    | StrOutputParser()
)

# --- Тестирование ---

questions = [
    "Как запустить несколько задач одновременно?",
    "asyncio.gather vs asyncio.wait, в чём разница?",
    "Что такое LangChain?",
    "Какие есть ORM вна языке Nim?",
]

for q in questions:
    print(f"\nВопрос: {q}")
    answer = chain.invoke({"question": q})
    print(f"Ответ: {answer}")
    print("-" * 60)

Вопрос: Как делать параллельные запросы?
Варианты запроса: ['параллельные HTTP запросы asyncio', 'concurrent requests Python', ...]
Кандидатов до re-ranking: 8
Топ-3 после re-ranking: [7.0, 6.5, 5.0]
Ответ: ...

Вопрос: А с таймаутом?
[Переформулировка с учётом контекста: "параллельные HTTP запросы с таймаутом"]
...