Привет, Хабр!
В этой статье я объясню, как работает технология 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 вы можете узнать в моем телеграмм канале:
Зависимости
Для начала установим 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_tavilytavily_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 субъектов.'.....На этом я закончу статью. Спасибо за прочтение!
Если вам было интересно и вы узнали что то новое, поставьте продвижение статье и подписывайтесь, чтобы не пропустить продолжение.
