Skip to main content

behest_core/
embedding.rs

1//! Provider-neutral embedding request and response types.
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::id::{ModelName, ProviderId};
7use crate::message::TokenUsage;
8
9/// Request for vector embeddings.
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub struct EmbeddingRequest {
12    /// Provider-specific embedding model name.
13    pub model: ModelName,
14    /// Inputs to embed.
15    pub input: Vec<EmbeddingInput>,
16    /// Optional target embedding dimension.
17    pub dimensions: Option<u32>,
18    /// Application metadata forwarded to provider adapters.
19    pub metadata: Value,
20}
21
22impl EmbeddingRequest {
23    /// Creates an embedding request from explicit inputs.
24    #[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    /// Creates an embedding request for one text input.
35    #[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    /// Sets the target embedding dimension.
41    #[must_use]
42    pub fn with_dimensions(mut self, dimensions: u32) -> Self {
43        self.dimensions = Some(dimensions);
44        self
45    }
46
47    /// Sets adapter metadata.
48    #[must_use]
49    pub fn with_metadata(mut self, metadata: Value) -> Self {
50        self.metadata = metadata;
51        self
52    }
53}
54
55/// A single embedding input.
56#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
57#[serde(rename_all = "snake_case", tag = "type")]
58#[non_exhaustive]
59pub enum EmbeddingInput {
60    /// Plain text input.
61    Text {
62        /// Text to embed.
63        text: String,
64    },
65    /// Tokenized input.
66    Tokens {
67        /// Provider-specific token IDs.
68        tokens: Vec<u32>,
69    },
70}
71
72impl EmbeddingInput {
73    /// Creates a text embedding input.
74    #[must_use]
75    pub fn text(text: impl Into<String>) -> Self {
76        Self::Text { text: text.into() }
77    }
78
79    /// Creates a token embedding input.
80    #[must_use]
81    pub fn tokens(tokens: Vec<u32>) -> Self {
82        Self::Tokens { tokens }
83    }
84}
85
86/// Response returned by an embedding provider.
87#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
88pub struct EmbeddingResponse {
89    /// Provider that produced the embeddings.
90    pub provider: ProviderId,
91    /// Model that produced the embeddings.
92    pub model: ModelName,
93    /// Embeddings in the same order as request inputs.
94    pub embeddings: Vec<Embedding>,
95    /// Token accounting, when supplied by the provider.
96    pub usage: Option<TokenUsage>,
97    /// Raw provider response for adapters that retain it.
98    pub raw: Option<Value>,
99}
100
101/// One embedding vector and its input index.
102#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct Embedding {
104    /// Input index associated with this vector.
105    pub index: usize,
106    /// Dense embedding vector.
107    pub vector: Vec<f32>,
108}
109
110impl Embedding {
111    /// Creates one embedding vector.
112    #[must_use]
113    pub fn new(index: usize, vector: Vec<f32>) -> Self {
114        Self { index, vector }
115    }
116}