profit/lakehouse
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Phase 40: Routing Engine + Policy

Browse Source 
- RoutingEngine with RouteDecision (model_pattern → provider)
- config/routing.toml: rules, fallback chain, cost gating
- Per-provider Usage tracking in /v1/usage response
- 12 gateway tests green
This commit is contained in:
root
2026-04-23 02:36:45 -05:00
parent e27a17e950
commit 55f8e0fe6e
5 changed files with 191 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/gateway/src/v1/mod.rs
+30 -9
View File
@@ -24,7 +24,7 @@ use axum::{
Json,
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::{collections::HashMap, sync::Arc};
use tokio::sync::RwLock;
#[derive(Clone)]
@@ -46,6 +46,16 @@ pub struct Usage {
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
#[serde(default)]
pub by_provider: std::collections::HashMap<String, ProviderUsage>,
}
#[derive(Default, Clone, Serialize)]
pub struct ProviderUsage {
pub requests: u64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
}
pub fn router(state: V1State) -> Router {
@@ -135,18 +145,22 @@ async fn chat(
let start_time = chrono::Utc::now();
let start_instant = std::time::Instant::now();
let resp = match provider.as_str() {
"ollama" | "local" | "" => ollama::chat(&state.ai_client, &req)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama local: {e}")))?,
let (resp, used_provider) = match provider.as_str() {
"ollama" | "local" | "" => {
let r = ollama::chat(&state.ai_client, &req)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama local: {e}")))?;
(r, "ollama".to_string())
}
"ollama_cloud" | "cloud" => {
let key = state.ollama_cloud_key.as_deref().ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"OLLAMA_CLOUD_KEY not configured — set env or populate /root/llm_team_config.json".to_string(),
"OLLAMA_CLOUD_KEY not configured".to_string(),
))?;
ollama_cloud::chat(key, &req)
let r = ollama_cloud::chat(key, &req)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama cloud: {e}")))?
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama cloud: {e}")))?;
(r, "ollama_cloud".to_string())
}
other => {
return Err((
@@ -168,7 +182,7 @@ async fn chat(
.map(|c| c.message.content.clone())
.unwrap_or_default();
lf.emit_chat(langfuse_trace::ChatTrace {
provider: provider.clone(),
provider: used_provider.clone(),
model: resp.model.clone(),
input: req.messages.clone(),
output,
@@ -183,12 +197,19 @@ async fn chat(
});
}
// Phase 40: per-provider usage tracking
{
let mut u = state.usage.write().await;
u.requests += 1;
u.prompt_tokens += resp.usage.prompt_tokens as u64;
u.completion_tokens += resp.usage.completion_tokens as u64;
u.total_tokens += resp.usage.total_tokens as u64;
let provider_usage = u.by_provider.entry(used_provider).or_default();
provider_usage.requests += 1;
provider_usage.prompt_tokens += resp.usage.prompt_tokens as u64;
provider_usage.completion_tokens += resp.usage.completion_tokens as u64;
provider_usage.total_tokens += resp.usage.total_tokens as u64;
}
Ok(Json(resp))