feat: dynamic pipeline routing

Cargo.toml

@@ -15,7 +15,7 @@ axum-extra = "0.9.4"
 tracing = "0.1"
 tracing-subscriber = "0.3"
 serde_yaml = "0.9"
-tower = "0.5.1"
+tower = { version = "0.5.1", features = ["full"]}
 anyhow = "1.0.92"
 tower-http = { version = "0.6.1", features = ["full"]}
 uuid = "1.11.0"

config-example.yaml

@@ -27,7 +27,6 @@ pipelines:
       - tracing:
           enabled: true
           endpoint: "https://api.traceloop.com/v1/traces"
-          api_key: "<your-traceloop-api-key>"
       - model-router:
           models:
             - gpt-4o-azure

src/ai_models/instance.rs

@@ -3,7 +3,6 @@ use crate::models::chat::{ChatCompletionRequest, ChatCompletionResponse};
 use crate::models::completion::{CompletionRequest, CompletionResponse};
 use crate::models::embeddings::{EmbeddingsRequest, EmbeddingsResponse};
 use crate::providers::provider::Provider;
-use crate::state::AppState;
 use axum::http::StatusCode;
 use std::sync::Arc;
 
@@ -17,33 +16,26 @@ pub struct ModelInstance {
 impl ModelInstance {
     pub async fn chat_completions(
         &self,
-        state: Arc<AppState>,
         mut payload: ChatCompletionRequest,
     ) -> Result<ChatCompletionResponse, StatusCode> {
         payload.model = self.model_type.clone();
-        self.provider
-            .chat_completions(state, payload, &self.config)
-            .await
+        self.provider.chat_completions(payload, &self.config).await
     }
 
     pub async fn completions(
         &self,
-        state: Arc<AppState>,
         mut payload: CompletionRequest,
     ) -> Result<CompletionResponse, StatusCode> {
         payload.model = self.model_type.clone();
 
-        self.provider
-            .completions(state, payload, &self.config)
-            .await
+        self.provider.completions(payload, &self.config).await
     }
 
     pub async fn embeddings(
         &self,
-        state: Arc<AppState>,
         mut payload: EmbeddingsRequest,
     ) -> Result<EmbeddingsResponse, StatusCode> {
         payload.model = self.model_type.clone();
-        self.provider.embeddings(state, payload, &self.config).await
+        self.provider.embeddings(payload, &self.config).await
     }
 }

src/ai_models/registry.rs

@@ -6,6 +6,7 @@ use super::instance::ModelInstance;
 use crate::config::models::ModelConfig;
 use crate::providers::registry::ProviderRegistry;
 
+#[derive(Clone)]
 pub struct ModelRegistry {
     models: HashMap<String, Arc<ModelInstance>>,
 }

src/config/models.rs

@@ -39,38 +39,26 @@ pub enum PipelineType {
 pub struct Pipeline {
     pub name: String,
     pub r#type: PipelineType,
+    #[serde(with = "serde_yaml::with::singleton_map_recursive")]
     pub plugins: Vec<PluginConfig>,
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone)]
-#[serde(untagged)]
+#[serde(rename_all = "kebab-case")]
 pub enum PluginConfig {
     Logging {
-        logging: LoggingConfig,
+        #[serde(default = "default_log_level")]
+        level: String,
     },
     Tracing {
-        tracing: TracingConfig,
+        endpoint: String,
+        api_key: String,
     },
     ModelRouter {
-        #[serde(rename = "model-router")]
-        model_router: ModelRouterConfig,
+        models: Vec<String>,
     },
 }
 
-#[derive(Debug, Deserialize, Serialize, Clone)]
-pub struct LoggingConfig {
-    pub enabled: bool,
-    pub level: String,
-}
-
-#[derive(Debug, Deserialize, Serialize, Clone)]
-pub struct TracingConfig {
-    pub enabled: bool,
-    pub endpoint: String,
-    pub api_key: String,
-}
-
-#[derive(Debug, Deserialize, Serialize, Clone)]
-pub struct ModelRouterConfig {
-    pub models: Vec<String>,
+fn default_log_level() -> String {
+    "warning".to_string()
 }

src/lib.rs

@@ -1,6 +1,5 @@
 pub mod ai_models;
 pub mod config;
-pub mod handlers;
 pub mod models;
 pub mod pipelines;
 pub mod providers;

src/main.rs

@@ -1,11 +1,10 @@
 use gateway::{config::lib::load_config, routes, state::AppState};
 use std::sync::Arc;
 use tracing::info;
-use tracing_subscriber::fmt::Subscriber;
 
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
-    Subscriber::builder().init();
+    tracing_subscriber::fmt::init();
 
     info!("Starting the application...");
 
@@ -21,7 +20,6 @@ async fn main() -> Result<(), anyhow::Error> {
         .unwrap();
 
     info!("Server is running on port {}", port);
-
     axum::serve(listener, app).await.unwrap();
 
     Ok(())

src/pipelines/pipeline.rs

@@ -1,35 +1,93 @@
-use crate::config::models::{Pipeline, PipelineType};
-use axum::http::HeaderMap;
-
-pub fn select_pipeline<'a>(
-    pipelines: &'a [Pipeline],
-    pipeline_type: PipelineType,
-    headers: &HeaderMap,
-) -> Option<&'a Pipeline> {
-    // Filter pipelines by type
-    let matching_pipelines: Vec<&Pipeline> = pipelines
-        .iter()
-        .filter(|p| p.r#type == pipeline_type)
-        .collect();
-
-    if matching_pipelines.is_empty() {
-        return None;
+use std::sync::Arc;
+
+use crate::config::models::PipelineType;
+use crate::models::completion::{CompletionRequest, CompletionResponse};
+use crate::models::embeddings::{EmbeddingsRequest, EmbeddingsResponse};
+use crate::{
+    ai_models::registry::ModelRegistry,
+    config::models::{Pipeline, PluginConfig},
+    models::chat::{ChatCompletionRequest, ChatCompletionResponse},
+};
+use axum::{extract::State, http::StatusCode, routing::post, Json, Router};
+use tower_http::trace::TraceLayer;
+
+pub fn create_pipeline(pipeline: &Pipeline, model_registry: &ModelRegistry) -> Router {
+    let mut router = Router::new();
+
+    for plugin in pipeline.plugins.clone() {
+        router = match plugin {
+            PluginConfig::Tracing { .. } => router.layer(TraceLayer::new_for_http()),
+            PluginConfig::ModelRouter { models } => match pipeline.r#type {
+                PipelineType::Chat => router.route(
+                    "/chat/completions",
+                    post(move |state, payload| chat_completions(state, payload, models)),
+                ),
+                PipelineType::Completion => router.route(
+                    "/completions",
+                    post(move |state, payload| completions(state, payload, models)),
+                ),
+                PipelineType::Embeddings => router.route(
+                    "/embeddings",
+                    post(move |state, payload| embeddings(state, payload, models)),
+                ),
+            },
+            _ => router,
+        };
+    }
+
+    return router.with_state(Arc::new(model_registry.clone()));
+}
+
+pub async fn chat_completions(
+    State(model_registry): State<Arc<ModelRegistry>>,
+    Json(payload): Json<ChatCompletionRequest>,
+    model_keys: Vec<String>,
+) -> Result<Json<ChatCompletionResponse>, StatusCode> {
+    for model_key in model_keys {
+        let model = model_registry.get(&model_key).unwrap();
+
+        if payload.model == model.model_type {
+            let response = model.chat_completions(payload.clone()).await.unwrap();
+
+            return Ok(Json(response));
+        }
     }
 
-    if matching_pipelines.len() == 1 {
-        return Some(matching_pipelines[0]);
+    Err(StatusCode::SERVICE_UNAVAILABLE)
+}
+
+pub async fn completions(
+    State(model_registry): State<Arc<ModelRegistry>>,
+    Json(payload): Json<CompletionRequest>,
+    model_keys: Vec<String>,
+) -> Result<Json<CompletionResponse>, StatusCode> {
+    for model_key in model_keys {
+        let model = model_registry.get(&model_key).unwrap();
+
+        if payload.model == model.model_type {
+            let response = model.completions(payload.clone()).await.unwrap();
+
+            return Ok(Json(response));
+        }
     }
 
-    // Check for pipeline specification in headers
-    if let Some(pipeline_name) = headers
-        .get("x-traceloop-pipeline")
-        .and_then(|h| h.to_str().ok())
-    {
-        matching_pipelines
-            .into_iter()
-            .find(|p| p.name == pipeline_name)
-    } else {
-        // Default to pipeline named "default"
-        matching_pipelines.into_iter().find(|p| p.name == "default")
+    Err(StatusCode::SERVICE_UNAVAILABLE)
+}
+
+pub async fn embeddings(
+    State(model_registry): State<Arc<ModelRegistry>>,
+    Json(payload): Json<EmbeddingsRequest>,
+    model_keys: Vec<String>,
+) -> Result<Json<EmbeddingsResponse>, StatusCode> {
+    for model_key in model_keys {
+        let model = model_registry.get(&model_key).unwrap();
+
+        if payload.model == model.model_type {
+            let response = model.embeddings(payload.clone()).await.unwrap();
+
+            return Ok(Json(response));
+        }
     }
+
+    Err(StatusCode::SERVICE_UNAVAILABLE)
 }

src/pipelines/services/mod.rs

@@ -1 +1 @@
-pub mod model_router;
+