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    /// Validates the model configuration.
117    fn validate_model(&self) -> Result<()> {
118        Ok(())
119    }
120
121    /// Returns the custom guidance string for system prompt injection, if set.
122    fn custom_guidance(&self) -> Option<&str> {
123        None
124    }
125
126    /// Sends a chat completion request to the provider's API (HTTP-only, no retry).
127    #[allow(private_interfaces)]
128    async fn send_request_inner(
129        &self,
130        request: &ChatCompletionRequest,
131    ) -> Result<ChatCompletionResponse> {
132        self::http::send_request_inner(self, request).await
133    }
134
135    /// Sends a chat completion request and parses the response with retry logic.
136    #[allow(private_interfaces)]
137    async fn send_and_parse<T: serde::de::DeserializeOwned + Send>(
138        &self,
139        request: &ChatCompletionRequest,
140    ) -> Result<(T, AiStats, Vec<String>)> {
141        self::http::send_and_parse(self, request).await
142    }
143
144    /// Analyzes a GitHub issue using the provider's API.
145    async fn analyze_issue(&self, issue: &IssueDetails) -> Result<crate::ai::AiResponse> {
146        self::triage::analyze_issue(self, issue).await
147    }
148
149    /// Builds the system prompt for issue triage.
150    #[must_use]
151    fn build_system_prompt(custom_guidance: Option<&str>) -> String {
152        self::triage::build_system_prompt(custom_guidance)
153    }
154
155    /// Builds the system prompt for issue creation/formatting.
156    #[must_use]
157    fn build_create_system_prompt(custom_guidance: Option<&str>) -> String {
158        self::create::build_create_system_prompt_fn(custom_guidance)
159    }
160
161    /// Creates a formatted GitHub issue using the provider's API.
162    async fn create_issue(
163        &self,
164        title: &str,
165        body: &str,
166        repo: &str,
167    ) -> Result<(CreateIssueResponse, AiStats)> {
168        self::create::create_issue(self, title, body, repo).await
169    }
170
171    /// Estimates the initial size of a PR review prompt in characters.
172    #[must_use]
173    fn estimate_pr_size(
174        pr: &crate::ai::types::PrDetails,
175        ast_context: &str,
176        call_graph: &str,
177    ) -> usize {
178        self::review::estimate_pr_size(pr, ast_context, call_graph)
179    }
180
181    /// Reviews a pull request using the provider's API.
182    #[allow(unused_assignments)]
183    async fn review_pr(
184        &self,
185        ctx: crate::ai::review_context::ReviewContext,
186        review_config: &crate::config::ReviewConfig,
187    ) -> Result<(PrReviewResponse, AiStats, Vec<String>)> {
188        self::review::review_pr(self, ctx, review_config).await
189    }
190
191    /// Suggests labels for a pull request using the provider's API.
192    async fn suggest_pr_labels(
193        &self,
194        title: &str,
195        body: &str,
196        file_paths: &[String],
197    ) -> Result<(Vec<String>, AiStats)> {
198        self::label::suggest_pr_labels(self, title, body, file_paths).await
199    }
200
201    /// Builds the system prompt for PR review.
202    #[must_use]
203    fn build_pr_review_system_prompt(custom_guidance: Option<&str>) -> String {
204        self::review::build_pr_review_system_prompt_fn(custom_guidance)
205    }
206
207    /// Builds the user prompt for PR review.
208    #[must_use]
209    fn build_pr_review_user_prompt(ctx: &mut crate::ai::review_context::ReviewContext) -> String {
210        self::review::build_pr_review_user_prompt(ctx)
211    }
212
213    /// Builds the system prompt for PR label suggestion.
214    #[must_use]
215    fn build_pr_label_system_prompt(custom_guidance: Option<&str>) -> String {
216        self::label::build_pr_label_system_prompt_fn(custom_guidance)
217    }
218
219    /// Builds the user prompt for PR label suggestion.
220    #[must_use]
221    fn build_pr_label_user_prompt(title: &str, body: &str, file_paths: &[String]) -> String {
222        self::label::build_pr_label_user_prompt(title, body, file_paths)
223    }
224}
225
226#[cfg(test)]
227pub(crate) mod test_utils {
228    use super::*;
229
230    pub(crate) static TEST_PROVIDER_CONFIG: ProviderConfig = ProviderConfig {
231        name: "test",
232        display_name: "Test",
233        api_url: "https://test.example.com",
234        api_key_env: "TEST_API_KEY",
235        model: "test-model",
236        max_tokens: 2048,
237        temperature: 0.3,
238    };
239
240    #[derive(Debug, serde::Deserialize)]
241    pub(crate) struct ErrorTestResponse {
242        pub(crate) _message: String,
243    }
244
245    pub(crate) struct TestProvider;
246
247    impl AiProvider for TestProvider {
248        fn config(&self) -> &ProviderConfig {
249            &TEST_PROVIDER_CONFIG
250        }
251
252        fn http_client(&self) -> &Client {
253            unimplemented!()
254        }
255
256        fn api_key(&self) -> &SecretString {
257            unimplemented!()
258        }
259    }
260}