Skip to main content

elizaos_plugin_gateway/
lib.rs

1#![allow(missing_docs)]
2
3pub mod client;
4pub mod config;
5pub mod error;
6pub mod types;
7
8// Re-export commonly used types for convenience
9pub use client::GatewayClient;
10pub use error::{GatewayError, Result};
11pub use types::{
12    EmbeddingParams, GatewayConfig, ImageDescriptionParams, ImageDescriptionResult,
13    ImageGenerationParams, ImageGenerationResult, TextGenerationParams,
14};
15
16use anyhow::Result as AnyhowResult;
17
18pub struct GatewayPlugin {
19    client: GatewayClient,
20}
21
22impl GatewayPlugin {
23    pub fn new(config: GatewayConfig) -> Result<Self> {
24        let client = GatewayClient::new(config)?;
25        Ok(Self { client })
26    }
27
28    pub async fn generate_text(&self, prompt: &str) -> Result<String> {
29        let params = TextGenerationParams::new(prompt);
30        self.client.generate_text(&params).await
31    }
32
33    pub async fn generate_text_with_system(&self, prompt: &str, system: &str) -> Result<String> {
34        let params = TextGenerationParams::new(prompt).system(system);
35        self.client.generate_text(&params).await
36    }
37
38    pub async fn generate_text_with_params(&self, params: &TextGenerationParams) -> Result<String> {
39        self.client.generate_text(params).await
40    }
41
42    pub async fn stream_text(
43        &self,
44        params: &TextGenerationParams,
45    ) -> Result<impl futures::Stream<Item = Result<String>>> {
46        self.client.stream_text(params).await
47    }
48
49    pub async fn stream_text_simple(
50        &self,
51        prompt: &str,
52    ) -> Result<impl futures::Stream<Item = Result<String>>> {
53        let params = TextGenerationParams::new(prompt);
54        self.client.stream_text(&params).await
55    }
56
57    pub async fn create_embedding(&self, text: &str) -> Result<Vec<f32>> {
58        let params = EmbeddingParams::new(text);
59        self.client.create_embedding(&params).await
60    }
61
62    pub async fn generate_image(
63        &self,
64        params: &ImageGenerationParams,
65    ) -> Result<Vec<ImageGenerationResult>> {
66        self.client.generate_image(params).await
67    }
68
69    pub async fn describe_image(
70        &self,
71        params: &ImageDescriptionParams,
72    ) -> Result<ImageDescriptionResult> {
73        self.client.describe_image(params).await
74    }
75
76    /// Generate a structured JSON object.
77    pub async fn generate_object(&self, prompt: &str) -> Result<serde_json::Value> {
78        self.client.generate_object(prompt, None).await
79    }
80
81    pub fn client(&self) -> &GatewayClient {
82        &self.client
83    }
84}
85
86pub fn get_gateway_plugin() -> AnyhowResult<GatewayPlugin> {
87    let config = GatewayConfig::from_env()
88        .map_err(|e| anyhow::anyhow!("Failed to load Gateway config: {}", e))?;
89
90    GatewayPlugin::new(config)
91        .map_err(|e| anyhow::anyhow!("Failed to create Gateway plugin: {}", e))
92}
93
94pub const PLUGIN_NAME: &str = "gateway";
95pub const PLUGIN_DESCRIPTION: &str =
96    "Vercel AI Gateway plugin with text, embedding, and image generation support";
97pub const PLUGIN_VERSION: &str = env!("CARGO_PKG_VERSION");