chat-api/src/api/types.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use utoipa::ToSchema;
use uuid::Uuid;

// ------ Models ------

#[derive(Debug, Serialize, ToSchema)]
pub struct ApiModelInfo {
    pub name: String,
    pub family: Option<String>,
    pub parameter_size: Option<String>,
    pub metadata: ApiModelMetadata,
}

#[derive(Debug, Clone, Default, Serialize, ToSchema)]
pub struct ApiModelMetadata {
    pub extra: HashMap<String, String>,
}

// ------ Endpoint: /models ------

#[derive(Debug, Serialize, ToSchema)]
pub struct ApiModelsResponse {
    pub models: Vec<ApiModelInfo>,
}

// ------ Endpoint: /models/{model}/load ------

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ApiLoadModelResponse {
    pub model: String,
    pub status: String,
    pub keep_alive: String,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ApiLoadModelRequest {
    pub keep_alive: String,
}

// ------ Endpoint: /models/{model}/unload ------

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ApiUnloadModelResponse {
    pub model: String,
    pub status: String,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ApiUnloadModelRequest {
    pub keep_alive: String,
}

// ------ Completions ------

#[derive(Debug, Deserialize, Serialize, ToSchema, Default)]
pub struct ApiLlmOptions {
    #[serde(default)]
    pub stream: bool,

    pub temperature: Option<f32>,
    pub top_p: Option<f32>,

    pub top_k: Option<u32>,
    pub repeat_penalty: Option<f32>,
    pub seed: Option<i64>,

    pub num_ctx: Option<u32>,
    pub num_predict: Option<u32>,

    pub stop: Option<Vec<String>>,

    pub keep_alive: Option<String>,

    pub context_depth: Option<u32>,
}

// ------ Endpoint: /completions ------

#[derive(Debug, Deserialize, Serialize, ToSchema)]
pub struct ApiCompletionRequest {
    #[serde(flatten)]
    pub options: ApiLlmOptions,

    pub model: String,
    pub prompt: String,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ApiCompletionResponse {
    pub id: Uuid,
    pub object: ApiCompletionObject,
    pub created: String,
    pub model: String,
    pub choices: Vec<Choice>,
    pub usage: Usage,
}

#[derive(Debug, Serialize, Deserialize, ToSchema, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ApiCompletionObject {
    TextCompletion,
}

// ------ Endpoint: /chat/completions ------

#[derive(Debug, Serialize, Deserialize, ToSchema, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ApiFinishReason {
    Stop,
    Length,
    ContentFilter,
    ToolCalls,
    Error,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct Choice {
    pub text: String,
    pub index: u32,
    pub finish_reason: ApiFinishReason,
}

#[derive(Debug, Serialize, Deserialize, ToSchema, Clone, Copy)]
pub struct Usage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
}

// #[derive(Debug, Serialize, Deserialize, ToSchema)]
// pub struct CompletionChunk {
//     pub id: String,
//     pub object: String,
//     pub choices: Vec<Choice>,
// }

#[derive(Debug, Deserialize, Serialize, ToSchema)]
pub struct ApiChatRequest {
    #[serde(flatten)]
    pub base: ApiLlmOptions,

    pub model: String,
    pub message: ApiMessage,

    // Non standard Open AI
    pub conversation_id: Option<Uuid>,
    pub parent_id: Option<Uuid>, // used when branching, regenerate, etc
}

#[derive(Debug, Deserialize, Serialize, ToSchema)]
pub struct ApiMessage {
    pub role: ApiRole,
    pub content: String,
}

#[derive(Debug, Deserialize, Serialize, ToSchema, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ApiRole {
    System,
    User,
    Assistant,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ApiChatResponse {
    pub id: Uuid,
    pub object: ApiCompletionObject,
    pub created: String,
    pub model: String,
    pub choices: Vec<ApiChatChoice>,
    pub usage: Usage,

    pub conversation_id: Uuid,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ApiChatChoice {
    pub index: u32,
    pub message: ApiMessage,
    pub finish_reason: ApiFinishReason,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ChatCompletionChunk {
    pub id: String,
    pub object: String,
    pub choices: Vec<ChatChunkChoice>,
    pub usage: Option<Usage>,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ChatChunkChoice {
    pub index: u32,
    pub delta: Delta,
    pub finish_reason: Option<ApiFinishReason>,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct Delta {
    pub content: Option<String>,
    pub role: Option<ApiRole>,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct StartEventData {
    pub conversation_id: Uuid,
    pub created: u64,
    pub id: Uuid,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct EndEventData {
    pub created: u64,
    pub id: Uuid,
    pub usage: Usage,
}

#[derive(Debug, Serialize, Deserialize, ToSchema)]
#[serde(tag = "type", content = "data")]
pub enum StreamEvent {
    #[serde(rename = "start")]
    Start(StartEventData),
    #[serde(rename = "end")]
    End(EndEventData),
    #[serde(rename = "delta")]
    Delta(ChatCompletionChunk),
}

// ------ Api Key ------

#[derive(Debug, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ApiKeyScope {
    User,
    Admin,
}

#[derive(serde::Deserialize)]
pub struct CreateApiKeyRequest {
    pub name: String,
    pub scopes: Vec<ApiKeyScope>,
}

#[derive(serde::Serialize)]
pub struct CreateApiKeyResponse {
    pub api_key: String,
}

// ------ Fetch database ------
// --- Shared ---
#[derive(Debug, Deserialize)]
pub struct CursorPage {
    pub limit: Option<u32>,
    pub before: Option<chrono::DateTime<chrono::Utc>>,
}

// --- Conversation ---

#[derive(Debug, Serialize)]
pub struct ConversationSummary {
    pub id: Uuid,
    pub title: String,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub updated_at: chrono::DateTime<chrono::Utc>,
}

#[derive(Debug, Serialize)]
pub struct ConversationListResponse {
    pub conversations: Vec<ConversationSummary>,
    pub has_more: bool,
}

// --- Message ---

#[derive(Debug, Serialize)]
pub enum ApiChatRole {
    System,
    User,
    Assistant,
}

#[derive(Debug, Serialize)]
pub struct MessageSummary {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub role: ApiChatRole,
    pub content: String,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub tokens: Option<i32>,
}

#[derive(Debug, Serialize)]
pub struct MessageListResponse {
    pub messages: Vec<MessageSummary>,
    pub has_more: bool,
}