Как я сделал RAG для своей компании

Появилась свободное время, и я решил сделать RAG (Retrieval Augmented Generation) для нашей компании. Компания небольшая, но документации технической и бизнес накопилось очень много, в основном на wiki.

Цель - подключить бота в slack, который быстро может выдать инфу по нужной теме.

Источник знаний:

1. Wiki

2. JIRA

3. Slack

Сегодня покроем только создание RAG из Wiki + создание бота. Писалось все за 2 дня, поэтому жутко не оптимально, но бот уже работает и выполняет свою функцию, дальше будет улучшать.

Для начал чуть-чуть теории. Из чего сделана любая RAG?

• Парсинг данных из источника в текст (например в json)

• Конвертер из текстовых данных (json) в векторную базу при помощи embed модели и векторной БД

• Получение запрос пользователя (в моем случае в Slack)

• Поиск векторов схожих с запросом пользователя в векторной БД

• Отправка топ схожих векторов в generator model

• Получение ответа от generator model

• Отправка ответ пользователю (в моем случае в Slack)

Векторная база — это хранилище, где данные представлены в виде векторов, позволяющих быстро находить похожие объекты с помощью поиска ближайших соседей.

Пример схожести текстов в векторе:

«Как работает векторная база данных?»
«Что такое векторный поиск в базах данных?»

Векторное представление (условное, для примера):

import numpy as np

vector1 = np.array([0.12, 0.85, 0.64, 0.33, 0.92])  # Вектор для первого текста
vector2 = np.array([0.14, 0.83, 0.67, 0.35, 0.91])  # Вектор для второго текста

# Вычисляем косинусное сходство (показывает, насколько тексты похожи)
cosine_similarity = np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))

print(f"Косинусное сходство: {cosine_similarity:.3f}")  # Близко к 1 → тексты похожи

Вывод: Векторы близки, потому что тексты имеют схожий смысл. Векторная база позволяет быстро находить такие похожие тексты по запросу пользователя.

Начнем с парсинга данных

С wiki можно получить список страниц.

# URL to get the list of pages in the SQA space
wiki_url = "https://company.atlassian.net/wiki/rest/api/content"

# Request parameters to get pages from the SQA space
params = {
    'spaceKey': 'SQA',  # space Name
    'limit': 100,       # Limit the number of pages
    'expand': 'space',  # Expand the space information
    'type': 'page',     # Only pages
    'format': 'json'    # JSON response format
}

username = os.getenv('WIKI_USERNAME')
apikey = os.getenv('WIKI_APIKEY')

# Headers
headers = {
    'Accept': 'application/json'
}

# Perform the request with basic authorization to get the list of pages
response = requests.get(wiki_url, headers=headers, params=params, auth=(username, apikey))

А затем сами страницы.

# Check the response
if response.status_code == 200:
    print("Successful request to get the list of pages.")
    pages = response.json()['results']  # Extract pages from the response

    # Open the file for writing
    with open('./Data/wiki_pages.json', 'w', encoding='utf-8') as f:
        # Write the list of pages to the file
        json.dump(pages, f, ensure_ascii=False, indent=4)

    print(f"The list of pages has been written to 'wiki_pages.json'. Starting to load page content...")

    # Load the content of each page
    for page in pages:
        page_id = page['id']
        page_url = f"https://company.atlassian.net/wiki/rest/api/content/{page_id}"

        # Request parameters to get the page content
        content_params = {
            'expand': 'body.storage',  # Get the HTML content of the page
            'format': 'json'
        }

        # Perform the request to get the page content
        content_response = requests.get(page_url, headers=headers, params=content_params, auth=(username, apikey))

        # Check the response
        if content_response.status_code == 200:
            content = content_response.json()
            page_title = page['title']
            page_body = content['body']['storage']['value']

            # Clean the HTML content
            cleaned_title = clean_html(page_title)
            cleaned_body = clean_html(page_body)

            # Save the content to a separate file
            with open(f"./Data/wiki_page_{page_id}.json", 'w', encoding='utf-8') as page_file:
                json.dump({'title': cleaned_title, 'content': cleaned_body}, page_file, ensure_ascii=False, indent=4)

            print(f"Page content '{page_title}' saved to './Data/wiki_page_{page_id}.json'.")
        else:
            print(f"Error loading content for page {page_id}: {content_response.status_code}")

