golang

RAG-сервис на Go с Ламой

  • вторник, 10 июня 2025 г. в 00:00:08
https://habr.com/ru/articles/916900/

Привет, Хабр! Меня зовут Александр Белышев. Хочу поделиться своим недавним опытом работы над RAG‑сервисом и думаю, что эта тема может быть интересна другим коллегам.

У меня на работе возникла задача по реализации RAG (Retrieval‑Augmented Generation) сервиса. Хотя без моего участия эту задачу успешно решили на Python, изучая их код я задумался: а как можно сделать что‑то подобное на Go?
Результатом этого исследования и моей работы стало следующее решение...

Ссылочку на репозиторий оставлю тут для контекста https://github.com/xman12/rag-api , далее будут примеры из него.

Немного теории

Retrieval‑Augmented Generation (RAG) — это подход, который объединяет методы поиска релевантного контекста и генеративные модели ИИ. В рассматриваемом решении на языке Go используется база данных PostgreSQL с расширением pgvector для хранения и поиска векторных представлений документов. Сервис используется для обработки запросов, поиска по контексту и генерации ответов с помощью моделей векторизации и генерации.

Архитектура

Сервис состоит из нескольких ключевых компонентов:

  • Векторная база данных (Vector Database): Реализована с использованием PostgreSQL и расширения pgvector. Функции вставки и поиска документированных векторов обеспечивают быстрый поиск по сходству в больших наборах данных.

  • Встраивание векторов (Embeddings): Для преобразования текстовых данных в векторное представление используется backend‑решение, на основе Ollama.

  • Генерация ответов: Система использует генеративную модель (llama3) для создания ответов на основе объединённого контекста и исходного запроса.

  • REST API: Приложение использует фреймворк Gin для предоставления API, с маршрутами для сохранения данных и поиска.

Ниже показана упрощённая диаграмма архитектуры:

диаграмма архитектуры
диаграмма архитектуры

Реализация RAG на Go


Сервис написан на языке Go с использованием следующих технологий и библиотек:

  • Go и go.mod: Управление зависимостями и компиляция.

  • PostgreSQL + pgvector: Хранение векторов для поиска по семантическому сходству.

  • Gin: Фреймворк для реализации REST API.

Ключевые функции реализованы в файлах, таких как pkg/db/pgvector.go для взаимодействия с базой данных и internal/app/server.go для обработки запросов.

Немного разберем код добавления данных для поиска:

	headers := map[string]string{
		"Content-Type": "application/json",
	}

	embeddingBackend := backend.NewOllamaBackend(emdHost, embModel, time.Duration(20*time.Second))
	vectorDB, err := db.NewPGVector(databaseURL)
	if err != nil {
		log.Fatalf("Error initializing vector database: %v", err)
	}
	defer vectorDB.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	dataEmbedding, err := embeddingBackend.Embed(ctx, data.Value, headers)
	if err != nil {
		log.Fatalf("Error generating data embedding: %v", err)
	}
	log.Println("Vector embeddings generated")


	err = vectorDB.InsertDocument(ctx, data.Value, dataEmbedding)
	if err != nil {
		log.Fatalf("Error inserting data into vector database: %v", err)
	}

в embeddingBackend мы создаем подключение к ИИ, далее в dataEmbedding передаем данные из запроса, функция Embed отправляет запрос в ИИ в ответ получаем массив координат, которые уже передаем в БД в строке

err = vectorDB.InsertDocument(ctx, data.Value, dataEmbedding)

Код для поиска:

var (
	host        = "http://localhost:11434"
	emdHost     = "http://localhost:11434"
	embModel    = "mxbai-embed-large"
	genModel    = "llama3"
	databaseURL = "postgres://user:password@localhost:5432/dbname?sslmode=disable"
)

