Не так давно стала задача создать персональный чат-бот ассистент для компании занимающейся разработкой ПО. Система должна была иметь как Backend, отвечающий за работу с локальной нейросетью, так и простой FrontEnd, виджет на JavaScript, который можно подключить на любой из страниц компании. Ресурсов описывающих работу RAG-систем полно, однако руководств, которые расскажут и поэтапно проведут разработчика через все необходимые шаги я не нашел. Тем самым, постараюсь восполнить пробел в данной статье.

Введение

Для разработки персонального помощника оптимальным выбором станет система Retrieval-Augmented Generation (RAG) — это сочетание возможностей нейросети и базы знаний, которую мы сформируем на основе данных о компании. В качестве базы мы будем использовать ChromaDB которая работает с использованием Sqlite. В качестве локальной нейросети будем использовать бесплатную сетку Llama [https://ollama.com/download](https://ollama.com/download). Выбор количества параметров оставляю на ваш вкус, лично я выбрал Llama 3.2.

Общая последовательность действий выглядит так:  загружаем данные о компании из сети -> преобразуем их в md формат -> разбиваем информацию на логические блоки -> формируем и сохраняем векторное представление информации в БД -> реализуем алгоритм выборки интересующей нас информации из БД и формирование на их основе ответа от нейросети. Оформляем всю эту красоту в виде API для коммуникации с внешним миром и допиливаем UI.

Писать мы будем на Python 3.12. В качестве API фреймворка выбираю FastAPI.

Для нетерпеливых, все исходники можно найти на моем github.

Итак, приступим!

Структура проекта

Создаем директорию в которой будем реализовывать наш проект, назовем ее ai-chatbot-js. Внутри создадим дополнительно несколько директорий:

data – положим сюда файл links.txt который будет содержать массив ссылок для обработки
docs – здесь будут храниться «сырые» данные скачанные из массива ссылок links.txt
db_metadata_v5 – директория где будет создано локальное хранилище векторной базы данных
providers – здесь реализуем различные провайдеры для работы с нейросетью, в нашем случае либо локальной либо сторонним сервисом
public – наш UI

Подготовка проекта

Вкратце пробежим по шагам необходимым для настройки окружения и необходимых сервисов. Я устанавливал версию Python 3.12.2

• Шаг 1: Создадим виртуальное окружение.

  cd ai-chatbot-js
  python3 -m venv env

• Шаг 2: Подготовка модели Llama с использованием Ollama.
  

  • Переходим на сайт проекта, https://ollama.com/ качаем клиент, тут все просто

  • Далее нам будут необходимы 2 нейросетки одна непосредственно LLM которая будет генерировать ответ пользователю, вторая нейросеть будет отвечать за создание embeddings

  ollama pull llama3.2
  ollama pull mxbai-embed-large

  Обратите внимание на размер нейросетей, llama3.2 занимает около 2Gb, mxbai-embed-large еще 670Mb. Учитывайте это и освободите необходимое место на диске.

  Наше окружение готово, пора приступать к кодингу и начнем мы с загрузки данных о компании из сети

Загрузка данных

В нашем проекте в качестве источника «сырых» данных мы будем использовать сайт компании.

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

Получить ссылки можно различными способами, в моем случае я пользовался онлайн сервисом который на основании root домена распарсил и выдал интересующее меня число внутренних ссылок сайта. Мы же, для простоты примера, позаимствуем несколько ссылок компании Microsoft.

Создадим текстовый файл (data/links.txt) куда мы положим интересующие нас ссылки. Формат файла, полная ссылка на источник разделенные новой строкой:

https://www.microsoft.com/en-us/about
https://news.microsoft.com/facts-about-microsoft/
https://news.microsoft.com/source/

Создадим файл обработчик scrapper.py:

from langchain_community.document_loaders import AsyncHtmlLoader
from langchain_community.document_transformers import Html2TextTransformer, BeautifulSoupTransformer


FILE_TO_PARSE = "data/links.txt"
DIR_TO_STORE = "docs"


def getLinks2Parse() -> list:
    try:
        with open(FILE_TO_PARSE, "r") as f:
            return [link.strip() for link in f.readlines()]
    except:
        return []


def asyncLoader(links):
    loader = AsyncHtmlLoader(links)
    docs = loader.load()

    # Transform
    bs_transformer = BeautifulSoupTransformer()

    for doc in docs:
        doc.page_content = bs_transformer.remove_unwanted_classnames(doc.page_content,
                                                                     ['new-footer', 'main-header',
                                                                      'main-top-block', 'callback__form',
                                                                      'new-footer-bottom', 'blog-article-share', 'blog-article-slider',
                                                                      'blog-article-menu', 'blog__subscribe',
                                                                      'main-top-block__info', 'breadcrumbs'])

    html2text = Html2TextTransformer(ignore_links=True, ignore_images=True)
    docs_transformed = html2text.transform_documents(docs)

    for idx, doc in enumerate(docs_transformed):
        with open(f"{DIR_TO_STORE}/document_{idx}.txt", "w+", encoding="utf-8") as f:
            f.write(doc.page_content)
            print(f"File {DIR_TO_STORE}/document_{idx}.txt saved")


if __name__ == "__main__":
    ls = getLinks2Parse()
    asyncLoader(ls)

Логика довольно проста, при помощью библиотеки langchain_community. document_loaders, в асинхронном режиме скачиваем ссылки, удаляем не интересующие нас данные (строки 25-30) и сохраняем в формате MD. Формат выбран не случайно, далее это нам поможет более качественно семантически разбить содержимое.

Перед запуском, нам понадобится установить несколько зависимостей:

pip3 install langchain_community
pip3 install bs4
pip3 install html2text

Мы готовы, запускаем наш скрипт:

 python3 scrapper.py

В итоге в нашей папке documents появится 3 файла в формате MD.

document_0.txt
document_1.txt
document_2.txt

Создание векторной базы данных

Мы выгрузили информацию о компании, почистили от ненужных данных и преобразовали в интересующий нас формат. Далее одним из наиболее важных шагов преобразование информации в векторное представление.

Создадим файл обработчик ingest.py:

# Langchain dependencies
import hashlib
import os
import shutil

from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import MarkdownTextSplitter
from langchain.schema import Document
from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaEmbeddings


# Path to the directory to save Chroma database
CHROMA_PATH = "./db_metadata_v5"
DATA_PATH = "./docs"
global_unique_hashes = set()


def walk_through_files(path, file_extension='.txt'):
    for (dir_path, dir_names, filenames) in os.walk(path):
        for filename in filenames:
            if filename.endswith(file_extension):
                yield os.path.join(dir_path, filename)


def load_documents():
    """
    Load documents from the specified directory
    Returns:
    List of Document objects:
    """
    documents = []
    for f_name in walk_through_files(DATA_PATH):
        document_loader = TextLoader(f_name, encoding="utf-8")
        documents.extend(document_loader.load())

    return documents


def hash_text(text):
    # Generate a hash value for the text using SHA-256
    hash_object = hashlib.sha256(text.encode())
    return hash_object.hexdigest()


def split_text(documents: list[Document]):
    """
    Split the text content of the given list of Document objects into smaller chunks.
    Args:
    documents (list[Document]): List of Document objects containing text content to split.
    Returns:
    list[Document]: List of Document objects representing the split text chunks.
    """
    # Initialize text splitter with specified parameters
    """
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,  # Size of each chunk in characters
        chunk_overlap=100,  # Overlap between consecutive chunks
        length_function=len,  # Function to compute the length of the text
        add_start_index=True,  # Flag to add start index to each chunk
    )
    """
    text_splitter = MarkdownTextSplitter(
        chunk_size=500,  # Size of each chunk in characters
        chunk_overlap=100,  # Overlap between consecutive chunks
        length_function=len,  # Function to compute the length of the text
    )

    # Split documents into smaller chunks using text splitter
    chunks = text_splitter.split_documents(documents)
    print(f"Split {len(documents)} documents into {len(chunks)} chunks.")

    # Deduplication mechanism
    unique_chunks = []
    for chunk in chunks:
        chunk_hash = hash_text(chunk.page_content)
        if chunk_hash not in global_unique_hashes:
            unique_chunks.append(chunk)
            global_unique_hashes.add(chunk_hash)

    # Print example of page content and metadata for a chunk
    print(f"Unique chunks equals {len(unique_chunks)}.")
    # print(unique_chunks[:-5])

    return unique_chunks  # Return the list of split text chunks


def save_to_chroma(chunks: list[Document]):
    """
    Save the given list of Document objects to a Chroma database.
    Args:
    chunks (list[Document]): List of Document objects representing text chunks to save.
    Returns:
    None
    """
    # Clear out the existing database directory if it exists
    if os.path.exists(CHROMA_PATH):
        shutil.rmtree(CHROMA_PATH)

    # Create a new Chroma database from the documents using OpenAI embeddings
    db = Chroma.from_documents(
        documents=chunks,
        embedding=OllamaEmbeddings(model="mxbai-embed-large"),
        persist_directory=CHROMA_PATH
    )

    # Persist the database to disk
    db.persist()
    print(f"Saved {len(chunks)} chunks to {CHROMA_PATH}.")


def generate_data_store():
    """
    Function to generate vector database in chroma from documents.
    """
    documents = load_documents()  # Load documents from a source
    chunks = split_text(documents)  # Split documents into manageable chunks
    save_to_chroma(chunks)  # Save the processed data to a data store


if __name__ == "__main__":
    generate_data_store()

В данном файле функция split_text требует небольшого пояснения. Параметром она принимает список документов загруженных из нашей папки docs и проводит несколько преобразований:

• При помощи MarkdownTextSplitter создаем экземпляр класса который будет разбивать информацию на логические блоки

• Пробегаем по массиву блоков вычисляя хэш и оставляем только уникальные куски информации

Далее полученные данные сохраняем в Chroma.

Устанавливаем зависимости необходимые для работы:

pip3 install langchain_ollama
pip3 install chromadb

Мы готовы, запускаем наш скрипт:

python3 ingest.py

Как итог, в консоли появится сопутствующая информация:
Saved 18 chunks to ./db_metadata_v5.
а данные будут сохранены в папке db_metadata_v5.

Поздравляю, мы на верном пути!

Дружим LLM и ChromaDB

Нам осталось реализовать выбор подходящего содержания в зависимости от вопроса пользователя (работа с данными в ChromaDB) и преобразование его в удобоваримый ответ от LLM.

Для начала создадим нашу модель models/index.py:

from pydantic import BaseModel


class ChatMessage(BaseModel):
    question: str

Создаем файл обработчик providers/ollama.py:

from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage
from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain.chains.combine_documents import create_stuff_documents_chain

from models.index import ChatMessage


CHROMA_PATH = "./db_metadata_v5"

# Initialize OpenAI chat model
model = OllamaLLM(model="llama3.2:latest", temperature=0.1)


# YOU MUST - Use same embedding function as before
embedding_function = OllamaEmbeddings(model="mxbai-embed-large")

# Prepare the database
db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
chat_history = {}  # approach with AiMessage/HumanMessage

prompt_template = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """
                [INST]You are a sales manager with the name 'AI Assistant'. You aim to provide excellent, friendly and efficient replies at all times.
                You will provide me with answers from the given info.
                If the answer is not included, say exactly “Hmm, I am not sure. Let me check and get back to you.”
                Refuse to answer any question not about the info.
                Never break character.
                No funny stuff.
                If a question is not clear, ask clarifying questions.
                Make sure to end your replies with a positive note.
                Do not be pushy.
                Answer should be in MD format.
                If someone asks for the price, cost, quote or similar, then reply “In order to provide you with a customized and reasonable quote, I would need a 15 minute call.
                Ready for an online meeting?[/INST]
                [INST]Answer the question based only on the following context:
                {context}[/INST]
            """
        ),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{question}")
    ]
)