else:
    print(f"Error requesting the list of pages: {response.status_code}")
    print(response.text)

Функция очистки html страницы.

def clean_html(content: str) -> str:
    """
    Cleans HTML content by removing tags, scripts, styles, and unnecessary spaces.
    """
    # Parse HTML with BeautifulSoup
    soup = BeautifulSoup(content, 'html.parser')

    # Remove all script and style tags
    for script_or_style in soup(['script', 'style']):
        script_or_style.decompose()

    # Add spaces after block-level elements
    for tag in soup.find_all(['p', 'div', 'br', 'li', 'td', 'th', 'tr', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']):
        tag.insert_after(' ')  

    # Get clean text
    clean_text = soup.get_text()

    # Remove multiple spaces and line breaks
    clean_text = re.sub(r's+', ' ', clean_text)  # Replace multiple spaces with one
    clean_text = clean_text.strip()  # Remove extra spaces at the beginning and end

    return clean_text

Конвертер в векторную БД

Инициализируем ChromaDB (векторная БД) и embed модель. Создаем коллекцию в БД.

# Initialize the model for text vectorization
model = SentenceTransformer("all-MiniLM-L6-v2")

# Initialize ChromaDB
client = chromadb.PersistentClient(path="./chroma_db")

# Create collection
collection = client.get_or_create_collection(name="wiki_embeddings")

Embed-модель — это нейросеть, которая преобразует текст, изображение или другой тип данных в векторное представление (эмбеддинг). Эти векторы помогают находить похожие объекты с помощью поиска ближайших соседей.

По идее качество RAG системы будет зависеть в том числе и от выбранной RAG модели, но пока я этого не заметил. Использовал самую простую, но обязательно испытаю разные.

Функция обработки файлов json.

def process_files():
    """Processes JSON files, extracts data, and adds it to ChromaDB."""
    for filename in os.listdir(DATA_DIR):
        if filename.endswith(".json"):
            file_path = os.path.join(DATA_DIR, filename)
            with open(file_path, "r", encoding="utf-8") as file:
                data = json.load(file)
                
                if isinstance(data, list):
                    for idx, item in enumerate(data):
                        process_entry(item, f"{filename}_{idx}")
                else:
                    process_entry(data, filename)

Функция создание одной записи в БД.

def process_entry(data, doc_id):
    """Processes a single JSON entry."""
    title = data.get("title", "Untitled")
    content = data.get("content", "")
    
    if content.strip():  # Skip empty content
        full_text = f"{title}n{content}"  # Combine title and content
        embedding = model.encode(full_text).tolist()
        
        collection.add(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[full_text],
            metadatas=[{"title": title, "filename": doc_id}]
        )
        print(f"Document added: {doc_id}")
        print(f"TITLE >>>>>>", full_text.splitlines()[0])

Внимание: здесь я сознательно добавил title из wiki в content и создаю вектор на основе title + content. Иначе в RAG не хватало данных для ответа на вопросы по title.

Локальный генератор

Качаем модель huggingface, и пробуем использовать ее локально.

from huggingface_hub import hf_hub_download

hf_hub_download(repo_id="TheBloke/Mistral-7B-Instruct-v0.1-GGUF", 
                filename="mistral-7b-instruct-v0.1.Q4_K_M.gguf", 
                local_dir="models/")

Я выбрал скомпилированную модель на C++ из-за скорости. Не уверен, что эта модель поддерживает русский язык, у меня wiki на английском.

Делаем цикл вопрос-ответ в command line.

def main():
    model_path = "models/mistral-7b-instruct-v0.1.Q4_K_M.gguf"  # Path to the model
    chroma_client = chromadb.PersistentClient(path="./chroma_db")
    llm = load_llm(model_path)
    
    while True:
        query = input("Query: ")
        if query.lower() in ["exit", "quit"]:
            break
        
        context_docs = query_chroma(chroma_client, query)
        context = "n".join(context_docs) if context_docs else "Info not found."
        response = generate_response(llm, query, context)
        print("Answer:", response.strip())

Запрос в ChromaDB. Просим вернуть топ 5 (top_k=5) совпадающих вектором. Функция ниже выдаст Similarity score (% совпадения вектора с запросом), по котором можно отлаживать ответы модели.

def query_chroma(chroma_client, query_text, top_k=5):
    collection = chroma_client.get_collection("wiki_embeddings")  # Collection name
    results = collection.query(query_texts=[query_text], n_results=top_k)

    if results["documents"]:
        for doc, score in zip(results["documents"][0], results["distances"][0]):
            print(f"Similarity Score: {score:.4f}, Document: {doc}")
            print(".")
        return results["documents"][0]
    return []

Генерация ответа. Это важный код, тут идет вызов модели Генератора.

Мы передаем модели prompt, скармливая топ полученных векторов, схожих с запросом, в виде Context. Здесь достаточно безопасно применять внешние модели, не боясь, что они получат все ваши данные, потому что получат они только топ-векторы по конкретному запросу.

Важный параметр temperature=0. Ноль, потому что я хочу, чтобы модель не фантазировала, а работала как справочник. Если поставить 1, то модель будет максимально креативной, что не есть хорошо для RAG.

def generate_response(llm, query, context):
    prompt = f"""
    Use the following context to answer the question:
    {context}nnQuestion: {query}nAnswer:n"""
    return llm(prompt, max_tokens=512, temperature=0)["choices"][0]["text"]

Загрузка LLM. Грузим с максимальный context window (n_ctx=8196), чтобы вошло 5 векторов.

def load_llm(model_path):
    return Llama(model_path=model_path, n_ctx=8196, n_threads=8, log_level="OFF")

Локальный генератор работает достаточно медленно на MacBook Pro M3. Качества ответов на 3 из 5. Можно подобрать и скачать более подходящую модель. Можно любую модель потюнить параметрами.

Удаленный генератор

Пробуем Gemini от Google. ChromaDB здесь инициализируем как persistent client, иначе могут быть тонкости с получением того, что в кэше, вместо реального вектора.

import google.generativeai as genai

# Load Gemini API Key from environment variables
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

# Initialize Gemini client
genai.configure(api_key=GEMINI_API_KEY)

# Initialize ChromaDB
chroma_client = chromadb.PersistentClient(path="./chroma_db")

Генератор ответа. Я использовал gemini-2.0-flash. Она быстрая, точная и имеет большой лимит бесплатных токенов. Ответы на 5 из 5.

def generate_response(query, context):
    """Generate a response using Google's Gemini model."""
    prompt = f"""
    Use the following context to answer the question:
    {context}
    
    Question: {query}
    Answer:
    """
    
    model = genai.GenerativeModel("gemini-2.0-flash")
    response = model.generate_content(prompt)
    
    return response.text.strip()

Тот же самый генератор ответа с OpenAPI. Модель chatGPT-3.5-turbo. Работает медленнее, чем Gemini, и отвечает на 4 из 5. И бесплатных токенов мало. Но хотим купить токены для chatGPT-4.0. По качеству ответов отпишусь.

def generate_response(query, context):
    """Generate a response using OpenAI's GPT-4 model (OpenAI v1.0+ API)."""
    prompt = f"""
    Use the following context to answer the question:
    {context}
    
    Question: {query}
    Answer:
    """

    response = client.chat.completions.create(
        model="gpt-4o-mini",  # Change to "gpt-3.5-turbo" if needed
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.7
    )
    
    return response.choices[0].message.content.strip()

Потом прикручиваем Slack. Или любой другой чат, который нужен, и где есть API.

Как оно работает?

Первичное тестирование показало отличные результаты на Gemini. Быстро и точно отвечает. Если информации нет, так и говорит, что информация не найдена.

Бота в Slack назвали R2D2 . Когда я у него спросил: "Напиши список сотрудников", он написал всех и добавил себя в конце. Я не мог понять, откуда он себя взял?! Оказалось, на wiki был описан некий R2D2, тоже бот .

Wiki можно парсить overnight, для обновления векторной БД новыми знаниями.

Что еще нужно сделать?

1. Добавить знания из Slack

2. Добавить парсинг PDF

3. Попробовать дообучить модель на наших данных вместо подачи контекста модели-генератора (долго, дорого и не факт, что будет лучше)

4. Поробовать Grok-3 от Elon Musk

5. Попробовать embed+generator модель одного производителя (не факт, что будет лучше)

Что пока не будем делать?

Подключать фреймворк Langchain. Пока не понял всех его преимуществ.

Langchain — это фреймворк для разработки приложений с использованием языковых моделей, таких как GPT или другие, который предоставляет удобные инструменты для интеграции различных компонентов.

Если было интересно, дайте знать — напишу продолжение по следующим шагам.