diff --git a/Cargo.lock b/Cargo.lock index 97cb2b6..441b7b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -360,6 +360,7 @@ dependencies = [ "axum", "axum-extra", "clap", + "headers", "reqwest", "schemars", "secrecy", diff --git a/Cargo.toml b/Cargo.toml index 7139e4c..8779b20 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ edition = "2024" axum = { version = "0.8.9", features = ["http2", "multipart", "ws"] } axum-extra = { version = "0.12.6", features = ["typed-header", "typed-routing"] } clap = { version = "4.6.7", features = ["derive", "env"] } +headers = "0.4.2" reqwest = { version = "0.13.5", features = ["brotli", "cookies", "deflate", "form", "gzip", "json", "multipart", "query", "stream", "zstd"] } schemars = "1.2.2" secrecy = { version = "0.10.3", features = ["serde"] } diff --git a/src/main.rs b/src/main.rs index c280d54..0982222 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,20 @@ +use std::net::{IpAddr, SocketAddr}; + use clap::Parser; use secrecy::SecretString; use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan}; mod openai; +mod server; #[derive(Parser, Debug)] struct AppArgs { + #[arg(long, env)] + ip: IpAddr, + + #[arg(long, env)] + port: u16, + #[arg(long)] cloudflare_workers_api_key: Vec, @@ -48,12 +57,16 @@ async fn main() { .with_env_filter(env_filter) .init(); tracing::debug!(?app_args, "using"); - let AppArgs { - cloudflare_workers_api_key, - google_ai_studio_api_key, - groq_api_key, - hugging_face_api_key, - mistral_api_key, - openrouter_api_key, - } = app_args; + + let addr = SocketAddr::new(app_args.ip, app_args.port); + + let app = server::create_app( + server::AppState::builder() + .app_args(app_args.into()) + .build(), + ); + + let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO"); + + axum::serve(listener, app).await.expect("TODO"); } diff --git a/src/openai.rs b/src/openai.rs index bb8b467..2474a26 100644 --- a/src/openai.rs +++ b/src/openai.rs @@ -1,81 +1,44 @@ -use std::marker::PhantomData; +use schemars::Schema; +use serde::{Deserialize, Serialize}; -use schemars::{JsonSchema, Schema, SchemaGenerator}; -use serde::{Deserialize, Serialize, ser::SerializeMap as _}; -use typed_builder::TypedBuilder; -use url::Url; - -#[derive(Debug, TypedBuilder, Serialize, Deserialize)] -struct Message { +#[derive(Debug, Serialize, Deserialize)] +pub struct Message { role: String, content: String, } -struct StructuredOutputJson(PhantomData); - -impl Serialize for StructuredOutputJson -where - S: JsonSchema, -{ - fn serialize( - &self, - serializer: Serializer, - ) -> Result - where - Serializer: serde::Serializer, - { - let mut map = serializer.serialize_map(Some(3))?; - map.serialize_entry("strict", &true)?; - map.serialize_entry("name", &S::schema_name())?; - map.serialize_entry( - "schema", - &SchemaGenerator::default().into_root_schema_for::(), - )?; - map.end() - } +#[derive(Debug, Deserialize, Serialize)] +pub struct ResponseFormatJsonSchema { + pub strict: bool, + pub name: String, + pub schema: Schema, } -#[derive(Debug)] -enum ResponseFormat { - Json { schema: PhantomData }, +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "type")] +pub enum ResponseFormat { + JsonSchema { + json_schema: ResponseFormatJsonSchema, + }, } -impl Serialize for ResponseFormat -where - OutputFormat: JsonSchema, -{ - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - let mut map = serializer.serialize_map(Some(2))?; - map.serialize_entry("type", "json_schema")?; - map.serialize_entry( - "json_schema", - &StructuredOutputJson::(Default::default()), - )?; - map.end() - } +#[derive(Debug, Serialize, Deserialize)] +pub struct ChatCompletionsRequest { + pub model: Model, + + pub messages: Vec, + + pub response_format: ResponseFormat, } -#[derive(Debug, TypedBuilder, Serialize)] -#[serde(bound(serialize = "ResponseFormat: Serialize"))] -struct ChatCompletionsRequest { - model: String, +#[derive(Debug, Serialize, Deserialize)] +pub struct Choice { + pub finish_reason: String, // I don't care - messages: Vec, - - response_format: ResponseFormat, + pub message: Message, } -#[derive(Debug, Deserialize)] -struct Choice { - finish_reason: String, // I don't care - - message: Message, -} - -#[derive(Debug, Deserialize)] -struct ChatCompletionsResponse { - choices: Vec, +#[derive(Debug, Serialize, Deserialize)] +pub struct ChatCompletionsResponse { + pub choices: Vec, } diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..2b1b35b --- /dev/null +++ b/src/server.rs @@ -0,0 +1,46 @@ +use std::sync::Arc; + +use crate::AppArgs; +use crate::openai; +use axum::Json; +use axum::Router; +use axum::extract::State; +use axum_extra::routing::RouterExt; +use axum_extra::routing::TypedPath; +use serde::Deserialize; +use serde::Serialize; +use typed_builder::TypedBuilder; + +#[derive(Debug, Clone, TypedBuilder)] +pub struct AppState { + app_args: Arc, +} + +#[derive(Deserialize, TypedPath)] +#[typed_path("/v1/chat/completions")] +struct V1ChatCompletions; + +#[derive(Debug, Serialize, Deserialize)] +enum Shortcut { + Xs, + S, + M, + L, + XL +} + +async fn post_v1_chat_completions( + _: V1ChatCompletions, + State(state): State, + Json(request): Json>, +) -> Result, ()> { + todo!(); +} + +pub fn create_router() -> Router { + Router::new().typed_post(post_v1_chat_completions) +} + +pub fn create_app(state: AppState) -> Router<()> { + create_router().with_state(state) +}