Предисловие
Прежде всего хочу сказать, что я не являюсь никаким специалистом, даже джуновского лвла, просто безработный студент, пишущий на коленке свои пет-проекты. И код, и тем более архитектура далеки от идеала. Однако, я думаю, некоторые моменты, о которых я буду рассказывать далее в статье, могут быть интересны полноценным разработчикам как бэкенда, так и ИИ-агентов.
RAG и Go
С технологией RAG я познакомился около года назад на хакатоне, посвященном обработке естественного языка. Там мы с командой разработали его простейшую имплементацию, с которой и заняли «почетное» 5-е место. Подробно об этой технологии в данной статье рассказывать я не буду, так как статья не о ней; вкратце - RAG позволяет генерировать ответы LLM на основании базы контекста, необходимый фрагмент которой вместе с запросом передается языковой модели на вход.
Шло время, мои навыки росли, я полностью пересел с Python на Go, начал интересоваться больше бэкенд-разработкой и думал какой бы пет-проект мне написать. Идея приложения, связанного с ИИ на Go кажется сперва странной: язык предназначен для совершенно других целей, отсутствуют хорошие библиотеки сообщества, вроде Langchain (langchaingo слишком слаба). И с одной стороны, если бы идеей было обучить собственную языковую модель - это было бы действительно глупо. Но чем больше я думал о разработке RAG и ИИ-агентов, тем больше понимал, что это чисто бэкенд задача, с нулем машинного обучения под капотом.
Эта мысль натолкнула меня на следующие рассуждения: зачем использовать низкопроизводительный Python, если можно создать более эффективное ИИ-приложение на Go, к тому же с лучшей масштабируемостью? Конечно, в основном это дело привычки и наличия в питоне необходимых библиотек, да и мало кто из-за небольшого прироста производительности пойдет переписывать всех ИИ-агентов на Go, Rust или C++. Но для меня это и стало хорошей идеей для своего странненького пет-проекта.
Архитектура
Архитектуру будущего приложения полностью определил язык, на котором оно будет написано. Кроме того, мне хотелось отточить свои навыки связывания микросервисов, а также попробовать применить брокеры сообщений.
Сперва я определил, какие сервисы должны находиться в приложении и как они должны быть связаны. Всего я выделил 5 отдельных сервисов:
API Gateway. Единая точка входа для всех внешних запросов. Маршрутизацию RESTful запросов.
User Service. Управление пользователями и доступом. Регистрация/аутентификация (JWT). Управление API-ключами. Организация коллекций документов.
Generation Service. Контекстно-зависимая генерация.
Splitter Service. Обработка массивов документов - расщепление массива на отдельные документы. Пул горутин для параллельной обработки.
Storage Service. Работа с векторными данными.
Загрузка данных
Первым сложным моментом для меня было разработать пайплайн загрузки данных. Идея была в том, что пользователи могут загруждать огромные JSON'ы с массивом, состоящим из тысяч отдельных документов - фрагментов контекста. Не думаю, что в конечном итоге я сделал правильно, но было принято решение поступить следующим образом:
Сначала поступивший JSON буферизируется в первый топик Kafka;
JSON забирает сервис разделения на документы, работая в несколько параллельных горутин расщепляет его и отсылает во второй топик Kafka;
Сервис хранилища, в котором запущена группа консьюмеров параллельно забирает по несколько документов и занимается их обработкой;
Не думаю, что данный подход был хоть немного эффективным. Узким горлом стал последний этап, где всего на одну локальную Ollama, которая должна была малой языковой моделью из текстовых документов генерировать их эмбеддинги (векторное представление текста), приходило по нескольку запросов единовременно, которые она не успела обрабатывать и возникали заторы. На этапе MVP я не стал думать над решением этого вопроса, однако в последующем, если я захочу дорабатывать приложение - это станет большой проблемой.
Генерация данных
Вторым сложным и очень интересным моментом стал пайплайн генерации данных. Для самой генерации я использовал сырой Langchaingo, у которого, все же, неплохая интеграция с Ollama.
Кроме того, порадовала фича стриминга токенов, которые генерирует LLM, такого в Langchain я не видел (может плохо смотрел). Суть в том, что LLM генерирует токены поочередно, вы это можете увидеть, написав запрос к любой модели в браузере: ответ будет генерироваться на ходу. Сам стриминг задается функцией, мне показалось лучшим вариантом реализовать функцию так, чтобы токены передавались через выходной канал. Код сервиса генерации представлен ниже.
func (gs *GenerationService) Generate(ctx context.Context, query string, collection string) (<-chan string, error) {
docs, err := gs.storage.Search(ctx, query, 2, collection)
if err != nil {
return nil, err
}
contextJSON, err := prepareContext(docs)
if err != nil {
return nil, err
}
out := make(chan string)
go func() {
gs.initGenerating(ctx, out, query, contextJSON)
}()
return out, nil
}
func (gs *GenerationService) initGenerating(ctx context.Context, out chan<- string, query string, contextJSON string) {
llm, err := ollama.New(ollama.WithModel(gs.model), ollama.WithServerURL(gs.ollamaAddress))
if err != nil {
log.Printf("generation failed: %v", err)
}
_, err = llms.GenerateFromSinglePrompt(
ctx,
llm,
fmt.Sprintf(defaultPrompt, contextJSON, query),
llms.WithTemperature(0.8),
llms.WithStreamingFunc(func(ctx context.Context, chunk []byte) error {
streamingFunc(ctx, out, chunk)
return nil
}),
)
if err != nil {
log.Printf("generation failed: %v", err)
}
defer close(out)
}
func streamingFunc(ctx context.Context, out chan<- string, chunk []byte) {
if ctx.Err() != nil {
return
}
out <- string(chunk)
}Следующим встал вопрос стриминга токенов уже с сервиса генерации на сервис оркестрации. Здесь думать пришлось недолго - я уже знал о существовании gRPC потоков, осталось реализовать, что, на самом деле, оказалось довольно просто. Логика передачи токенов в сам gRPC эндпоинт проста: вызывается сервисная фукнкция генерации, которая возвращает канал, через который и происходит стриминг. Далее - gRPC поток читает канал и отправляет токены.
func (gs *GenerationService) Generate(query *pb.Query, stream grpc.ServerStreamingServer[pb.ResponseChunk]) error {
ctx := stream.Context()
chunks, err := gs.service.Generate(ctx, query.Query, query.CollectionName)
if err != nil {
return fmt.Errorf("generation service: failed to generate answer: %v", err)
}
log.Printf("INFO: generator service new query: %s, starting stream...", query)
for c := range chunks {
err := stream.Send(&pb.ResponseChunk{Chunk: c})
if err != nil {
return fmt.Errorf("generation service: error sending message to stream: %v", err)
}
}
return nil
}На стороне клиента это выглядит так:
func (g *Generator) Generate(ctx context.Context, query, collection string) (<-chan string, error) {
out := make(chan string)
go func() {
stream, err := g.generator.Generate(ctx, &pb.Query{
Query: query,
CollectionName: collection,
})
if err != nil {
log.Printf("geteration client stream error: %v", err)
}
defer close(out)
for {
msg, err := stream.Recv()
if err != nil {
log.Printf("geteration client stream error: %v", err)
break
}
out <- msg.Chunk
}
}()
return out, nil
}И наконец, самое (для меня) интересное - а как же с сервиса оркестрации стримить токены конечному пользователю? Первой идеей было использовать вебсокеты, и, с одной стороны, это вроде логично, но быстро становится понятно, что это слишком большая работа для реализации исключительно одностороннего стриминга. Немного погуглив я узнал про то, как часто реализуются такие потоки на реальных LLM проектах - SSE. Если кто-то, как и я тогда, не в курсе, SSE - это технология, позволяющая отправлять непрерывные события от сервера к клиенту. По сути - небольшая надстройка над HTTP, описывающая начало и завершение потока, а также схему ивентов. Самое сложное было реализовать передачу событий на сервере, так как примеров, в том числе на фреймворке gin (который я использовал для http роутера) мало. Возможно, код вышел избыточным.
func (h *Handler) generateAnswer(c *gin.Context) {
c.Writer.Flush()
query := c.Query("query")
collection := c.Param("collection")
// получение потока токенов
stream, err := h.grpclient.Generator.Generate(c, query, collection)
if err != nil {
log.Printf("Generation error: %v", err)
c.SSEvent("error", "Failed to start generation")
c.Writer.Flush()
return
}
// потоковая передача
c.Stream(func(w io.Writer) bool {
select {
case <-c.Writer.CloseNotify():
log.Println("Client disconnected")
return false
case chunk, ok := <-stream:
if !ok {
c.SSEvent("end", "[DONE]")
log.Println("Stream finished")
return false
}
// отправление чанка данных
chunkWithNbsp := strings.TrimLeftFunc(chunk, func(r rune) bool { return r == ' ' })
leadingSpaces := len(chunk) - len(chunkWithNbsp)
spaces := strings.Repeat(" ", leadingSpaces)
finalChunk := spaces + chunkWithNbsp
c.SSEvent("message", finalChunk)
return true
}
})
}
Проблемой снова стала локальная Ollama, на которой была запущена LLM. Моя видеокарта попросу не вывозила одновременную генерацию для двух и более запросов. Решение этой проблемы на поверхности - либо использовать платные API генеративных моделей, по типу ChatGPT, либо поставить Ollama на дорогущий сервер с высокопроизводительными видеокартами.
Заключение
В заключение хочется сказать, что как мне кажется, не смотря на текущую популярность Python для задач разработки ИИ-агентов, в будущем, при повышении нагрузки на них, по крайней мере часть таких приложений придется переписывать на более производительные языки, а также усложнять их архитектуру. Важно понимать, что тот же Langchain - это не ML библиотека, написанная на C, как PyTorch и ожидать от нее внушительной производительности - не стоит. В то же время, Rust или C++ могут быть избыточными для таких задач, так что у Go, может быть, появится возможность занять эту нишу.
