RAG с LLM на Rust

Для чего нужна данная статья? :

Добавить документ в базу знаний LLM.

Зачем Вам это уметь? :

RAG — один из наиболее масштабируемых способов настройки базы знаний LLM для взаимодействия, специфичного для предметной области.

Обзор сервера RAG

Сервер RAG обычно включает в себя:

Внедрение вопроса.
Извлечение релевантных документов с помощью поиска по сходству.

Формирование ответа на основе контекста, предоставленного найденными документами.

Компоненты и инструменты

1. LLM API : предоставляет языковую модель для генерации ответов. (Например, можно использовать OpenAI API или Google Gemini API).
2. Модель встраивания : преобразует текст в векторные встраивания.
3. Векторная база данных : эффективно хранит и извлекает векторные вложения. Например, Weaviate или Pinecone.

Создание RAG-сервера в Rust

1. Настройка проекта Rust

Сначала создайте новый проект Rust:

cargo new rag_server cd rag_server

зависимости Cargo.toml:

[dependencies]

tokio = { version = "1", features = ["full"] }

reqwest = "0.11"

serde = { version = "1.0", features = ["derive"] }

serde_json = "1.0"

warp = "0.3"

dotenv = "0.15"

2. Создание HTTP-сервера

Используйте warp для создания REST API. Вот базовая настройка:

use warp::Filter; use serde::{Deserialize, Serialize};

use reqwest::Client; use std::env;

#[derive(Debug, Serialize, Deserialize)] struct Document { text: String, }

#[derive(Debug, Serialize, Deserialize)] struct QueryRequest {

query: String, } #[derive(Debug, Serialize, Deserialize)]

struct QueryResponse {

answer: String,

}

async fn add_documents_handler(documents: Vec<Document>) -> Result<impl warp::Reply, warp::Rejection> {

// Logic for embedding and storing documents

// Example: Call an embedding model and a vector database

Ok(warp::reply::with_status("Documents added", warp::http::StatusCode::OK))

}

async fn query_handler(query: QueryRequest) -> Result<impl warp::Reply,

warp::Rejection> {

// Logic for querying and generating response

// Example: Embed query, retrieve relevant documents, call LLM

Ok(warp::reply::json(&QueryResponse {

answer: "Example answer".to_string() })) }

#[tokio::main] async fn main() {

dotenv::dotenv().ok();

let addr = "0.0.0.0:8080";

let add = warp::path("add")

.and(warp::post())

.and(warp::body::json())

.and_then(add_documents_handler);

let query = warp::path("query")

.and(warp::post())

.and(warp::body::json())

.and_then(query_handler);

let routes = add.or(query);

warp::serve(routes).run(addr.parse().unwrap()).await;

}

3. Взаимодействие с внешними службами

Вот как можно интегрироваться с внешним API для встраивания и LLM:

async fn get_embeddings(text: &str) -> Result<Vec<f32>,

reqwest::Error> {

let api_url =

env::var("EMBEDDING_API_URL").expect("EMBEDDING_API_URL must be set");

let api_key =

env::var("EMBEDDING_API_KEY").expect("EMBEDDING_API_KEY must be set");

let client = Client::new();

let response = client.post(&api_url)

.header("Authorization", format!("Bearer {}", api_key))

.json(&serde_json::json!({ "text": text }))

.send()

.await?;

let embeddings: Vec<f32> = response.json().await?;

Ok(embeddings)

}

async fn query_llm(question: &str) -> Result<String, reqwest::Error> {

let api_url = env::var("LLM_API_URL").expect("LLM_API_URL must be set");

let api_key = env::var("LLM_API_KEY").expect("LLM_API_KEY must be

set");

let client = Client::new();

let response = client.post(&api_url) .header("Authorization", format!

("Bearer {}", api_key)) .json(&serde_json::json!({ "query": question }))

.send() .await?; let answer: serde_json::Value =

response.json().await?;

Ok(answer["answer"].as_str().unwrap_or("No answer").to_string())

}

4. Интеграция модели внедрения и векторной базы данных

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

async fn add_to_vector_db(doc_id: &str, embedding: &[f32]) -> Result<(),

reqwest::Error> {

let api_url = env::var("VECTOR_DB_URL").expect("VECTOR_DB_URL must be set");

let api_key = env::var("VECTOR_DB_API_KEY").expect("VECTOR_DB_API_KEY must be set");

let client = Client::new();

client.post(format!("{}/documents", api_url))

.header("Authorization", format!("Bearer {}", api_key))

.json(&serde_json::json!({ "id": doc_id, "embedding": embedding }))

.send()

.await?;

Ok(())

}