chore: groundwork for embeddings

.gitignore

@@ -1 +1,3 @@
 /target
+
+.history/

src/openai/chat/completion.rs

@@ -1092,7 +1092,7 @@ pub struct ApproximateLocation {
 
 /// Represents a chat completion response returned by model, based on the provided input.
 #[derive(Debug, Clone, Deserialize, Serialize)]
-pub struct ChatCompletion {
+pub struct ChatCompletion<Model> {
     /// A unique identifier for the chat completion.
     pub id: String,
 
@@ -1103,7 +1103,7 @@ pub struct ChatCompletion {
     pub created: i64,
 
     /// The model used for the chat completion.
-    pub model: String,
+    pub model: Model,
 
     /// The object type, which is always `chat.completion`.
     pub object: String,

src/openai/embeddings.rs

@@ -0,0 +1,97 @@
+use serde::{Deserialize, Serialize};
+
+/// Represents the possible input types for embedding requests
+#[derive(Debug, Clone, Deserialize, Serialize)]
+#[serde(untagged)]
+pub enum EmbeddingInput {
+    /// A single string input
+    String(String),
+    /// An array of string inputs
+    StringArray(Vec<String>),
+    /// An array of token IDs (numbers)
+    NumberArray(Vec<i32>),
+    /// An array of arrays of token IDs (numbers)
+    NumberMatrix(Vec<Vec<i32>>),
+}
+
+/// Represents the encoding format for embeddings
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum EncodingFormat {
+    /// Return embeddings as an array of floats
+    Float,
+    /// Return embeddings as a base64-encoded string
+    Base64,
+}
+
+/// Request body for creating embeddings
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct CreateEmbeddingRequest<Model> {
+    /// Input text to embed
+    pub input: EmbeddingInput,
+
+    /// ID of the model to use for generating embeddings
+    pub model: Model,
+
+    /// Optional number of dimensions for the output embeddings
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub dimensions: Option<u32>,
+
+    /// Format to return the embeddings in
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub encoding_format: Option<EncodingFormat>,
+
+    /// Optional unique identifier for tracking end-user
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub user: Option<String>,
+}
+
+/// Represents a single embedding in the response
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct Embedding {
+    /// The embedding vector, returned as an array of floats or a base64 string
+    /// depending on the encoding_format parameter
+    pub embedding: EmbeddingOutput,
+
+    /// The index of this embedding in the list of embeddings
+    pub index: u32,
+
+    /// The object type, which is always "embedding"
+    pub object: String,
+}
+
+/// Represents the possible output formats for embeddings
+#[derive(Debug, Clone, Deserialize, Serialize)]
+#[serde(untagged)]
+pub enum EmbeddingOutput {
+    /// Array of floats when encoding_format is "float"
+    FloatArray(Vec<f32>),
+    /// Base64-encoded string when encoding_format is "base64"
+    Base64String(String),
+}
+
+/// Usage information for an embeddings request
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct Usage {
+    /// Number of tokens used by the prompt
+    pub prompt_tokens: u32,
+
+    /// Total number of tokens used by the request
+    pub total_tokens: u32,
+}
+
+/// Response from the embeddings API
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct CreateEmbeddingResponse {
+    /// List of generated embeddings
+    pub data: Vec<Embedding>,
+
+    /// The name of the model used to generate the embeddings
+    pub model: String,
+
+    /// The object type, which is always "list"
+    pub object: String,
+
+    /// Usage information for the request
+    pub usage: Usage,
+}

src/openai/mod.rs

@@ -1,2 +1,3 @@
 pub mod chat;
+pub mod embeddings;
 pub mod models;

src/server/chat/completions.rs

@@ -1,26 +1,24 @@
-use crate::AppArgs;
 use crate::CloudflareCredentials;
 use crate::openai;
+use crate::server::AppState;
+use crate::server::Provider;
+use crate::server::Upstream;
+use crate::server::take_random_by_swapping;
 use axum::Json;
-use axum::Router;
 use axum::extract::State;
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::response::Response;
 use axum::response::Sse;
-use axum_extra::routing::{RouterExt, TypedPath};
+use axum_extra::routing::TypedPath;
 use ecow::EcoString;
 use futures::StreamExt;
 use headers::Header;
 use headers::authorization::Bearer;
-use itertools::iproduct;
 use mitsein::vec1::Vec1;
-use rand::Rng;
-use rand::RngExt;
 use reqwest::Certificate;
 use reqwest_sse::EventSource as _;
 use secrecy::ExposeSecret;
-use secrecy::SecretString;
 use serde::Deserialize;
 use serde_with::DeserializeFromStr;
 use serde_with::SerializeDisplay;
@@ -28,37 +26,19 @@ use std::collections::VecDeque;
 use std::convert::Infallible;
 use std::fmt::Display;
 use std::str::FromStr;
-use std::sync::Arc;
 use strum::EnumString;
 use strum::VariantArray;
 use tokio::runtime::Handle;
-use typed_builder::TypedBuilder;
-use url::Url;
-
-#[derive(Debug, Clone, TypedBuilder)]
-pub struct AppState {
-    app_args: Arc<AppArgs>,
-}
 
 #[derive(Deserialize, TypedPath)]
 #[typed_path("/v1/chat/completions")]
-struct V1ChatCompletionsPath;
-
-#[derive(Debug, Clone)]
-enum Provider {
-    CloudflareWorkersAi,
-    GoogleAiStudio,
-    Groq,
-    HuggingFace,
-    Mistral,
-    OpenRouter,
-}
+pub struct V1ChatCompletionsPath;
 
 #[derive(
     Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
 )]
 #[strum(serialize_all = "lowercase")]
-enum Shortcut {
+pub enum Shortcut {
     XS,
     S,
     M,
@@ -69,7 +49,10 @@ enum Shortcut {
 }
 
 #[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)]
-struct ShortcutAndMaybeReasoning(Shortcut, Option<openai::chat::completion::ReasoningEffort>);
+pub struct ShortcutAndMaybeReasoning(
+    pub Shortcut,
+    pub Option<openai::chat::completion::ReasoningEffort>,
+);
 
 impl FromStr for ShortcutAndMaybeReasoning {
     type Err = Infallible; // TODO
@@ -195,28 +178,8 @@ impl Shortcut {
     }
 }
 
-fn take_random_by_swapping<T>(mut vec1: Vec1<T>, mut rng: impl Rng) -> (T, Option<Vec1<T>>) {
-    let end = vec1.len().get() - 1;
-    let index = rng.random_range(0..=end);
-
-    vec1.swap(index, end);
-    let (body, element) = vec1.into_rtail_and_head();
-
-    let body = body.try_into().ok();
-
-    (element, body)
-}
-
-#[derive(Debug, TypedBuilder)]
-struct Upstream {
-    base_url: Url,
-
-    #[builder(setter(strip_option))]
-    api_key: Option<SecretString>,
-}
-
-enum ChatCompletionResponse {
-    Immediate(openai::chat::completion::ChatCompletion),
+pub enum ChatCompletionResponse {
+    Immediate(openai::chat::completion::ChatCompletion<String>),
     Stream(reqwest::Response),
 }
 
@@ -258,7 +221,7 @@ impl IntoResponse for ChatCompletionResponse {
     }
 }
 
-enum ChatCompletionError {
+pub enum ChatCompletionError {
     NoneSucceeded,
 }
 
@@ -272,7 +235,7 @@ impl IntoResponse for ChatCompletionError {
 
 #[axum::debug_handler]
 #[tracing::instrument(skip(state, request))]
-async fn create_chat_completion(
+pub async fn create_chat_completion(
     _: V1ChatCompletionsPath,
     State(state): State<AppState>,
     Json(request): Json<
@@ -544,44 +507,3 @@ async fn create_chat_completion(
 
     Err(ChatCompletionError::NoneSucceeded)
 }
-
-#[derive(Deserialize, TypedPath)]
-#[typed_path("/v1/models")]
-struct V1ModelsPath;
-
-#[tracing::instrument]
-async fn list_models(
-    _: V1ModelsPath,
-) -> Json<openai::models::ListModelsResponse<ShortcutAndMaybeReasoning>> {
-    let shortcuts = Shortcut::VARIANTS.iter().cloned();
-    let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS
-        .iter()
-        .cloned()
-        .map(Some)
-        .chain(None);
-
-    let models = iproduct!(shortcuts, reasoning_efforts).map(|(shortcut, reasoning)| {
-        openai::models::Model {
-            id: ShortcutAndMaybeReasoning(shortcut, reasoning),
-            created: 0,
-            object: openai::models::ModelObject::Model,
-            owned_by: "nobody".into(),
-            shutdown_date: None,
-        }
-    });
-
-    Json(openai::models::ListModelsResponse {
-        data: models.collect(),
-        object: openai::models::ListObject::List,
-    })
-}
-
-pub fn create_router() -> Router<AppState> {
-    Router::new()
-        .typed_post(create_chat_completion)
-        .typed_get(list_models)
-}
-
-pub fn create_app(state: AppState) -> Router<()> {
-    create_router().with_state(state)
-}

src/server/chat/mod.rs

@@ -0,0 +1 @@
+pub mod completions;

src/server/embeddings.rs

@@ -0,0 +1,36 @@
+#[derive(Deserialize, TypedPath)]
+#[typed_path("/v1/chat/completions")]
+pub struct V1ChatCompletionsPath;
+
+#[derive(
+    Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
+)]
+#[strum(serialize_all = "lowercase")]
+pub enum Shortcut {
+    XS,
+    S,
+    M,
+    L,
+    XL,
+}
+
+impl Shortcut {
+    fn suggestions(&self) -> Vec<(Provider, EcoString)> {
+        use Provider as P;
+        use Shortcut as S;
+
+        match self {
+            S::XS => vec![],
+
+            S::S => vec![],
+
+            S::M => vec![],
+
+            S::L => vec![],
+
+            S::XL => vec![],
+
+            S::Experimental => vec![],
+        }
+    }
+}

src/server/mod.rs

@@ -0,0 +1,58 @@
+use crate::AppArgs;
+use axum::Router;
+use axum_extra::routing::RouterExt;
+use mitsein::vec1::Vec1;
+use rand::Rng;
+use rand::RngExt;
+use secrecy::SecretString;
+use std::sync::Arc;
+use typed_builder::TypedBuilder;
+use url::Url;
+
+mod chat;
+mod models;
+
+#[derive(Debug, Clone, TypedBuilder)]
+pub struct AppState {
+    app_args: Arc<AppArgs>,
+}
+
+#[derive(Debug, Clone)]
+enum Provider {
+    CloudflareWorkersAi,
+    GoogleAiStudio,
+    Groq,
+    HuggingFace,
+    Mistral,
+    OpenRouter,
+}
+
+fn take_random_by_swapping<T>(mut vec1: Vec1<T>, mut rng: impl Rng) -> (T, Option<Vec1<T>>) {
+    let end = vec1.len().get() - 1;
+    let index = rng.random_range(0..=end);
+
+    vec1.swap(index, end);
+    let (body, element) = vec1.into_rtail_and_head();
+
+    let body = body.try_into().ok();
+
+    (element, body)
+}
+
+#[derive(Debug, TypedBuilder)]
+struct Upstream {
+    base_url: Url,
+
+    #[builder(setter(strip_option))]
+    api_key: Option<SecretString>,
+}
+
+pub fn create_router() -> Router<AppState> {
+    Router::new()
+        .typed_post(chat::completions::create_chat_completion)
+        .typed_get(models::list_models)
+}
+
+pub fn create_app(state: AppState) -> Router<()> {
+    create_router().with_state(state)
+}

src/server/models.rs

@@ -0,0 +1,36 @@
+use crate::openai;
+use axum::Json;
+use axum_extra::routing::TypedPath;
+use itertools::iproduct;
+use serde::Deserialize;
+use strum::VariantArray;
+
+#[derive(Deserialize, TypedPath)]
+#[typed_path("/v1/models")]
+pub struct V1ModelsPath;
+
+#[tracing::instrument]
+pub async fn list_models(_: V1ModelsPath) -> Json<openai::models::ListModelsResponse<String>> {
+    let shortcuts = super::chat::completions::Shortcut::VARIANTS.iter().cloned();
+    let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS
+        .iter()
+        .cloned()
+        .map(Some)
+        .chain(None);
+
+    let models = iproduct!(shortcuts, reasoning_efforts).map(|(shortcut, reasoning)| {
+        openai::models::Model {
+            id: super::chat::completions::ShortcutAndMaybeReasoning(shortcut, reasoning)
+                .to_string(),
+            created: 0,
+            object: openai::models::ModelObject::Model,
+            owned_by: "nobody".into(),
+            shutdown_date: None,
+        }
+    });
+
+    Json(openai::models::ListModelsResponse {
+        data: models.collect(),
+        object: openai::models::ListObject::List,
+    })
+}