Skip to main content

infernum_core/
response.rs

1//! Response types for inference operations.
2
3use serde::{Deserialize, Serialize};
4
5use crate::types::{FinishReason, ModelId, RequestId, Usage};
6
7/// Response from text generation.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct GenerateResponse {
10    /// Request identifier.
11    pub request_id: RequestId,
12
13    /// Model used for generation.
14    pub model: ModelId,
15
16    /// Generated completions.
17    pub choices: Vec<Choice>,
18
19    /// Token usage statistics.
20    pub usage: Usage,
21
22    /// Time to first token in milliseconds.
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub time_to_first_token_ms: Option<f64>,
25
26    /// Total generation time in milliseconds.
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub total_time_ms: Option<f64>,
29}
30
31/// A single completion choice.
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct Choice {
34    /// Index of this choice.
35    pub index: u32,
36
37    /// Generated text.
38    pub text: String,
39
40    /// Reason generation stopped.
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub finish_reason: Option<FinishReason>,
43
44    /// Log probabilities (if requested).
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub logprobs: Option<LogProbs>,
47}
48
49/// Log probability information.
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct LogProbs {
52    /// Token strings.
53    pub tokens: Vec<String>,
54
55    /// Log probabilities for each token.
56    pub token_logprobs: Vec<f32>,
57
58    /// Top log probabilities at each position.
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub top_logprobs: Option<Vec<Vec<TopLogProb>>>,
61}
62
63/// Top log probability entry.
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct TopLogProb {
66    /// Token string.
67    pub token: String,
68
69    /// Log probability.
70    pub logprob: f32,
71}
72
73/// Information about a single generated token.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct TokenInfo {
76    /// Token ID.
77    pub id: u32,
78
79    /// Token text.
80    pub text: String,
81
82    /// Log probability.
83    #[serde(skip_serializing_if = "Option::is_none")]
84    pub logprob: Option<f32>,
85
86    /// Whether this is a special token.
87    #[serde(default)]
88    pub special: bool,
89}
90
91/// Response from embedding generation.
92#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct EmbedResponse {
94    /// Request identifier.
95    pub request_id: RequestId,
96
97    /// Model used for embeddings.
98    pub model: ModelId,
99
100    /// Generated embeddings.
101    pub data: Vec<Embedding>,
102
103    /// Token usage statistics.
104    pub usage: Usage,
105}
106
107/// A single embedding.
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct Embedding {
110    /// Index of this embedding.
111    pub index: u32,
112
113    /// Embedding vector.
114    pub embedding: EmbeddingData,
115}
116
117/// Embedding data in different formats.
118#[derive(Debug, Clone, Serialize, Deserialize)]
119#[serde(untagged)]
120pub enum EmbeddingData {
121    /// Float vector.
122    Float(Vec<f32>),
123
124    /// Base64 encoded binary.
125    Base64(String),
126}
127
128impl EmbeddingData {
129    /// Returns the embedding as a float vector.
130    ///
131    /// # Errors
132    ///
133    /// Returns an error if the data is base64 encoded and decoding fails.
134    pub fn as_floats(&self) -> Result<Vec<f32>, &'static str> {
135        match self {
136            Self::Float(v) => Ok(v.clone()),
137            Self::Base64(_) => Err("base64 decoding not implemented"),
138        }
139    }
140
141    /// Returns the dimensionality of the embedding.
142    #[must_use]
143    pub fn dimensions(&self) -> usize {
144        match self {
145            Self::Float(v) => v.len(),
146            Self::Base64(s) => {
147                // Each f32 is 4 bytes, base64 encoding is ~4/3 ratio
148                (s.len() * 3) / 16
149            },
150        }
151    }
152}