aither_copilot/
provider.rs

1//! Copilot provider for listing and instantiating models.
2
3use crate::{
4    Copilot,
5    auth::get_session_token,
6    constant::{COPILOT_BASE_URL, COPILOT_INTEGRATION_ID, EDITOR_VERSION},
7    error::CopilotError,
8};
9use aither_core::llm::{
10    LanguageModelProvider, model::Profile as ModelProfile, provider::Profile as ProviderProfile,
11};
12use aither_models::lookup as lookup_model_info;
13use serde::Deserialize;
14use std::{future::Future, sync::Arc};
15use zenwave::{Client, client, header};
16
17/// Provider capable of listing and instantiating `Copilot` models.
18#[derive(Clone, Debug)]
19pub struct CopilotProvider {
20    inner: Arc<ProviderConfig>,
21}
22
23impl CopilotProvider {
24    /// Create a new provider with the given session token.
25    pub fn new(token: impl Into<String>) -> Self {
26        Self {
27            inner: Arc::new(ProviderConfig {
28                token: token.into(),
29                base_url: COPILOT_BASE_URL.to_string(),
30                editor_version: EDITOR_VERSION.to_string(),
31                integration_id: COPILOT_INTEGRATION_ID.to_string(),
32                oauth_token: None,
33            }),
34        }
35    }
36
37    /// Set the base URL for API calls.
38    #[must_use]
39    pub fn base_url(mut self, url: impl Into<String>) -> Self {
40        Arc::make_mut(&mut self.inner).base_url = url.into();
41        self
42    }
43
44    /// Provide OAuth token so short-lived session tokens can be refreshed.
45    #[must_use]
46    pub fn oauth_token(mut self, token: impl Into<String>) -> Self {
47        Arc::make_mut(&mut self.inner).oauth_token = Some(token.into());
48        self
49    }
50}
51
52impl LanguageModelProvider for CopilotProvider {
53    type Model = Copilot;
54    type Error = CopilotError;
55
56    fn list_models(&self) -> impl Future<Output = Result<Vec<ModelProfile>, Self::Error>> + Send {
57        let cfg = self.inner.clone();
58        async move {
59            let mut cfg = cfg.as_ref().clone();
60            let response = match fetch_models(&cfg).await {
61                Ok(response) => response,
62                Err(err) if is_unauthorized(&err) => match refresh_provider_config(&cfg).await? {
63                    Some(refreshed) => {
64                        cfg = refreshed;
65                        fetch_models(&cfg).await?
66                    }
67                    None => return Err(err),
68                },
69                Err(err) => return Err(err),
70            };
71
72            Ok(response
73                .data
74                .into_iter()
75                .filter_map(|model| {
76                    let context_length = model
77                        .context_length
78                        .or(model.max_tokens)
79                        .or(model.input_token_limit)
80                        .or_else(|| lookup_model_info(&model.id).and_then(aither_models::ModelEntry::max_input_tokens));
81                    let Some(context_length) = context_length else {
82                        tracing::warn!(
83                            model = %model.id,
84                            "skip model from Copilot list: missing context_length and no aither-models fallback"
85                        );
86                        return None;
87                    };
88                    Some(ModelProfile::new(
89                        model.id.clone(),
90                        "copilot",
91                        model.id,
92                        model.owned_by.unwrap_or_else(|| "github".to_string()),
93                        context_length,
94                    ))
95                })
96                .collect())
97        }
98    }
99
100    fn get_model(
101        &self,
102        name: &str,
103    ) -> impl Future<Output = Result<Self::Model, Self::Error>> + Send {
104        let cfg = self.inner.clone();
105        let name = name.to_string();
106        async move {
107            let mut builder = Copilot::builder(cfg.token.clone())
108                .base_url(cfg.base_url.clone())
109                .model(name);
110            if let Some(oauth_token) = &cfg.oauth_token {
111                builder = builder.oauth_token(oauth_token.clone());
112            }
113            Ok(builder.build())
114        }
115    }
116
117    fn profile() -> ProviderProfile {
118        ProviderProfile::new("copilot", "GitHub Copilot API provider")
119    }
120}
121
122#[derive(Debug, Clone)]
123struct ProviderConfig {
124    token: String,
125    base_url: String,
126    editor_version: String,
127    integration_id: String,
128    oauth_token: Option<String>,
129}
130
131#[derive(Debug, Deserialize)]
132struct ModelListResponse {
133    data: Vec<ModelDescriptor>,
134}
135
136#[derive(Debug, Deserialize)]
137struct ModelDescriptor {
138    id: String,
139    #[serde(default)]
140    owned_by: Option<String>,
141    #[serde(default)]
142    context_length: Option<u32>,
143    #[serde(default)]
144    max_tokens: Option<u32>,
145    #[serde(default, alias = "inputTokenLimit", alias = "input_token_limit")]
146    input_token_limit: Option<u32>,
147}
148
149async fn fetch_models(cfg: &ProviderConfig) -> Result<ModelListResponse, CopilotError> {
150    let endpoint = format!("{}/models", cfg.base_url.trim_end_matches('/'));
151    let mut backend = client();
152    backend
153        .get(endpoint)
154        .map_err(CopilotError::Http)?
155        .header(
156            header::AUTHORIZATION.as_str(),
157            format!("Bearer {}", cfg.token),
158        )
159        .map_err(CopilotError::Http)?
160        .header(header::USER_AGENT.as_str(), "aither-copilot/0.1")
161        .map_err(CopilotError::Http)?
162        .header(header::ACCEPT.as_str(), "application/json")
163        .map_err(CopilotError::Http)?
164        .header("Editor-Version", cfg.editor_version.clone())
165        .map_err(CopilotError::Http)?
166        .header("Copilot-Integration-Id", cfg.integration_id.clone())
167        .map_err(CopilotError::Http)?
168        .json()
169        .await
170        .map_err(CopilotError::Http)
171}
172
173const fn is_unauthorized(err: &CopilotError) -> bool {
174    matches!(
175        err,
176        CopilotError::Http(zenwave::Error::Http { status, .. }) if status.as_u16() == 401
177    )
178}
179
180async fn refresh_provider_config(
181    cfg: &ProviderConfig,
182) -> Result<Option<ProviderConfig>, CopilotError> {
183    let Some(oauth_token) = cfg.oauth_token.as_deref() else {
184        return Ok(None);
185    };
186    let session = get_session_token(oauth_token).await?;
187    let mut refreshed = cfg.clone();
188    refreshed.token = session.token;
189    refreshed.base_url = session.api_endpoint;
190    Ok(Some(refreshed))
191}
aither_copilot/provider.rs

aither_copilot/
provider.rs