document_chain = create_stuff_documents_chain(llm=model, prompt=prompt_template)


def query_rag(message: ChatMessage, session_id: str = "") -> str:
    """
    Query a Retrieval-Augmented Generation (RAG) system using Chroma database and OpenAI.
    :param message: ChatMessage The text to query the RAG system with.
    :param session_id: str Session identifier
    :return str
    """

    if session_id not in chat_history:
        chat_history[session_id] = []

    # Generate response text based on the prompt
    response_text = document_chain.invoke({"context": db.similarity_search(message.question, k=3),
                                           "question": message.question,
                                           "chat_history": chat_history[session_id]})

    chat_history[session_id].append(HumanMessage(content=message.question))
    chat_history[session_id].append(AIMessage(content=response_text))

    return response_text

Здесь нам потребуется установить еще одну зависимость для работы с ChromaDB:

pip3 install langchain_chroma

Работа по поиску необходимых данных происходит вот тут - db.similarity_search(message.question, k=3) здесь мы производим поиск по DB похожих данных на наш запрос в указанном количестве. И в качестве контекста прокидываем в наш prompt message. Дополнительно мы отдаем историю диалога с пользователем для создания более релевантной среды.

Подытожим, скрипт создает 2 экземпляра классов для работы с embeddings и LLM, позволяет сохранять историю переписки и осуществляет поиск по документам. В итоге, ответом, мы получим сгенерированный текст на основании найденных кусочков информации и истории переписки с пользователем. Если данных для генерации ответа недостаточно, либо пользователь пытается перенаправить диалог вне заданной роли - мы получим стандартный ответ "Hmm, I am not sure. Let me check and get back to you." из нашего PROMPT.

