Skip to main content

openai_protocol/
embedding.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use super::common::{GenerationRequest, UsageInfo};
5
6// ============================================================================
7// Embedding API
8// ============================================================================
9
10#[serde_with::skip_serializing_none]
11#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)]
12pub struct EmbeddingRequest {
13    /// ID of the model to use
14    pub model: String,
15
16    /// Input can be a string, array of strings, tokens, or batch inputs
17    pub input: Value,
18
19    /// Optional encoding format (e.g., "float", "base64")
20    pub encoding_format: Option<String>,
21
22    /// Optional user identifier
23    pub user: Option<String>,
24
25    /// Optional number of dimensions for the embedding
26    pub dimensions: Option<u32>,
27
28    /// SGLang extension: request id for tracking
29    pub rid: Option<String>,
30}
31
32impl GenerationRequest for EmbeddingRequest {
33    fn rid(&self) -> Option<&str> {
34        self.rid.as_deref()
35    }
36
37    fn is_stream(&self) -> bool {
38        // Embeddings are non-streaming
39        false
40    }
41
42    fn get_model(&self) -> Option<&str> {
43        Some(&self.model)
44    }
45
46    fn extract_text_for_routing(&self) -> String {
47        // Best effort: extract text content for routing decisions
48        match &self.input {
49            Value::String(s) => s.clone(),
50            Value::Array(arr) => arr
51                .iter()
52                .filter_map(|v| v.as_str())
53                .collect::<Vec<_>>()
54                .join(" "),
55            _ => String::new(),
56        }
57    }
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
61pub struct EmbeddingObject {
62    pub object: String, // "embedding"
63    pub embedding: Vec<f32>,
64    pub index: u32,
65}
66
67#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
68pub struct EmbeddingResponse {
69    pub object: String, // "list"
70    pub data: Vec<EmbeddingObject>,
71    pub model: String,
72    pub usage: UsageInfo,
73}