Привет, Хабр!

В этой статье я объясню, как работает технология RAG (Retrieval-Augmented Generation), и покажу её базовые реализации. Для примеров я буду использовать фреймворк LangGraph — его основы я разбирал в предыдущей статье

В конце статьи вас ждет дополнительный пример, поэтому дочитывайте до конца.

Как устроен RAG

Технология RAG состоит из двух ключевых компонентов:

Индексация (Indexing)
- Загрузка данных
- Разбиение на фрагменты
- Векторизация
- Хранение
Поиск и генерация (Retrieval and Generation)
- Извлечение релевантной информации
- Обработка найденной информации
- Генерация ответа

Этап индексации (Indexing)

Процесс индексации базово состоит из пяти частей:

Загрузка данных (Load)
Прежде обработать данные, их нужно загрузить. Для этого можно использовать:
- Document Loader
- PyPDF
- CSVLoader
- и т.д.
  Выбор зависит от источника данных. Я буду использовать Document Loader, так как буду загружать текстовый файл.
Разбиение на фрагменты
Большие документы необходимо разделить на меньшие части. Это необходимо по причинам:
- Ограниченный контекст у языковых моделей
- Увеличение точности поиска
  Даже если вы используете современные LLM с контекстным окном в миллиарды токенов, не пропускайте этот шаг.
Векторизация (Emdeddings)
После получения обработанного текста, его необходимо преобразовать в векторное представление. Для этого используются embedding модели. Я буду использовать
"cointegrated/LaBSE-en-ru" с HuggingFace
Хранение (Store)
После индексации данные сохраняются в векторной БД (например, FAISS, Chroma)

Вторая часть. Retrieval and generation

Retrieve. Отвечает за поиск релевантных запросу фрагментов из базы данных. Для этого используются ретриверы. Подробнее о них можно почитать в моей статье
Generation. Заключительный шаг, на котором формируется итоговый prompt для модели, и генерируются ответ.
Вторая часть

Больше о AI и NLP вы можете узнать в моем телеграмм канале:

https://t.me/Viacheslav_Talks

Зависимости

Для начала установим langchain, langgraph

!pip install langgraph langchain langchain_core langchain_community

!pip install --quiet -U langchain_openai langchain_core langgraph langgraph-prebuilt

и библиотеки для работы с моделями

!pip install langchain-gigachat

!pip install langchain_huggingface

Начинаем начинать

Загрузка документа, Для примера я буду использовать текст о России, который состоит из 4000 символов.

Файл можно найти по ссылке

from langchain_community.document_loaders import TextLoader
from pprint import pprint


loader  = TextLoader("/content/Ruusai.txt")
documents  = loader.load() #получили обьект типа Document

Разделение на фрагменты. Для этого я воспользуюсь RecursiveCharacterTextSplitter с параметрами:

chunk_size=1000 (произвольное значение)
chunk_overlap=100 (произвольное значение)

Векторизация и хранение. Я буду использовать InMemoryVectorStore в качестве векторной базы.

from langchain_core.vectorstores import InMemoryVectorStore
from langchain_huggingface import HuggingFaceEmbeddings
import os


HF_TOKEN = 'ВАШ ТОКЕН'
os.environ['HF_TOKEN'] = HF_TOKEN
embeddings_model_name  = "cointegrated/LaBSE-en-ru"

model_embed = HuggingFaceEmbeddings(
    model_name=embeddings_model_name
)
vec_store  = InMemoryVectorStore(model_embed)

_ = vec_store.add_documents(chunks) #добавляем фрамгенты в базу

Определение графа

Перед этим создадим экземпляр языковой модели. Я буду использовать GigaChat.

from langchain_gigachat import GigaChat

llm  = GigaChat(
                verify_ssl_certs=False, 
                credentials="ВАЩ КЛЮЧ",
                model="GigaChat-2"
)

Определение состояния графа. В моем случае достаточно 3 поля:

Запрос пользователя
Найденный контекст
Сгенерированный ответ

from typing import TypedDict


class State(TypedDict):
  question: str
  context: list[str]
  answer: str

Определение узлов графа. Минимальная реализация требует двух узлов, соответствующих основным компонентам RAG:

Retrieval Node - поиск релевантной информации
Generation Node - формирование ответа на основе контекста

from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate


