apollo/providers/
rs_ai.rs

1//! rs_ai-backed provider adapter for first-class ChatGPT, Gemini, xAI,
2//! Cloudflare and generic OpenAI-compatible endpoints.
3
4use async_trait::async_trait;
5use rs_ai_core::{
6    GenerateOptions, GenerateResult, Message, Prompt, ToolCallRequest, ToolDefinition,
7};
8
9use crate::providers::traits::{
10    ChatRequest, ChatResponse, Provider, ProviderCapabilities, ToolCall, Usage,
11};
12
13/// Provider implementation backed by the `rs_ai` / `rs_ai_core` SDK.
14pub struct RsAiProvider {
15    provider_name: String,
16    model_id: String,
17    api_key: String,
18    base_url: Option<String>,
19    account_id: Option<String>,
20}
21
22impl RsAiProvider {
23    pub fn new(
24        provider_name: &str,
25        model_id: &str,
26        api_key: &str,
27        base_url: Option<String>,
28        account_id: Option<String>,
29    ) -> Self {
30        Self {
31            provider_name: provider_name.to_string(),
32            model_id: model_id.to_string(),
33            api_key: api_key.to_string(),
34            base_url,
35            account_id,
36        }
37    }
38
39    fn effective_model_id(&self) -> &str {
40        if !self.model_id.is_empty() {
41            return &self.model_id;
42        }
43        match self.provider_name.as_str() {
44            "chatgpt" | "openai" => "gpt-4o",
45            "gemini" => "gemini-2.5-flash",
46            "xai" | "grok" => "grok-4.20-reasoning",
47            "cloudflare" => "@cf/meta/llama-3.1-8b-instruct",
48            _ => "",
49        }
50    }
51
52    fn build_model(&self) -> anyhow::Result<Box<dyn rs_ai_core::LanguageModel>> {
53        use rs_ai_providers::{
54            ChatGptProvider, CloudflareProvider, GeminiProvider, OpenAiCompatibleConfig,
55            OpenAiCompatibleProvider, XaiProvider,
56        };
57
58        let model: Box<dyn rs_ai_core::LanguageModel> = match self.provider_name.as_str() {
59            "chatgpt" | "openai" => {
60                Box::new(ChatGptProvider::new(&self.api_key).model(self.effective_model_id()))
61            }
62            "gemini" => {
63                Box::new(GeminiProvider::new(&self.api_key).model(self.effective_model_id()))
64            }
65            "xai" | "grok" => {
66                Box::new(XaiProvider::new(&self.api_key).model(self.effective_model_id()))
67            }
68            "cloudflare" => {
69                let account_id = self
70                    .account_id
71                    .as_deref()
72                    .or(self.base_url.as_deref())
73                    .unwrap_or("")
74                    .to_string();
75                Box::new(
76                    CloudflareProvider::new(account_id, &self.api_key)
77                        .model(self.effective_model_id()),
78                )
79            }
80            other => {
81                let base_url = self
82                    .base_url
83                    .as_deref()
84                    .unwrap_or("https://api.openai.com/v1");
85                let config = OpenAiCompatibleConfig::new(base_url, &self.api_key);
86                let provider = OpenAiCompatibleProvider::new(config, other, other);
87                provider.language_model(self.effective_model_id())
88            }
89        };
90        Ok(model)
91    }
92}
93
94#[async_trait]
95impl Provider for RsAiProvider {
96    fn name(&self) -> &str {
97        &self.provider_name
98    }
99
100    fn capabilities(&self) -> ProviderCapabilities {
101        ProviderCapabilities {
102            native_tools: true,
103            streaming: true,
104            vision: matches!(
105                self.provider_name.as_str(),
106                "chatgpt" | "openai" | "gemini" | "xai" | "grok"
107            ),
108            max_context: 200_000,
109            native_web_search: false,
110        }
111    }
112
113    async fn chat(&self, request: &ChatRequest<'_>) -> anyhow::Result<ChatResponse> {
114        let model = self.build_model()?;
115
116        let messages: Vec<Message> = request
117            .messages
118            .iter()
119            .map(|m| match m.role.as_str() {
120                "system" => Message::system(&m.content),
121                "assistant" => Message::assistant(&m.content),
122                "tool_result" => {
123                    Message::tool_result(m.tool_use_id.as_deref().unwrap_or(""), &m.content)
124                }
125                _ => Message::user(&m.content),
126            })
127            .collect();
128
129        let prompt = Prompt::Messages(messages);
130
131        let mut options = GenerateOptions::default().with_temperature(request.temperature);
132        if let Some(max_tokens) = request.max_tokens {
133            options = options.with_max_tokens(max_tokens);
134        }
135
136        let tools: Vec<ToolDefinition> = request
137            .tools
138            .unwrap_or(&[])
139            .iter()
140            .map(|t| ToolDefinition {
141                name: t.name.clone(),
142                description: t.description.clone(),
143                parameters: t.parameters.clone(),
144                examples: None,
145            })
146            .collect();
147        if !tools.is_empty() {
148            options = options
149                .with_tools(tools)
150                .with_tool_choice(rs_ai_core::ToolChoice::Auto);
151        }
152
153        let result = model
154            .generate(prompt, options)
155            .await
156            .map_err(|e| anyhow::anyhow!("rs_ai provider error: {e}"))?;
157
158        Ok(map_generate_result(result)?)
159    }
160}
161
162fn map_generate_result(result: GenerateResult) -> anyhow::Result<ChatResponse> {
163    let tool_calls = result
164        .tool_calls
165        .iter()
166        .map(|tc: &ToolCallRequest| -> anyhow::Result<ToolCall> {
167            Ok(ToolCall {
168                id: tc.id.clone(),
169                name: tc.name.clone(),
170                arguments: serde_json::to_string(&tc.arguments)?,
171            })
172        })
173        .collect::<Result<Vec<_>, _>>()?;
174
175    let usage = Usage {
176        input_tokens: result.usage.prompt_tokens.unwrap_or(0) as u32,
177        output_tokens: result.usage.completion_tokens.unwrap_or(0) as u32,
178    };
179
180    Ok(ChatResponse {
181        text: result.text,
182        tool_calls,
183        usage: Some(usage),
184    })
185}
apollo/providers/rs_ai.rs

apollo/providers/
rs_ai.rs
