Как я сделал 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)

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

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

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

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

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}\n\nQuestion: {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