def retrieve(state: State):
  retrieved_docs: List[Document]  = vec_store.similarity_search_with_score(state["question"])
  retrieved_content  = [doc[0].page_content for doc in retrieved_docs]
  return {"context": retrieved_content}


def generate_answer(state: State): 
  system_prompt  = """
  Ты  - умный ассистент, который должен отвечать на вопрос пользователя, основываясь на найденном контексте.
  Для ответа используй только найденный контекст.

  Найденный контекст: {context}
  """

  prompt  = ChatPromptTemplate.from_messages(
      [
          ("system", system_prompt),
          ("human", state["question"])
      ]
  )

  chain  = prompt | llm
  answer  = chain.invoke({"context": state["context"]})
  return {"answer": answer}

Определим связи между узлами

from langgraph.graph import StateGraph, START, END


builder  = StateGraph(State)
builder.add_node("retrieve", retrieve)
builder.add_node("generate_answer", generate_answer)

builder.add_edge(START, "retrieve")
builder.add_edge("retrieve", "generate_answer")
builder.add_edge("generate_answer", END)

В результате получили последовательный граф

from IPython.display import Image, display


graph  = builder.compile()
display(Image(graph.get_graph().draw_mermaid_png()))

Результаты:

from pprint import pprint


question  = "в каких международных организациях состоит Россия?"

answer  = graph.invoke({"question": question})
pprint(answer)

content='Россия состоит в следующих международных организациях:\n- ООН\n- G20\n- ЕАЭС\n- СНГ\n- ОДКБ\n- ВТО\n- ОБСЕ\n- ШОС\n- АТЭС\n- БРИКС\n- МОК\n\nи других.' additional_kwargs={} response_metadata={'token_usage': {'prompt_tokens': 524, 'completion_tokens': 56, 'total_tokens': 580, 'precached_prompt_tokens': 2}, 'model_name': 'GigaChat-2:2.0.28.2', 'x_headers': {'x-request-id': 'c60ef6da-fc1a-424b-83a4-b99d441bcf22', 'x-session-id': '5dc8944f-3355-44e1-a57b-e91728509bca', 'x-client-id': None}, 'finish_reason': 'stop'} id='c60ef6da-fc1a-424b-83a4-b99d441bcf22' usage_metadata={'output_tokens': 56, 'input_tokens': 524, 'total_tokens': 580, 'input_token_details': {'cache_read': 2}}

