feat: server outline

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<SecretString>,
 
@@ -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");
 }

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<S>(PhantomData<S>);
-
-impl<S> Serialize for StructuredOutputJson<S>
-where
-    S: JsonSchema,
-{
-    fn serialize<Serializer>(
-        &self,
-        serializer: Serializer,
-    ) -> Result<Serializer::Ok, Serializer::Error>
-    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::<S>(),
-        )?;
-        map.end()
-    }
+#[derive(Debug, Deserialize, Serialize)]
+pub struct ResponseFormatJsonSchema {
+    pub strict: bool,
+    pub name: String,
+    pub schema: Schema,
 }
 
-#[derive(Debug)]
-enum ResponseFormat<OutputFormat> {
-    Json { schema: PhantomData<OutputFormat> },
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(tag = "type")]
+pub enum ResponseFormat {
+    JsonSchema {
+        json_schema: ResponseFormatJsonSchema,
+    },
 }
 
-impl<OutputFormat> Serialize for ResponseFormat<OutputFormat>
-where
-    OutputFormat: JsonSchema,
-{
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: serde::Serializer,
-    {
-        let mut map = serializer.serialize_map(Some(2))?;
-        map.serialize_entry("type", "json_schema")?;
-        map.serialize_entry(
-            "json_schema",
-            &StructuredOutputJson::<OutputFormat>(Default::default()),
-        )?;
-        map.end()
-    }
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ChatCompletionsRequest<Model> {
+    pub model: Model,
+
+    pub messages: Vec<Message>,
+
+    pub response_format: ResponseFormat,
 }
 
-#[derive(Debug, TypedBuilder, Serialize)]
-#[serde(bound(serialize = "ResponseFormat<OutputSchema>: Serialize"))]
-struct ChatCompletionsRequest<OutputSchema> {
-    model: String,
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Choice {
+    pub finish_reason: String, // I don't care
 
-    messages: Vec<Message>,
-
-    response_format: ResponseFormat<OutputSchema>,
+    pub message: Message,
 }
 
-#[derive(Debug, Deserialize)]
-struct Choice {
-    finish_reason: String, // I don't care
-
-    message: Message,
-}
-
-#[derive(Debug, Deserialize)]
-struct ChatCompletionsResponse {
-    choices: Vec<Choice>,
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ChatCompletionsResponse {
+    pub choices: Vec<Choice>,
 }

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<AppArgs>,
+}
+
+#[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<AppState>,
+    Json(request): Json<openai::ChatCompletionsRequest<Shortcut>>,
+) -> Result<Json<openai::ChatCompletionsResponse>, ()> {
+    todo!();
+}
+
+pub fn create_router() -> Router<AppState> {
+    Router::new().typed_post(post_v1_chat_completions)
+}
+
+pub fn create_app(state: AppState) -> Router<()> {
+    create_router().with_state(state)
+}