65 lines
1.8 KiB
Rust
65 lines
1.8 KiB
Rust
|
use anyhow::{anyhow, Result};
|
||
|
|
use pgvector::Vector;
|
||
|
|
use reqwest::Client;
|
||
|
|
use serde::{Deserialize, Serialize};
|
||
|
|
use std::env;
|
||
|
|
|
||
|
|
/// Embeddings client for Ollama
|
||
|
|
#[derive(Clone)]
|
||
|
|
pub struct EmbeddingsClient {
|
||
|
|
base_url: String,
|
||
|
|
model: String,
|
||
|
|
#[allow(dead_code)]
|
||
|
|
http: Client,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Debug, Serialize)]
|
||
|
|
struct EmbeddingRequest {
|
||
|
|
model: String,
|
||
|
|
input: Vec<String>,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Debug, Deserialize)]
|
||
|
|
struct EmbeddingResponse {
|
||
|
|
embeddings: Vec<Vec<f32>>,
|
||
|
|
model: String,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl EmbeddingsClient {
|
||
|
|
/// Create from environment
|
||
|
|
/// Uses api.riotpiao.com gateway (nomic-ai/nomic-embed-text-v2-moe model)
|
||
|
|
pub fn from_env() -> Result<Self> {
|
||
|
|
let base_url = env::var("LLM_API_BASE").unwrap_or_else(|_| "https://api.riotpiao.com".to_string());
|
||
|
|
let model = "nomic-ai/nomic-embed-text-v2-moe".to_string();
|
||
|
|
|
||
|
|
Ok(Self {
|
||
|
|
base_url,
|
||
|
|
model,
|
||
|
|
http: Client::new(),
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Embed a single text string
|
||
|
|
pub async fn embed(&self, text: &str) -> Result<Vector> {
|
||
|
|
let embeddings = self.embed_batch(&[text.to_string()]).await?;
|
||
|
|
Ok(embeddings.into_iter().next().ok_or_else(|| anyhow::anyhow!("empty embedding response"))?)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Embed multiple texts in a batch using api.riotpiao.com gateway
|
||
|
|
pub async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vector>> {
|
||
|
|
let req = EmbeddingRequest {
|
||
|
|
model: self.model.clone(),
|
||
|
|
input: texts.to_vec(),
|
||
|
|
};
|
||
|
|
|
||
|
|
let url = format!("{}/v1/embeddings", self.base_url);
|
||
|
|
let resp: EmbeddingResponse = self.http.post(&url).json(&req).send().await?.json().await?;
|
||
|
|
|
||
|
|
Ok(resp
|
||
|
|
.embeddings
|
||
|
|
.into_iter()
|
||
|
|
.map(Vector::from)
|
||
|
|
.collect())
|
||
|
|
}
|
||
|
|
}
|