{'answer': AIMessage(content='Россия состоит в следующих международных организациях:\n- ООН\n- G20\n- ЕАЭС\n- СНГ\n- ОДКБ\n- ВТО\n- ОБСЕ\n- ШОС\n- АТЭС\n- БРИКС\n- МОК\n\nи других.', additional_kwargs={}, response_metadata={'token_usage': {'prompt_tokens': 524, 'completion_tokens': 56, 'total_tokens': 580, 'precached_prompt_tokens': 2}, 'model_name': 'GigaChat-2:2.0.28.2', 'x_headers': {'x-request-id': 'c60ef6da-fc1a-424b-83a4-b99d441bcf22', 'x-session-id': '5dc8944f-3355-44e1-a57b-e91728509bca', 'x-client-id': None}, 'finish_reason': 'stop'}, id='c60ef6da-fc1a-424b-83a4-b99d441bcf22', usage_metadata={'output_tokens': 56, 'input_tokens': 524, 'total_tokens': 580, 'input_token_details': {'cache_read': 2}}),
 'context': ['Россия — многонациональное государство с широким этнокультурным '
             'многообразием[20]. Согласно результатам переписи населения '
             'России 2020—2021 года, в стране живут представители свыше 190 '
             ...........................................................

question  = "какими компаниями владеет Илон Маск?"

answer  = graph.invoke({"question": question})
pprint(answer)

 content='Из предоставленного контекста невозможно сформировать ответ на этот вопрос.' additional_kwargs={} response_metadata={'token_usage': {'prompt_tokens': 58, 'completion_tokens': 13, 'total_tokens': 71, 'precached_prompt_tokens': 2}, 'model_name': 'GigaChat-2:2.0.28.2', 'x_headers': {'x-request-id': '3537302b-ae38-4e82-97cb-affde8772260', 'x-session-id': '27a801e4-4f7e-4e20-a4c1-f62a88bcc8c6', 'x-client-id': None}, 'finish_reason': 'stop'} id='3537302b-ae38-4e82-97cb-affde8772260' usage_metadata={'output_tokens': 13, 'input_tokens': 58, 'total_tokens': 71, 'input_token_details': {'cache_read': 2}}

{'answer': AIMessage(content='Из предоставленного контекста невозможно сформировать ответ на этот вопрос.', additional_kwargs={}, response_metadata={'token_usage': {'prompt_tokens': 58, 'completion_tokens': 13, 'total_tokens': 71, 'precached_prompt_tokens': 2}, 'model_name': 'GigaChat-2:2.0.28.2', 'x_headers': {'x-request-id': '3537302b-ae38-4e82-97cb-affde8772260', 'x-session-id': '27a801e4-4f7e-4e20-a4c1-f62a88bcc8c6', 'x-client-id': None}, 'finish_reason': 'stop'}, id='3537302b-ae38-4e82-97cb-affde8772260', usage_metadata={'output_tokens': 13, 'input_tokens': 58, 'total_tokens': 71, 'input_token_details': {'cache_read': 2}}),
 'context': [],

Расширяем возможности

Я решил вставить в эту статью пример с использованием инструмента TavilySearch.

Я буду использовать его в качестве дополнительного узла в графе, который я использовал выше. Если информация не была найдена в векторной базе, то агент спросит у пользователя, можно ли воспользоваться поиском в интернете.

Установка зависимостей:

!pip install langchain_tavily

tavily_api_key  = "ВАШ КЛЮЧ. ЕГО МОЖНО ПОЛУЧИТЬ БЕСПЛАТНО НА САЙТЕ TAVILY"
os.environ["TAVILY_API_KEY"] = tavily_api_key

Затем я создам инструмент и проверю работу

from langchain_tavily import TavilySearch

search_tool  = TavilySearch(
    max_results = 5
)

answer_tool  = search_tool.invoke("Россия")
content  = [ans["content"] for ans in answer_tool["results"]]
pprint(content)

Определим состояние графа и узлы. Изменений:

Добавим еще один узел для поиска информации в интернете
Функцию, которая будет запрашивать подтверждение у пользователя
Порог релевантности я установлю равным 0.4

from typing import TypedDict
from langgraph.types import interrupt, Command


class State(TypedDict):
  question: str
  context: list[str]
  answer: str

Определим связи между узлами

builder  = StateGraph(State)
builder.add_node("retrieve", retrieve)
builder.add_node("generate_answer", generate_answer)
builder.add_node("web_search", web_search)

builder.add_edge(START, "retrieve")
#длбавляем возможность выбора следующего узла
builder.add_conditional_edges(
    "retrieve",
    use_search,
    {
        "web_search": "web_search",
        "generate_answer": "generate_answer"
    }
)

builder.add_edge("web_search", "generate_answer")
builder.add_edge("generate_answer", END)

graph  = builder.compile()
display(Image(graph.get_graph().draw_mermaid_png()))

В use_search я использовал interrupt. Благодаря interrupt выполнение узла графа прерывается. Эта функция похожа на input, но с одним существенным отличием. Input продолжает выполнение с места прерывания. Interrupt начинает выполнение с первой строки узла, в котором произошло прерывание.

Результаты. Для запуска такого графа нужно есть:

Для компилирования графа с interrupt необходимо использовать память, чтобы запоминать выполненные шаги. Я буду использовать InMemorySaver.
Мы должны сами проверять вызов прерывания. Если оно было вызвано, мы можем получить ответ от пользователя и заново запустить граф с полученным ответом.

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

config = {"configurable": {"thread_id": 1}}
result = graph.invoke({"question": "сколько субьектов в россии?"}, config=config)

if result.get("__interrupt__", None):
  age  = input("Документы не найдены. Использовать поиск в интернете (accept)?")
  final_result  = graph.invoke(Command(resume=str(age)), config=config)
  pprint(final_result)

content='В состав Российской Федерации входят 89 субъектов.'.....

На этом я закончу статью. Спасибо за прочтение!

Если вам было интересно и вы узнали что то новое, поставьте продвижение статье и подписывайтесь, чтобы не пропустить продолжение.

Создаем свой RAG: от загрузки данных до генерации ответов с LangGraph. Часть 2