operonx 0.6.2

High-performance Rust execution backend for Operon workflows
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

//! Embedding provider configuration.
//!
//! Mirrors Python [`operonx/providers/embeddings/config.py`](../../../../../operonx/providers/embeddings/config.py).
//! Python uses a flat `EmbeddingConfig` class with an `api_type` enum — we
//! match that shape exactly here.

use serde::{Deserialize, Serialize};

/// Supported embedding backends.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum EmbeddingType {
    #[serde(rename = "openai")]
    OpenAI,
    #[serde(rename = "azure")]
    Azure,
    #[serde(rename = "gemini")]
    Gemini,
    #[serde(rename = "tei")]
    Tei,
    #[serde(rename = "vllm")]
    Vllm,
    #[serde(rename = "hf")]
    HuggingFace,
    #[serde(rename = "onnx")]
    Onnx,
}

impl Default for EmbeddingType {
    fn default() -> Self {
        EmbeddingType::Vllm
    }
}

/// Flat embedding config — one struct covers every backend variant.
/// The factory branches on `api_type`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct EmbeddingConfig {
    #[serde(default)]
    pub api_type: EmbeddingType,
    #[serde(default)]
    pub api_key: Option<String>,
    #[serde(default)]
    pub base_url: Option<String>,
    #[serde(default)]
    pub api_version: Option<String>,
    #[serde(default)]
    pub model: Option<String>,
    #[serde(default)]
    pub embed_batch_size: Option<usize>,
    #[serde(default)]
    pub dimensions: Option<usize>,

    // ONNX-specific fields (mirrors Python `OnnxInferenceConfig` inherited
    // fields when `api_type=onnx`).
    #[serde(default)]
    pub model_path: Option<String>,
    #[serde(default)]
    pub tokenizer_path: Option<String>,
    #[serde(default)]
    pub max_length: Option<usize>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_openai_embedding() {
        let src = r#"{"api_type": "openai", "api_key": "sk", "model": "text-embedding-3-small", "dimensions": 1536}"#;
        let cfg: EmbeddingConfig = serde_json::from_str(src).unwrap();
        assert_eq!(cfg.api_type, EmbeddingType::OpenAI);
        assert_eq!(cfg.dimensions, Some(1536));
    }

    #[test]
    fn parses_onnx_embedding() {
        let src =
            r#"{"api_type": "onnx", "model_path": "/tmp/m.onnx", "tokenizer_path": "/tmp/t"}"#;
        let cfg: EmbeddingConfig = serde_json::from_str(src).unwrap();
        assert_eq!(cfg.api_type, EmbeddingType::Onnx);
        assert_eq!(cfg.model_path.as_deref(), Some("/tmp/m.onnx"));
    }
}