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/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,
+    })
+}