Подсмотреть подключение внешнего сервиса с более мощной сетью для production можно здесь. Общий принцип тот че что и в работе с локальной LLM, единственное что придется прокинуть API ключ.

Связь с внешним миром, API

Пришло время открыть наше творение внешнему миру. Для этого нам следует написать простейший API. Использовать мы будем FastAPI.

Создадим файл main.py:

from fastapi.middleware.cors import CORSMiddleware

from models.index import ChatMessage
from providers.ollama import query_rag

from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


app.mount("/public", StaticFiles(directory="public"), name="public")


@app.get("/")
async def read_root():
    return {"Hello": "world"}


@app.post("/chat/{chat_id}")
async def ask(chat_id: str, message: ChatMessage):
    return {"response": query_rag(message, chat_id)}

Создаем endpoint POST "/chat/{chat_id}", именно он будет получать идентификатор чата именно он у нас будет отвечать за уникальность диалога между пользователем и LLM.

Не забываем установить необходимые зависимости:

pip3 install "fastapi[standard]"

Запускаем наш API:

fastapi run

Сервер должен запуститься без каких либо проблем. Проверяем наш бэкенд

curl --location 'http://localhost:8000/chat/abc' \
--header 'Content-Type: application/json' \
--data '{
    "question": "Hi, who are you?"
}'

Ответом будет сообщение от нашего вежливого AI

{"response":"*I smile warmly and extend a helping hand*\n\nI'm AI Assistant, your friendly sales manager. I'll be happy to help you with any questions or concerns you may have about our amazing products and services. How can I assist you today?\n\nHave a great day!"}

Поздравляю основная часть работы над персональным ассистентом подошла к концу!

Web Интерфейс

Не буду подробно останавливаться на реализации html/js виджета. В моем случае я попросил сгенерировать окно чата у ChatGPT. Путем несложного допиливания, все привел в удобную для использования форму. Код можно подсмотреть на github.

Заключение

Итак, мы подошли к финалу! Я постарался как можно более проще объяснить шаги через которые прошел сам. После прочтения этой статьи мы научились делать кучу полезных вещей:

1. Работать с данными — процесс сбора информации из интернета, её очистки и преобразования в удобный формат.

2. Векторная база данных — теперь вы знаете, как превратить текстовые данные в удобный для поиска формат с помощью ChromaDB и Embeddings.

3. Интегрирование нейросети — вы научились подключать локальные и внешние LLM (например, Llama) для генерации умных ответов.

Если есть пожелания или правки, дайте знать!