plano/crates/brightstaff/src/main.rs

use brightstaff::handlers::agent_chat_completions::agent_chat;
use brightstaff::handlers::chat_completions::chat;
use brightstaff::handlers::models::list_models;
use brightstaff::router::llm_router::RouterService;
use brightstaff::utils::tracing::init_tracer;
use bytes::Bytes;
use common::configuration::Configuration;
use common::consts::{CHAT_COMPLETIONS_PATH, MESSAGES_PATH};
use http_body_util::{combinators::BoxBody, BodyExt, Empty};
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use opentelemetry::trace::FutureExt;
use opentelemetry::{global, Context};
use opentelemetry_http::HeaderExtractor;
use std::sync::Arc;
use std::{env, fs};
use tokio::net::TcpListener;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

pub mod router;

const BIND_ADDRESS: &str = "0.0.0.0:9091";
const DEFAULT_ROUTING_LLM_PROVIDER: &str = "arch-router";
const DEFAULT_ROUTING_MODEL_NAME: &str = "Arch-Router";

// Utility function to extract the context from the incoming request headers
fn extract_context_from_request(req: &Request<Incoming>) -> Context {
    global::get_text_map_propagator(|propagator| {
        propagator.extract(&HeaderExtractor(req.headers()))
    })
}

fn empty() -> BoxBody<Bytes, hyper::Error> {
    Empty::<Bytes>::new()
        .map_err(|never| match never {})
        .boxed()
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let _tracer_provider = init_tracer();
    let bind_address = env::var("BIND_ADDRESS").unwrap_or_else(|_| BIND_ADDRESS.to_string());

    info!(
        "current working directory: {}",
        env::current_dir().unwrap().display()
    );
    // loading arch_config.yaml file
    let arch_config_path = env::var("ARCH_CONFIG_PATH_RENDERED")
        .unwrap_or_else(|_| "./arch_config_rendered.yaml".to_string());
    info!("Loading arch_config.yaml from {}", arch_config_path);

    let config_contents =
        fs::read_to_string(&arch_config_path).expect("Failed to read arch_config.yaml");

    let config: Configuration =
        serde_yaml::from_str(&config_contents).expect("Failed to parse arch_config.yaml");

    let arch_config = Arc::new(config);

    let llm_providers = Arc::new(RwLock::new(arch_config.model_providers.clone()));
    let agents_list = Arc::new(RwLock::new(arch_config.agents.clone()));
    let listeners = Arc::new(RwLock::new(arch_config.listeners.clone()));

    debug!(
        "arch_config: {:?}",
        &serde_json::to_string(arch_config.as_ref()).unwrap()
    );

    let llm_provider_url =
        env::var("LLM_PROVIDER_ENDPOINT").unwrap_or_else(|_| "http://localhost:12001".to_string());

    info!("llm provider url: {}", llm_provider_url);
    info!("listening on http://{}", bind_address);
    let listener = TcpListener::bind(bind_address).await?;

    let routing_model_name: String = arch_config
        .routing
        .as_ref()
        .and_then(|r| r.model.clone())
        .unwrap_or_else(|| DEFAULT_ROUTING_MODEL_NAME.to_string());

    let routing_llm_provider = arch_config
        .routing
        .as_ref()
        .and_then(|r| r.model_provider.clone())
        .unwrap_or_else(|| DEFAULT_ROUTING_LLM_PROVIDER.to_string());

    let router_service: Arc<RouterService> = Arc::new(RouterService::new(
        arch_config.model_providers.clone(),
        llm_provider_url.clone() + CHAT_COMPLETIONS_PATH,
        routing_model_name,
        routing_llm_provider,
    ));

    let model_aliases = Arc::new(arch_config.model_aliases.clone());

    loop {
        let (stream, _) = listener.accept().await?;
        let peer_addr = stream.peer_addr()?;
        let io = TokioIo::new(stream);

        let router_service: Arc<RouterService> = Arc::clone(&router_service);
        let model_aliases = Arc::clone(&model_aliases);
        let llm_provider_url = llm_provider_url.clone();

        let llm_providers = llm_providers.clone();
        let agents_list = agents_list.clone();
        let listeners = listeners.clone();
        let service = service_fn(move |req| {
            let router_service = Arc::clone(&router_service);
            let parent_cx = extract_context_from_request(&req);
            let llm_provider_url = llm_provider_url.clone();
            let llm_providers = llm_providers.clone();
            let model_aliases = Arc::clone(&model_aliases);
            let agents_list = agents_list.clone();
            let listeners = listeners.clone();

            async move {
                match (req.method(), req.uri().path()) {
                    (&Method::POST, CHAT_COMPLETIONS_PATH | MESSAGES_PATH) => {
                        let fully_qualified_url =
                            format!("{}{}", llm_provider_url, req.uri().path());
                        chat(req, router_service, fully_qualified_url, model_aliases)
                            .with_context(parent_cx)
                            .await
                    }
                    (&Method::POST, "/agents/v1/chat/completions") => {
                        let fully_qualified_url =
                            format!("{}{}", llm_provider_url, req.uri().path());
                        agent_chat(
                            req,
                            router_service,
                            fully_qualified_url,
                            agents_list,
                            listeners,
                        )
                        .with_context(parent_cx)
                        .await
                    }
                    (&Method::GET, "/v1/models" | "/agents/v1/models") => {
                        Ok(list_models(llm_providers).await)
                    }
                    // hack for now to get openw-web-ui to work
                    (&Method::OPTIONS, "/v1/models" | "/agents/v1/models") => {
                        let mut response = Response::new(empty());
                        *response.status_mut() = StatusCode::NO_CONTENT;
                        response
                            .headers_mut()
                            .insert("Allow", "GET, OPTIONS".parse().unwrap());
                        response
                            .headers_mut()
                            .insert("Access-Control-Allow-Origin", "*".parse().unwrap());
                        response.headers_mut().insert(
                            "Access-Control-Allow-Headers",
                            "Authorization, Content-Type".parse().unwrap(),
                        );
                        response.headers_mut().insert(
                            "Access-Control-Allow-Methods",
                            "GET, POST, OPTIONS".parse().unwrap(),
                        );
                        response
                            .headers_mut()
                            .insert("Content-Type", "application/json".parse().unwrap());

                        Ok(response)
                    }
                    _ => {
                        debug!("No route for {} {}", req.method(), req.uri().path());
                        let mut not_found = Response::new(empty());
                        *not_found.status_mut() = StatusCode::NOT_FOUND;
                        Ok(not_found)
                    }
                }
            }
        });

        tokio::task::spawn(async move {
            debug!("Accepted connection from {:?}", peer_addr);
            if let Err(err) = http1::Builder::new()
                // .serve_connection(io, service_fn(chat_completion))
                .serve_connection(io, service)
                .await
            {
                warn!("Error serving connection: {:?}", err);
            }
        });
    }
}