Skip to main content

oxirs_embed/api/
types.rs

1//! Request and response types for the API
2//!
3//! This module contains all the data structures used for API requests and responses.
4
5use crate::{ModelStats, TrainingStats, Vector};
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use uuid::Uuid;
10
11/// Basic embedding request
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct EmbeddingRequest {
14    /// Entity ID to get embedding for
15    pub entity_id: String,
16    /// Entity ID (alias for compatibility)
17    pub entity: String,
18    /// Optional model ID to use
19    pub model_id: Option<Uuid>,
20    /// Optional model version to use
21    pub model_version: Option<String>,
22    /// Use cached result if available
23    pub use_cache: Option<bool>,
24}
25
26/// Basic embedding response
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct EmbeddingResponse {
29    /// Entity ID
30    pub entity_id: String,
31    /// Entity ID (alias for compatibility)
32    pub entity: String,
33    /// Generated embedding vector
34    pub embedding: Vector,
35    /// Embedding dimensions
36    pub dimensions: usize,
37    /// Model ID used
38    pub model_id: Uuid,
39    /// Model version used
40    pub model_version: String,
41    /// Whether result came from cache
42    pub from_cache: bool,
43    /// Generation time in milliseconds
44    pub generation_time_ms: f64,
45}
46
47/// Batch embedding request
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct BatchEmbeddingRequest {
50    /// List of entity IDs
51    pub entity_ids: Vec<String>,
52    /// List of entities (alias for compatibility)
53    pub entities: Vec<String>,
54    /// Optional model ID to use
55    pub model_id: Option<Uuid>,
56    /// Optional model version to use
57    pub model_version: Option<String>,
58    /// Use cached result if available
59    pub use_cache: Option<bool>,
60    /// Batch processing options
61    pub options: Option<BatchOptions>,
62}
63
64/// Batch processing options
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct BatchOptions {
67    /// Use cached results if available
68    pub use_cache: Option<bool>,
69    /// Parallel processing batch size
70    pub batch_size: Option<usize>,
71}
72
73/// A single entity that could not be embedded in a batch request.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct FailedEmbedding {
76    /// The entity identifier that failed
77    pub entity: String,
78    /// Human-readable reason the embedding could not be produced
79    pub error: String,
80}
81
82/// Batch embedding response
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct BatchEmbeddingResponse {
85    /// Embedding results
86    pub embeddings: Vec<EmbeddingResponse>,
87    /// Entities that could not be embedded (e.g. unknown entities). Populated so
88    /// callers can distinguish partial success from silent data loss instead of
89    /// receiving fewer results than requested with no explanation.
90    pub failed: Vec<FailedEmbedding>,
91    /// Total processing time in milliseconds
92    pub total_time_ms: f64,
93    /// Number of cache hits
94    pub cache_hits: usize,
95    /// Number of cache misses
96    pub cache_misses: usize,
97    /// Model ID used
98    pub model_id: Uuid,
99}
100
101/// Text embedding request
102#[derive(Debug, Clone, Serialize, Deserialize)]
103pub struct TextEmbeddingRequest {
104    /// Text to embed
105    pub text: String,
106    /// Optional text type hint
107    pub text_type: Option<String>,
108    /// Optional model ID to use
109    pub model_id: Option<Uuid>,
110    /// Language hint (ISO 639-1)
111    pub language: Option<String>,
112    /// Use cached result if available
113    pub use_cache: Option<bool>,
114}
115
116/// Text embedding response
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct TextEmbeddingResponse {
119    /// Original text
120    pub text: String,
121    /// Generated embedding vector
122    pub embedding: Vector,
123    /// Detected language (if applicable)
124    pub detected_language: Option<String>,
125    /// Model ID used
126    pub model_id: Uuid,
127    /// Whether result came from cache
128    pub from_cache: bool,
129    /// Generation time in milliseconds
130    pub generation_time_ms: f64,
131}
132
133/// Multi-modal embedding request
134#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct MultiModalRequest {
136    /// Text content
137    pub text: Option<String>,
138    /// Knowledge graph entities
139    pub entities: Option<Vec<String>>,
140    /// Optional model ID to use
141    pub model_id: Option<Uuid>,
142    /// Fusion strategy
143    pub fusion_strategy: Option<String>,
144    /// Use cached result if available
145    pub use_cache: Option<bool>,
146}
147
148/// Multi-modal embedding response
149#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct MultiModalResponse {
151    /// Generated unified embedding
152    pub embedding: Vector,
153    /// Individual component embeddings
154    pub component_embeddings: HashMap<String, Vector>,
155    /// Fusion strategy used
156    pub fusion_strategy: String,
157    /// Model ID used
158    pub model_id: Uuid,
159    /// Whether result came from cache
160    pub from_cache: bool,
161    /// Generation time in milliseconds
162    pub generation_time_ms: f64,
163}
164
165/// Stream embedding request
166#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct StreamEmbeddingRequest {
168    /// Stream of items to embed
169    pub items: Vec<StreamEmbeddingItem>,
170    /// Optional model ID to use
171    pub model_id: Option<Uuid>,
172    /// Streaming options
173    pub options: Option<BatchOptions>,
174}
175
176/// Individual item in streaming request
177#[derive(Debug, Clone, Serialize, Deserialize)]
178pub struct StreamEmbeddingItem {
179    /// Item ID
180    pub id: String,
181    /// Item content (text or entity ID)
182    pub content: String,
183    /// Content type
184    pub content_type: String,
185}
186
187/// Triple scoring request
188#[derive(Debug, Clone, Serialize, Deserialize)]
189pub struct TripleScoreRequest {
190    /// Subject entity
191    pub subject: String,
192    /// Predicate relation
193    pub predicate: String,
194    /// Object entity
195    pub object: String,
196    /// Optional model ID to use
197    pub model_id: Option<Uuid>,
198    /// Optional model version to use
199    pub model_version: Option<String>,
200    /// Use cached result if available
201    pub use_cache: Option<bool>,
202}
203
204/// Triple scoring response
205#[derive(Debug, Clone, Serialize, Deserialize)]
206pub struct TripleScoreResponse {
207    /// Subject entity
208    pub subject: String,
209    /// Predicate relation
210    pub predicate: String,
211    /// Object entity
212    pub object: String,
213    /// Triple (subject, predicate, object)
214    pub triple: (String, String, String),
215    /// Plausibility score
216    pub score: f64,
217    /// Model ID used
218    pub model_id: Uuid,
219    /// Model version used
220    pub model_version: String,
221    /// Whether result came from cache
222    pub from_cache: bool,
223    /// Computation time in milliseconds
224    pub computation_time_ms: f64,
225    /// Scoring time in milliseconds
226    pub scoring_time_ms: f64,
227}
228
229/// Prediction request
230#[derive(Debug, Clone, Serialize, Deserialize)]
231pub struct PredictionRequest {
232    /// Input entities
233    pub entities: Vec<String>,
234    /// Prediction type
235    pub prediction_type: PredictionType,
236    /// Number of predictions to return
237    pub top_k: Option<usize>,
238    /// Optional model ID to use
239    pub model_id: Option<Uuid>,
240    /// Use cached result if available
241    pub use_cache: Option<bool>,
242}
243
244/// Types of predictions
245#[derive(Debug, Clone, Serialize, Deserialize)]
246pub enum PredictionType {
247    /// Predict missing objects
248    Objects { subject: String, predicate: String },
249    /// Predict missing subjects
250    Subjects { predicate: String, object: String },
251    /// Predict missing relations
252    Relations { subject: String, object: String },
253}
254
255/// Prediction response
256#[derive(Debug, Clone, Serialize, Deserialize)]
257pub struct PredictionResponse {
258    /// Input entities
259    pub input: Vec<String>,
260    /// Prediction type
261    pub prediction_type: String,
262    /// Predictions with scores
263    pub predictions: Vec<(String, f64)>,
264    /// Model version used
265    pub model_version: String,
266    /// Whether result came from cache
267    pub from_cache: bool,
268    /// Prediction time in milliseconds
269    pub prediction_time_ms: f64,
270}
271
272/// Model information request
273#[derive(Debug, Clone, Serialize, Deserialize)]
274pub struct ModelInfoRequest {
275    /// Optional specific model ID
276    pub model_id: Option<Uuid>,
277}
278
279/// Model information response
280#[derive(Debug, Clone, Serialize, Deserialize)]
281pub struct ModelInfoResponse {
282    /// Model statistics
283    pub stats: ModelStats,
284    /// Model health status
285    pub health: ModelHealth,
286    /// Available operations
287    pub capabilities: Vec<String>,
288    /// Last training statistics
289    pub last_training: Option<TrainingStats>,
290}
291
292/// Model health information
293#[derive(Debug, Clone, Serialize, Deserialize)]
294pub struct ModelHealth {
295    /// Current health status
296    pub status: HealthStatus,
297    /// Last health check timestamp
298    pub last_check: DateTime<Utc>,
299    /// Performance metrics
300    pub metrics: HealthMetrics,
301}
302
303/// Health status enumeration
304#[derive(Debug, Clone, Serialize, Deserialize)]
305pub enum HealthStatus {
306    /// Model is healthy and operational
307    Healthy,
308    /// Model is operational but with degraded performance
309    Degraded,
310    /// Model is not operational
311    Unhealthy,
312}
313
314/// Performance metrics for health checking
315#[derive(Debug, Clone, Serialize, Deserialize)]
316pub struct HealthMetrics {
317    /// Average response time in milliseconds
318    pub avg_response_time_ms: f64,
319    /// Number of requests in the last hour
320    pub requests_last_hour: u64,
321    /// Error rate percentage
322    pub error_rate_percent: f64,
323    /// Memory usage in MB
324    pub memory_usage_mb: f64,
325}
326
327/// Query parameters for API endpoints
328#[derive(Debug, Clone, Serialize, Deserialize)]
329pub struct QueryParams {
330    /// Optional limit for results
331    pub limit: Option<usize>,
332    /// Optional offset for pagination
333    pub offset: Option<usize>,
334    /// Optional model ID filter
335    pub model_id: Option<Uuid>,
336    /// Optional format specification
337    pub format: Option<String>,
338    /// Include detailed information
339    pub detailed: Option<bool>,
340}