func handleSearch(c *gin.Context) {

	var query Query
	if err := c.ShouldBindUri(&query); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error": "Invalid query parameter",
		})
		return
	}

	headers := map[string]string{
		"Content-Type": "application/json",
	}

	embeddingBackend := backend.NewOllamaBackend(emdHost, embModel, time.Duration(20*time.Second))
	generationBackend := backend.NewOllamaBackend(host, genModel, time.Duration(20*time.Second))

	vectorDB, err := db.NewPGVector(databaseURL)
	if err != nil {
		log.Fatalf("Error initializing vector database: %v", err)
	}
	defer vectorDB.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	defer vectorDB.Close()

	queryEmbedding, err := embeddingBackend.Embed(ctx, query.Value, headers)
	if err != nil {
		log.Fatalf("Error generating query embedding: %v", err)
	}
	log.Println("Vector embeddings generated")

	retrievedDocs, err := vectorDB.QueryRelevantDocuments(ctx, queryEmbedding, "ollama")
	if err != nil {
		log.Fatalf("Error retrieving relevant documents: %v", err)
	}

	augmentedQuery := db.CombineQueryWithContext(query.Value, retrievedDocs)

	prompt := backend.NewPrompt().
		AddMessage("system", "Вы — помощник ИИ. Используйте предоставленный контекст, чтобы ответить на вопрос пользователя как можно точнее. Не отвечает на вопросы не загруженного контекста. Пиши только по русски. После ответа сбрасывай контекст общения. Если результат не найден, то верни текст 'Данные не найдены'").
		AddMessage("user", augmentedQuery).
		SetParameters(backend.Parameters{
			MaxTokens:   150, // Supported by LLaMa
			Temperature: 0.7, // Supported by LLaMa
			TopP:        0.9, // Supported by LLaMa
		})

	response, err := generationBackend.Generate(ctx, prompt)

	c.JSON(http.StatusOK, gin.H{
		"query":    query.Value,
		"response": response,
	})
}

Если не заострять внимание на деталях, то данный код работает по следующему принципу:

  • Поиск релевантных значений в БД

  • Передача релевантных значений в ИИ с промптом

  • Получение ответа от ИИ.

Область применения

Данный RAG‑сервис на Go может найти применение в следующих областях:

  • Системы поиска и рекомендации: Поиск релевантной информации на основе схожих векторных представлений.

  • Чат‑боты и виртуальные помощники: Обогащение запроса пользователя дополнительным контекстом для более точного ответа.

  • Анализ данных: Семантический разбор текстов и аналитика, где требуется высокая точность поиска по смысловому сходству.

Развертывание

Для того чтобы запустить сервис нужно иметь на борту:

  • Docker

  • Go (1.20 и выше)

  • Ollama

Чтобы запустить и ознакомиться с кодом склонируйте репозиторий

git clone git@github.com:xman12/rag-api.git

далее переходим в папку rag-api и выполняем go mod tidy.

После установки зависимостей нужно поднять БД. Можно воспользоваться докером или же локальной Postgres.
Если используем докер то достаточно выполнить команду:

docker-compose up -d

Если использовать локальную Postgres, тогда придется установить расширение pgvector

делается это так

CREATE EXTENSION IF NOT EXISTS vector;

После установки нужно создать таблицу и индекс к ней.

CREATE TABLE ollama_embeddings (
    id SERIAL PRIMARY KEY,
    doc_id TEXT NOT NULL,
    embedding VECTOR(1024),
    metadata JSONB
);

CREATE INDEX ollama_embeddings_idx ON ollama_embeddings
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

После этого отредактировать файл internal/app/server.go, изменить подключение к БД на свои, строка 18.

Ollama

Для работы с ИИ нужно установить ollama. После установки нужно установить две модели:

ollama pull mxbai-embed-large

ollama pull llama3

На этом настройка проекта завершена, можно запускать API go run cmd/api/main.go и пробовать щупать как все работает.
Для того, чтобы обогатить данными наш RAG сервис достаточно отправить запрос, пример:

Чтобы поискать, то делаем GET запрос

http://localhost:8080/search/кто возглавлял русское войско 9-го сентября 1380 года

Ответ:

{"query":"кто возглавлял русское войско 9-го сентября 1380 года","response":"Возглавлял русское войско Великий князь Дмитрий Донской."}

Немного итогов

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