Skip to main content

roder_ext_openai_embeddings/
lib.rs

1use std::sync::Arc;
2
3use roder_api::capabilities::CapabilityRequest;
4use roder_api::extension::{
5    ExtensionManifest, ExtensionRegistryBuilder, ProvidedService, RoderExtension,
6};
7use semver::Version;
8
9pub mod provider;
10
11pub use provider::OpenAiEmbeddingProvider;
12
13pub struct OpenAiEmbeddingsExtension {
14    api_key: Option<String>,
15}
16
17impl OpenAiEmbeddingsExtension {
18    pub fn from_env() -> Self {
19        Self {
20            api_key: std::env::var("OPENAI_API_KEY").ok(),
21        }
22    }
23
24    pub fn with_api_key(api_key: impl Into<String>) -> Self {
25        Self {
26            api_key: Some(api_key.into()),
27        }
28    }
29}
30
31impl RoderExtension for OpenAiEmbeddingsExtension {
32    fn manifest(&self) -> ExtensionManifest {
33        ExtensionManifest {
34            id: "roder-ext-openai-embeddings".to_string(),
35            name: "OpenAI Embeddings".to_string(),
36            version: Version::new(0, 1, 0),
37            api_version: "0.1.0".to_string(),
38            description: Some("OpenAI text-embedding-3-large provider".to_string()),
39            provides: vec![ProvidedService::EmbeddingProvider("openai".to_string())],
40            required_capabilities: vec![
41                CapabilityRequest::new("network.openai"),
42                CapabilityRequest::new("secrets.openai-api-key"),
43            ],
44        }
45    }
46
47    fn install(&self, registry: &mut ExtensionRegistryBuilder) -> anyhow::Result<()> {
48        registry.embedding_provider(Arc::new(OpenAiEmbeddingProvider::new(self.api_key.clone())));
49        Ok(())
50    }
51}
52
53pub fn extension() -> OpenAiEmbeddingsExtension {
54    OpenAiEmbeddingsExtension::from_env()
55}