From 4f5c71ffc0c6dcda4527eff680f7df8a6499955d Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil.hafeez@gmail.com>
Date: Wed, 28 May 2025 14:47:26 -0700
Subject: [PATCH] add support for openwebui

---
 .../brightstaff/src/handlers/chat_completions.rs   | 14 +++++++++-----
 crates/brightstaff/src/handlers/models.rs          |  6 ------
 demos/use_cases/llm_routing/arch_config.yaml       |  3 ++-
 demos/use_cases/llm_routing/docker-compose.yaml    |  4 ++--
 4 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/crates/brightstaff/src/handlers/chat_completions.rs b/crates/brightstaff/src/handlers/chat_completions.rs
index 50e65915..cb3094fc 100644
--- a/crates/brightstaff/src/handlers/chat_completions.rs
+++ b/crates/brightstaff/src/handlers/chat_completions.rs
@@ -3,7 +3,6 @@ use std::sync::Arc;
 use bytes::Bytes;
 use common::api::open_ai::ChatCompletionsRequest;
 use common::consts::ARCH_PROVIDER_HINT_HEADER;
-use common::utils::shorten_string;
 use http_body_util::combinators::BoxBody;
 use http_body_util::{BodyExt, Full, StreamBody};
 use hyper::body::Frame;
@@ -39,7 +38,7 @@ pub async fn chat_completions(
                 let v: Value = serde_json::from_slice(&chat_request_bytes).unwrap();
                 let err_msg = format!("Failed to parse request body: {}", err);
                 warn!("{}", err_msg);
-                warn!("request body: {}", v.to_string());
+                warn!("arch-router request body: {}", v.to_string());
                 let mut bad_request = Response::new(full(err_msg));
                 *bad_request.status_mut() = StatusCode::BAD_REQUEST;
                 return Ok(bad_request);
@@ -47,8 +46,8 @@ pub async fn chat_completions(
         };
 
     debug!(
-        "request body: {}",
-        shorten_string(&serde_json::to_string(&chat_completion_request).unwrap())
+        "arch-router request body: {}",
+        &serde_json::to_string(&chat_completion_request).unwrap()
     );
 
     let trace_parent = request_headers
@@ -56,7 +55,7 @@ pub async fn chat_completions(
         .find(|(ty, _)| ty.as_str() == "traceparent")
         .map(|(_, value)| value.to_str().unwrap_or_default().to_string());
 
-    let selected_llm = match router_service
+    let mut selected_llm = match router_service
         .determine_route(&chat_completion_request.messages, trace_parent.clone())
         .await
     {
@@ -69,6 +68,11 @@ pub async fn chat_completions(
         }
     };
 
+    if selected_llm.is_none() {
+        debug!("No LLM model selected, using default from request");
+        selected_llm = Some(chat_completion_request.model.clone());
+    }
+
     info!(
         "sending request to llm provider: {} with llm model: {:?}",
         llm_provider_endpoint, selected_llm
diff --git a/crates/brightstaff/src/handlers/models.rs b/crates/brightstaff/src/handlers/models.rs
index 8187292f..7f3427a0 100644
--- a/crates/brightstaff/src/handlers/models.rs
+++ b/crates/brightstaff/src/handlers/models.rs
@@ -2,7 +2,6 @@ use bytes::Bytes;
 use common::api::open_ai::Models;
 use common::configuration::LlmProvider;
 use http_body_util::{combinators::BoxBody, BodyExt, Full};
-use hyper::body::Body;
 use hyper::{Response, StatusCode};
 use serde_json;
 use std::sync::Arc;
@@ -12,11 +11,6 @@ pub async fn list_models(
 ) -> Response<BoxBody<Bytes, hyper::Error>> {
     let prov = llm_providers.clone();
     let providers = (*prov).clone();
-    let providers = providers
-        .iter()
-        .filter(|provider| provider.name == "gpt-4o-mini")
-        .cloned()
-        .collect::<Vec<_>>();
     let openai_models = Models::from(providers);
 
     match serde_json::to_string(&openai_models) {
diff --git a/demos/use_cases/llm_routing/arch_config.yaml b/demos/use_cases/llm_routing/arch_config.yaml
index d0074642..896fb795 100644
--- a/demos/use_cases/llm_routing/arch_config.yaml
+++ b/demos/use_cases/llm_routing/arch_config.yaml
@@ -8,6 +8,7 @@ listeners:
     timeout: 30s
 
 llm_providers:
+
   - name: gpt-4o-mini
     access_key: $OPENAI_API_KEY
     provider_interface: openai
@@ -17,6 +18,7 @@ llm_providers:
     access_key: $OPENAI_API_KEY
     provider_interface: openai
     model: gpt-4o
+    default: true
 
   - name: ministral-3b
     access_key: $MISTRAL_API_KEY
@@ -27,7 +29,6 @@ llm_providers:
     access_key: $ANTHROPY_API_KEY
     provider_interface: claude
     model: claude-3-7-sonnet-latest
-    default: true
 
   - name: claude-sonnet-4
     access_key: $ANTHROPY_API_KEY
diff --git a/demos/use_cases/llm_routing/docker-compose.yaml b/demos/use_cases/llm_routing/docker-compose.yaml
index fad02191..6c5d2b68 100644
--- a/demos/use_cases/llm_routing/docker-compose.yaml
+++ b/demos/use_cases/llm_routing/docker-compose.yaml
@@ -3,13 +3,13 @@ services:
 
   open-web-ui:
     image: ghcr.io/open-webui/open-webui:main
-    container_name: open-webui
     restart: always
     ports:
       - "8080:8080"
     environment:
       - DEFAULT_MODEL=gpt-4o-mini
-      - OPENAI_API_KEY=your_secret_key
+      - ENABLE_OPENAI_API=true
+      - OPENAI_API_BASE_URL=http://host.docker.internal:12000/v1
 
   jaeger:
     build: