Skip to main content

gproxy_protocol/protocol/openai/
embeddings.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5use super::common::*;
6
7pub type EmbeddingWireModel = OpenAiWireModel<EmbeddingRequest, EmbeddingResponse>;
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10pub struct EmbeddingRequest {
11    pub input: EmbeddingInput,
12    pub model: OpenAiModelId,
13    #[serde(skip_serializing_if = "Option::is_none")]
14    pub dimensions: Option<u32>,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub encoding_format: Option<EmbeddingEncodingFormat>,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub user: Option<String>,
19    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
20    pub extra: Extra,
21}
22
23#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
24#[serde(untagged)]
25pub enum EmbeddingInput {
26    Text(String),
27    TextList(Vec<String>),
28    TokenList(Vec<i64>),
29    TokenLists(Vec<Vec<i64>>),
30}
31
32#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
33pub struct EmbeddingResponse {
34    pub data: Vec<Embedding>,
35    pub model: OpenAiModelId,
36    pub object: ListObjectType,
37    pub usage: EmbeddingUsage,
38    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
39    pub extra: Extra,
40}
41
42#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
43pub struct Embedding {
44    pub embedding: Vec<f64>,
45    pub index: u32,
46    pub object: EmbeddingObjectType,
47    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
48    pub extra: Extra,
49}
50
51#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52pub struct EmbeddingUsage {
53    pub prompt_tokens: u32,
54    pub total_tokens: u32,
55    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
56    pub extra: Extra,
57}