1use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::id::{ModelName, ProviderId};
7use crate::message::TokenUsage;
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub struct EmbeddingRequest {
12 pub model: ModelName,
14 pub input: Vec<EmbeddingInput>,
16 pub dimensions: Option<u32>,
18 pub metadata: Value,
20}
21
22impl EmbeddingRequest {
23 #[must_use]
25 pub fn new(model: ModelName, input: Vec<EmbeddingInput>) -> Self {
26 Self {
27 model,
28 input,
29 dimensions: None,
30 metadata: Value::Null,
31 }
32 }
33
34 #[must_use]
36 pub fn from_text(model: ModelName, text: impl Into<String>) -> Self {
37 Self::new(model, vec![EmbeddingInput::text(text)])
38 }
39
40 #[must_use]
42 pub fn with_dimensions(mut self, dimensions: u32) -> Self {
43 self.dimensions = Some(dimensions);
44 self
45 }
46
47 #[must_use]
49 pub fn with_metadata(mut self, metadata: Value) -> Self {
50 self.metadata = metadata;
51 self
52 }
53}
54
55#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
57#[serde(rename_all = "snake_case", tag = "type")]
58#[non_exhaustive]
59pub enum EmbeddingInput {
60 Text {
62 text: String,
64 },
65 Tokens {
67 tokens: Vec<u32>,
69 },
70}
71
72impl EmbeddingInput {
73 #[must_use]
75 pub fn text(text: impl Into<String>) -> Self {
76 Self::Text { text: text.into() }
77 }
78
79 #[must_use]
81 pub fn tokens(tokens: Vec<u32>) -> Self {
82 Self::Tokens { tokens }
83 }
84}
85
86#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
88pub struct EmbeddingResponse {
89 pub provider: ProviderId,
91 pub model: ModelName,
93 pub embeddings: Vec<Embedding>,
95 pub usage: Option<TokenUsage>,
97 pub raw: Option<Value>,
99}
100
101#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct Embedding {
104 pub index: usize,
106 pub vector: Vec<f32>,
108}
109
110impl Embedding {
111 #[must_use]
113 pub fn new(index: usize, vector: Vec<f32>) -> Self {
114 Self { index, vector }
115 }
116}