Skip to main content

aptu_core/ai/provider/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! AI provider trait and shared implementations.
4//!
5//! Defines the `AiProvider` trait that all AI providers must implement,
6//! along with default implementations for shared logic like prompt building,
7//! request sending, and response parsing.
8
9pub mod create;
10pub mod http;
11pub mod label;
12pub mod parse;
13pub mod review;
14pub mod triage;
15
16use anyhow::Result;
17use async_trait::async_trait;
18use reqwest::Client;
19use secrecy::SecretString;
20
21use crate::ai::registry::ProviderConfig;
22use crate::ai::types::{
23    ChatCompletionRequest, ChatCompletionResponse, CreateIssueResponse, IssueDetails,
24    PrReviewResponse,
25};
26use crate::history::AiStats;
27
28pub(crate) use crate::ai::provider::parse::{SCHEMA_PREAMBLE, sanitize_prompt_field};
29
30/// Maximum length for issue body to stay within token limits.
31pub const MAX_BODY_LENGTH: usize = 4000;
32
33/// Maximum number of comments to include in the prompt.
34pub const MAX_COMMENTS: usize = 5;
35
36/// Maximum number of files to include in PR review prompt.
37pub const MAX_FILES: usize = 20;
38
39/// Maximum number of labels to include in the prompt.
40pub const MAX_LABELS: usize = 30;
41
42/// Maximum number of milestones to include in the prompt.
43pub const MAX_MILESTONES: usize = 10;
44
45/// AI provider trait for issue triage and creation.
46///
47/// Defines the interface that all AI providers must implement.
48/// Default implementations are provided for shared logic.
49#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
50#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
51pub trait AiProvider: Send + Sync {
52    /// Returns the provider configuration.
53    fn config(&self) -> &ProviderConfig;
54
55    /// Returns the name of the provider (e.g., "gemini", "openrouter").
56    fn name(&self) -> &str {
57        self.config().name
58    }
59
60    /// Returns the API URL for this provider.
61    fn api_url(&self) -> &str {
62        self.config().api_url
63    }
64
65    /// Returns the environment variable name for the API key.
66    fn api_key_env(&self) -> &str {
67        self.config().api_key_env
68    }
69
70    /// Returns the HTTP client for making requests.
71    fn http_client(&self) -> &Client;
72
73    /// Returns the API key for authentication.
74    fn api_key(&self) -> &SecretString;
75
76    /// Returns the model name.
77    fn model(&self) -> &str {
78        self.config().model
79    }
80
81    /// Returns the maximum tokens for API responses.
82    fn max_tokens(&self) -> u32 {
83        self.config().max_tokens
84    }
85
86    /// Returns the temperature for API requests.
87    fn temperature(&self) -> f32 {
88        self.config().temperature
89    }
90
91    /// Returns whether this provider is Anthropic-compatible and supports
92    /// `cache_control` on message blocks.
93    fn is_anthropic(&self) -> bool {
94        self.name() == crate::ai::registry::PROVIDER_ANTHROPIC
95    }
96
97    /// Returns the maximum retry attempts for rate-limited requests.
98    fn max_attempts(&self) -> u32 {
99        3
100    }
101
102    /// Returns the circuit breaker for this provider (optional).
103    fn circuit_breaker(&self) -> Option<&crate::ai::CircuitBreaker> {
104        None
105    }
106
107    /// Builds HTTP headers for API requests.
108    fn build_headers(&self) -> reqwest::header::HeaderMap {
109        let mut headers = reqwest::header::HeaderMap::new();
110        if let Ok(val) = "application/json".parse() {
111            headers.insert("Content-Type", val);
112        }
113        headers
114    }
115
116    /// Provider-specific request body extensions (e.g., `OpenRouter` data controls).
117    fn provider_body_extensions(&self) -> Option<serde_json::Value> {
118        None
119    }
120
121    /// Validates the model configuration.
122    fn validate_model(&self) -> Result<()> {
123        Ok(())
124    }
125
126    /// Returns the custom guidance string for system prompt injection, if set.
127    fn custom_guidance(&self) -> Option<&str> {
128        None
129    }
130
131    /// Sends a chat completion request to the provider's API (HTTP-only, no retry).
132    #[allow(private_interfaces)]
133    async fn send_request_inner(
134        &self,
135        request: &ChatCompletionRequest,
136    ) -> Result<ChatCompletionResponse> {
137        self::http::send_request_inner(self, request).await
138    }
139
140    /// Sends a chat completion request and parses the response with retry logic.
141    #[allow(private_interfaces)]
142    async fn send_and_parse<T: serde::de::DeserializeOwned + Send>(
143        &self,
144        request: &ChatCompletionRequest,
145    ) -> Result<(T, AiStats, Vec<String>)> {
146        self::http::send_and_parse(self, request).await
147    }
148
149    /// Analyzes a GitHub issue using the provider's API.
150    async fn analyze_issue(&self, issue: &IssueDetails) -> Result<crate::ai::AiResponse> {
151        self::triage::analyze_issue(self, issue).await
152    }
153
154    /// Builds the system prompt for issue triage.
155    #[must_use]
156    fn build_system_prompt(custom_guidance: Option<&str>) -> String {
157        self::triage::build_system_prompt(custom_guidance)
158    }
159
160    /// Builds the system prompt for issue creation/formatting.
161    #[must_use]
162    fn build_create_system_prompt(custom_guidance: Option<&str>) -> String {
163        self::create::build_create_system_prompt_fn(custom_guidance)
164    }
165
166    /// Creates a formatted GitHub issue using the provider's API.
167    async fn create_issue(
168        &self,
169        title: &str,
170        body: &str,
171        repo: &str,
172    ) -> Result<(CreateIssueResponse, AiStats)> {
173        self::create::create_issue(self, title, body, repo).await
174    }
175
176    /// Reviews a pull request using the provider's API.
177    #[allow(unused_assignments)]
178    async fn review_pr(
179        &self,
180        ctx: crate::ai::review_context::ReviewContext,
181        review_config: &crate::config::ReviewConfig,
182    ) -> Result<(PrReviewResponse, AiStats, Vec<String>)> {
183        self::review::review_pr(self, ctx, review_config).await
184    }
185
186    /// Suggests labels for a pull request using the provider's API.
187    async fn suggest_pr_labels(
188        &self,
189        title: &str,
190        body: &str,
191        file_paths: &[String],
192    ) -> Result<(Vec<String>, AiStats)> {
193        self::label::suggest_pr_labels(self, title, body, file_paths).await
194    }
195
196    /// Builds the system prompt for PR review.
197    #[must_use]
198    fn build_pr_review_system_prompt(custom_guidance: Option<&str>) -> String {
199        self::review::build_pr_review_system_prompt_fn(custom_guidance)
200    }
201
202    /// Builds the user prompt for PR review.
203    #[must_use]
204    fn build_pr_review_user_prompt(ctx: &mut crate::ai::review_context::ReviewContext) -> String {
205        self::review::build_pr_review_user_prompt(ctx)
206    }
207
208    /// Builds the system prompt for PR label suggestion.
209    #[must_use]
210    fn build_pr_label_system_prompt(custom_guidance: Option<&str>) -> String {
211        self::label::build_pr_label_system_prompt_fn(custom_guidance)
212    }
213
214    /// Builds the user prompt for PR label suggestion.
215    #[must_use]
216    fn build_pr_label_user_prompt(title: &str, body: &str, file_paths: &[String]) -> String {
217        self::label::build_pr_label_user_prompt(title, body, file_paths)
218    }
219}
220
221#[cfg(test)]
222pub(crate) mod test_utils {
223    use super::*;
224
225    pub(crate) static TEST_PROVIDER_CONFIG: ProviderConfig = ProviderConfig {
226        name: "test",
227        display_name: "Test",
228        api_url: "https://test.example.com",
229        api_key_env: "TEST_API_KEY",
230        model: "test-model",
231        max_tokens: 2048,
232        temperature: 0.3,
233    };
234
235    #[derive(Debug, serde::Deserialize)]
236    pub(crate) struct ErrorTestResponse {
237        pub(crate) _message: String,
238    }
239
240    pub(crate) struct TestProvider;
241
242    impl AiProvider for TestProvider {
243        fn config(&self) -> &ProviderConfig {
244            &TEST_PROVIDER_CONFIG
245        }
246
247        fn http_client(&self) -> &Client {
248            unimplemented!()
249        }
250
251        fn api_key(&self) -> &SecretString {
252            unimplemented!()
253        }
254    }
255}