Phase 3: AI integration with Ollama via Python sidecar
- sidecar: FastAPI app with /embed, /generate, /rerank hitting Ollama - sidecar: Dockerfile, env var config (EMBED_MODEL, GEN_MODEL, RERANK_MODEL) - aibridge: reqwest HTTP client with typed request/response structs - aibridge: Axum proxy endpoints (POST /ai/embed, /ai/generate, /ai/rerank) - gateway: wires AiClient with SIDECAR_URL env var - e2e verified: nomic-embed-text returns 768d vectors, qwen2.5 generates text Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -10,3 +10,4 @@ axum = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
|
||||
/// HTTP client for the Python AI sidecar.
|
||||
#[derive(Clone)]
|
||||
pub struct AiClient {
|
||||
client: Client,
|
||||
base_url: String,
|
||||
}
|
||||
|
||||
// -- Request/Response types --
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct EmbedRequest {
|
||||
pub texts: Vec<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize, Clone)]
|
||||
pub struct EmbedResponse {
|
||||
pub embeddings: Vec<Vec<f64>>,
|
||||
pub model: String,
|
||||
pub dimensions: usize,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct GenerateRequest {
|
||||
pub prompt: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub system: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_tokens: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize, Clone)]
|
||||
pub struct GenerateResponse {
|
||||
pub text: String,
|
||||
pub model: String,
|
||||
pub tokens_evaluated: Option<u64>,
|
||||
pub tokens_generated: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct RerankRequest {
|
||||
pub query: String,
|
||||
pub documents: Vec<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize, Clone)]
|
||||
pub struct ScoredDocument {
|
||||
pub index: usize,
|
||||
pub text: String,
|
||||
pub score: f64,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize, Clone)]
|
||||
pub struct RerankResponse {
|
||||
pub results: Vec<ScoredDocument>,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
impl AiClient {
|
||||
pub fn new(base_url: &str) -> Self {
|
||||
let client = Client::builder()
|
||||
.timeout(Duration::from_secs(120))
|
||||
.build()
|
||||
.expect("failed to build HTTP client");
|
||||
Self {
|
||||
client,
|
||||
base_url: base_url.trim_end_matches('/').to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn health(&self) -> Result<serde_json::Value, String> {
|
||||
let resp = self.client
|
||||
.get(format!("{}/health", self.base_url))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("sidecar unreachable: {e}"))?;
|
||||
resp.json().await.map_err(|e| format!("invalid response: {e}"))
|
||||
}
|
||||
|
||||
pub async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, String> {
|
||||
let resp = self.client
|
||||
.post(format!("{}/embed", self.base_url))
|
||||
.json(&req)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("embed request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("embed error ({}): {text}", text.len()));
|
||||
}
|
||||
resp.json().await.map_err(|e| format!("embed parse error: {e}"))
|
||||
}
|
||||
|
||||
pub async fn generate(&self, req: GenerateRequest) -> Result<GenerateResponse, String> {
|
||||
let resp = self.client
|
||||
.post(format!("{}/generate", self.base_url))
|
||||
.json(&req)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("generate request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("generate error: {text}"));
|
||||
}
|
||||
resp.json().await.map_err(|e| format!("generate parse error: {e}"))
|
||||
}
|
||||
|
||||
pub async fn rerank(&self, req: RerankRequest) -> Result<RerankResponse, String> {
|
||||
let resp = self.client
|
||||
.post(format!("{}/rerank", self.base_url))
|
||||
.json(&req)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("rerank request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("rerank error: {text}"));
|
||||
}
|
||||
resp.json().await.map_err(|e| format!("rerank parse error: {e}"))
|
||||
}
|
||||
}
|
||||
@@ -1 +1,2 @@
|
||||
pub mod client;
|
||||
pub mod service;
|
||||
|
||||
@@ -1,9 +1,60 @@
|
||||
use axum::{Router, routing::get};
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::State,
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
};
|
||||
|
||||
pub fn router() -> Router {
|
||||
Router::new().route("/health", get(health))
|
||||
use crate::client::{
|
||||
AiClient, EmbedRequest, GenerateRequest, RerankRequest,
|
||||
};
|
||||
|
||||
pub fn router(client: AiClient) -> Router {
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/embed", post(embed))
|
||||
.route("/generate", post(generate))
|
||||
.route("/rerank", post(rerank))
|
||||
.with_state(client)
|
||||
}
|
||||
|
||||
async fn health() -> &'static str {
|
||||
"aibridge ok"
|
||||
async fn health(State(client): State<AiClient>) -> impl IntoResponse {
|
||||
match client.health().await {
|
||||
Ok(info) => Ok(Json(info)),
|
||||
Err(e) => Err((StatusCode::BAD_GATEWAY, format!("sidecar down: {e}"))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn embed(
|
||||
State(client): State<AiClient>,
|
||||
Json(req): Json<EmbedRequest>,
|
||||
) -> impl IntoResponse {
|
||||
tracing::info!("embedding {} texts", req.texts.len());
|
||||
match client.embed(req).await {
|
||||
Ok(resp) => Ok(Json(resp)),
|
||||
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn generate(
|
||||
State(client): State<AiClient>,
|
||||
Json(req): Json<GenerateRequest>,
|
||||
) -> impl IntoResponse {
|
||||
tracing::info!("generating with prompt len={}", req.prompt.len());
|
||||
match client.generate(req).await {
|
||||
Ok(resp) => Ok(Json(resp)),
|
||||
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn rerank(
|
||||
State(client): State<AiClient>,
|
||||
Json(req): Json<RerankRequest>,
|
||||
) -> impl IntoResponse {
|
||||
tracing::info!("reranking {} documents", req.documents.len());
|
||||
match client.rerank(req).await {
|
||||
Ok(resp) => Ok(Json(resp)),
|
||||
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,12 +22,16 @@ async fn main() {
|
||||
// Query engine — DataFusion over catalog-registered Parquet
|
||||
let engine = queryd::context::QueryEngine::new(registry.clone(), store.clone());
|
||||
|
||||
// AI sidecar client
|
||||
let sidecar_url = std::env::var("SIDECAR_URL").unwrap_or_else(|_| "http://localhost:3200".to_string());
|
||||
let ai_client = aibridge::client::AiClient::new(&sidecar_url);
|
||||
|
||||
let app = Router::new()
|
||||
.route("/health", get(health))
|
||||
.nest("/storage", storaged::service::router(store))
|
||||
.nest("/catalog", catalogd::service::router(registry))
|
||||
.nest("/query", queryd::service::router(engine))
|
||||
.nest("/ai", aibridge::service::router())
|
||||
.nest("/ai", aibridge::service::router(ai_client))
|
||||
.layer(TraceLayer::new_for_http());
|
||||
|
||||
let addr = "0.0.0.0:3100";
|
||||
|
||||
Reference in New Issue
Block a user