copilot-client 0.1.0

//! # GitHub Copilot Client Library
//!
//! This library provides a client for interacting with the GitHub Copilot API. It handles token
//! management, model retrieval, chat completions, and embeddings requests. The client uses
//! [reqwest](https://crates.io/crates/reqwest) for HTTP requests and [serde](https://crates.io/crates/serde)
//! for JSON serialization/deserialization.
//!
//! ## Features
//!
//! - Retrieve a GitHub token from the environment or configuration files.
//! - Fetch available Copilot models and agents.
//! - Send chat completion requests and receive responses.
//! - Request embeddings for provided input strings.

use reqwest::{
    header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION, USER_AGENT},
    Client as HttpClient,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{env, error::Error, fmt, fs, path::Path};

/// Represents errors that can occur when interacting with the GitHub Copilot API.
#[derive(Debug)]
pub enum CopilotError {
    /// An invalid model was specified.
    InvalidModel(String),
    /// An error occurred while retrieving or parsing the GitHub token.
    TokenError(String),
    /// An HTTP error occurred during the API call.
    HttpError(String),
    /// Other errors.
    Other(String),
}

impl fmt::Display for CopilotError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CopilotError::InvalidModel(model) => {
                write!(f, "Invalid model specified: {model}")
            }
            CopilotError::TokenError(msg) => write!(f, "Token error: {msg}"),
            CopilotError::HttpError(msg) => write!(f, "HTTP error: {msg}"),
            CopilotError::Other(msg) => write!(f, "{msg}"),
        }
    }
}

impl Error for CopilotError {}

/// Response from the GitHub Copilot token endpoint.
///
/// The `expires_at` field is a Unix timestamp.
#[derive(Debug, Serialize, Deserialize)]
pub struct CopilotTokenResponse {
    /// The Copilot token.
    pub token: String,
    /// Expiration time as a Unix timestamp.
    pub expires_at: u64,
}

/// Represents an agent returned by the GitHub Copilot API.
#[derive(Debug, Serialize, Deserialize)]
pub struct Agent {
    /// The agent's identifier.
    pub id: String,
    /// The agent's name.
    pub name: String,
    /// An optional description for the agent.
    pub description: Option<String>,
}

/// Response payload for retrieving agents.
#[derive(Debug, Serialize, Deserialize)]
pub struct AgentsResponse {
    /// List of agents.
    pub agents: Vec<Agent>,
}

/// Represents a model available for GitHub Copilot.
#[derive(Debug, Serialize, Deserialize)]
pub struct Model {
    /// The model identifier.
    pub id: String,
    /// The model name.
    pub name: String,
    /// The version of the model, if available.
    pub version: Option<String>,
    /// The tokenizer used by the model, if available.
    pub tokenizer: Option<String>,
    /// Maximum number of input tokens allowed.
    pub max_input_tokens: Option<u32>,
    /// Maximum number of output tokens allowed.
    pub max_output_tokens: Option<u32>,
}

/// Response payload for retrieving models.
#[derive(Debug, Serialize, Deserialize)]
pub struct ModelsResponse {
    /// List of models.
    pub data: Vec<Model>,
}

/// Represents a chat message.
///
/// The `role` field typically contains values such as `"system"`, `"user"`, or `"assistant"`.
#[derive(Debug, Serialize, Deserialize)]
pub struct Message {
    /// The role of the message sender.
    pub role: String,
    /// The content of the message.
    pub content: String,
}

/// Request payload for a chat completion.
#[derive(Debug, Serialize, Deserialize)]
pub struct ChatRequest {
    /// The model identifier to use.
    pub model: String,
    /// The list of messages to send.
    pub messages: Vec<Message>,
    /// Number of chat completions to generate.
    pub n: u32,
    /// Nucleus sampling probability.
    pub top_p: f64,
    /// Whether to stream the response.
    pub stream: bool,
    /// Sampling temperature.
    pub temperature: f64,
    /// Optional maximum number of tokens to generate.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,
}

/// Represents a single choice in a chat completion response.
#[derive(Debug, Serialize, Deserialize)]
pub struct ChatChoice {
    /// The message generated by the model.
    pub message: Message,
    /// The reason why the generation finished.
    pub finish_reason: Option<String>,
    /// Optional token usage information.
    pub usage: Option<TokenUsage>,
}

/// Information about token usage in a chat response.
#[derive(Debug, Serialize, Deserialize)]
pub struct TokenUsage {
    /// Total tokens used.
    pub total_tokens: u32,
}

/// Response payload for a chat completion request.
#[derive(Debug, Serialize, Deserialize)]
pub struct ChatResponse {
    /// List of generated chat choices.
    pub choices: Vec<ChatChoice>,
}

/// Request payload for an embeddings request.
#[derive(Debug, Serialize, Deserialize)]
pub struct EmbeddingRequest {
    /// The dimensions of the embedding vector.
    pub dimensions: u32,
    /// List of input strings to embed.
    pub input: Vec<String>,
    /// The model identifier to use for embeddings.
    pub model: String,
}

/// Represents an individual embedding.
#[derive(Debug, Serialize, Deserialize)]
pub struct Embedding {
    /// The index corresponding to the input.
    pub index: usize,
    /// The embedding vector.
    pub embedding: Vec<f64>,
}

/// Response payload for an embeddings request.
#[derive(Debug, Serialize, Deserialize)]
pub struct EmbeddingResponse {
    /// List of embeddings.
    pub data: Vec<Embedding>,
}

/// Client for interacting with the GitHub Copilot API.
///
/// This client handles GitHub token retrieval, fetching available models,
/// and sending API requests for chat completions and embeddings.
pub struct CopilotClient {
    http_client: HttpClient,
    github_token: String,
    editor_version: String,
    /// List of available models.
    models: Vec<Model>,
}

impl CopilotClient {
    /// Creates a new `CopilotClient` by retrieving the GitHub token from environment variables
    /// or configuration files, and then fetching the list of available models.
    ///
    /// # Arguments
    ///
    /// * `editor_version` - The version of the editor (e.g., "1.0.0").
    ///
    /// # Errors
    ///
    /// Returns a `CopilotError` if the token retrieval or model fetching fails.
    pub async fn from_env_with_models(editor_version: String) -> Result<Self, CopilotError> {
        let github_token =
            get_github_token().map_err(|e| CopilotError::TokenError(e.to_string()))?;
        Self::new_with_models(github_token, editor_version).await
    }

    /// Creates a new `CopilotClient` with the provided GitHub token and editor version,
    /// and fetches the list of available models.
    ///
    /// # Arguments
    ///
    /// * `github_token` - The GitHub token for authentication.
    /// * `editor_version` - The version of the editor.
    ///
    /// # Errors
    ///
    /// Returns a `CopilotError` if the model fetching fails.
    pub async fn new_with_models(
        github_token: String,
        editor_version: String,
    ) -> Result<Self, CopilotError> {
        let http_client = HttpClient::new();
        let mut client = CopilotClient {
            http_client,
            github_token,
            editor_version,
            models: Vec::new(),
        };
        // Fetch and store the available models.
        let models = client.get_models().await?;
        client.models = models;
        Ok(client)
    }

    /// Constructs the HTTP headers required for GitHub Copilot API requests.
    ///
    /// This includes the authentication token, editor version information,
    /// and other necessary headers.
    async fn get_headers(&self) -> Result<HeaderMap, CopilotError> {
        let token = self.get_copilot_token().await?;
        let mut headers = HeaderMap::new();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(&format!("Bearer {token}"))
                .map_err(|e| CopilotError::Other(e.to_string()))?,
        );
        headers.insert(
            "Editor-Version",
            HeaderValue::from_str(&self.editor_version)
                .map_err(|e| CopilotError::Other(e.to_string()))?,
        );
        headers.insert(
            "Editor-Plugin-Version",
            HeaderValue::from_static("CopilotChat.nvim/*"),
        );
        headers.insert(
            "Copilot-Integration-Id",
            HeaderValue::from_static("vscode-chat"),
        );
        headers.insert(USER_AGENT, HeaderValue::from_static("CopilotChat.nvim"));
        headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
        Ok(headers)
    }

    /// Retrieves a GitHub Copilot token using the stored GitHub token.
    ///
    /// # Errors
    ///
    /// Returns a `CopilotError` if the HTTP request fails or the response cannot be parsed.
    async fn get_copilot_token(&self) -> Result<String, CopilotError> {
        let url = "https://api.github.com/copilot_internal/v2/token";
        let mut headers = HeaderMap::new();
        headers.insert(USER_AGENT, HeaderValue::from_static("CopilotChat.nvim"));
        headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
        headers.insert(
            "Authorization",
            HeaderValue::from_str(&format!("Token {}", self.github_token))
                .map_err(|e| CopilotError::Other(e.to_string()))?,
        );
        let res = self
            .http_client
            .get(url)
            .headers(headers)
            .send()
            .await
            .map_err(|e| CopilotError::HttpError(e.to_string()))?
            .error_for_status()
            .map_err(|e| CopilotError::HttpError(e.to_string()))?;
        let token_response: CopilotTokenResponse = res
            .json()
            .await
            .map_err(|e| CopilotError::Other(e.to_string()))?;
        Ok(token_response.token)
    }

    /// Fetches the list of agents from the GitHub Copilot API.
    ///
    /// # Errors
    ///
    /// Returns a `CopilotError` if the HTTP request fails or the response cannot be parsed.
    pub async fn get_agents(&self) -> Result<Vec<Agent>, CopilotError> {
        let url = "https://api.githubcopilot.com/agents";
        let headers = self.get_headers().await?;
        let res = self
            .http_client
            .get(url)
            .headers(headers)
            .send()
            .await
            .map_err(|e| CopilotError::HttpError(e.to_string()))?
            .error_for_status()
            .map_err(|e| CopilotError::HttpError(e.to_string()))?;
        let agents_response: AgentsResponse = res
            .json()
            .await
            .map_err(|e| CopilotError::Other(e.to_string()))?;
        Ok(agents_response.agents)
    }

    /// Fetches the list of available models from the GitHub Copilot API.
    ///
    /// # Errors
    ///
    /// Returns a `CopilotError` if the HTTP request fails or the response cannot be parsed.
    pub async fn get_models(&self) -> Result<Vec<Model>, CopilotError> {
        let url = "https://api.githubcopilot.com/models";
        let headers = self.get_headers().await?;
        let res = self
            .http_client
            .get(url)
            .headers(headers)
            .send()
            .await
            .map_err(|e| CopilotError::HttpError(e.to_string()))?
            .error_for_status()
            .map_err(|e| CopilotError::HttpError(e.to_string()))?;
        let models_response: ModelsResponse = res
            .json()
            .await
            .map_err(|e| CopilotError::Other(e.to_string()))?;
        Ok(models_response.data)
    }

    /// Sends a chat completion request to the GitHub Copilot API.
    ///
    /// # Arguments
    ///
    /// * `messages` - A vector of chat messages to send.
    /// * `model_id` - The identifier of the model to use.
    ///
    /// # Errors
    ///
    /// Returns a `CopilotError::InvalidModel` error if the specified model is not available,
    /// or another `CopilotError` if the HTTP request or response parsing fails.
    pub async fn chat_completion(
        &self,
        messages: Vec<Message>,
        model_id: String,
    ) -> Result<ChatResponse, CopilotError> {
        // Check if the specified model is available.
        if !self.models.iter().any(|m| m.id == model_id) {
            return Err(CopilotError::InvalidModel(model_id));
        }
        let url = "https://api.githubcopilot.com/chat/completions";
        let headers = self.get_headers().await?;
        let request_body = ChatRequest {
            model: model_id,
            messages,
            n: 1,
            top_p: 1.0,
            stream: false,
            temperature: 0.5,
            max_tokens: None,
        };
        let res = self
            .http_client
            .post(url)
            .headers(headers)
            .json(&request_body)
            .send()
            .await
            .map_err(|e| CopilotError::HttpError(e.to_string()))?
            .error_for_status()
            .map_err(|e| CopilotError::HttpError(e.to_string()))?;
        let chat_response: ChatResponse = res
            .json()
            .await
            .map_err(|e| CopilotError::Other(e.to_string()))?;
        Ok(chat_response)
    }

    /// Sends an embeddings request to the GitHub Copilot API.
    ///
    /// # Arguments
    ///
    /// * `inputs` - A vector of input strings to generate embeddings for.
    ///
    /// # Errors
    ///
    /// Returns a `CopilotError` if the HTTP request fails or the response cannot be parsed.
    pub async fn get_embeddings(
        &self,
        inputs: Vec<String>,
    ) -> Result<Vec<Embedding>, CopilotError> {
        let url = "https://api.githubcopilot.com/embeddings";
        let headers = self.get_headers().await?;
        let request_body = EmbeddingRequest {
            dimensions: 512,
            input: inputs,
            model: "text-embedding-3-small".to_string(),
        };
        let res = self
            .http_client
            .post(url)
            .headers(headers)
            .json(&request_body)
            .send()
            .await
            .map_err(|e| CopilotError::HttpError(e.to_string()))?
            .error_for_status()
            .map_err(|e| CopilotError::HttpError(e.to_string()))?;
        let embedding_response: EmbeddingResponse = res
            .json()
            .await
            .map_err(|e| CopilotError::Other(e.to_string()))?;
        Ok(embedding_response.data)
    }
}

/// Retrieves the GitHub token from the `GITHUB_TOKEN` environment variable or from a configuration file.
///
/// # Errors
///
/// Returns an error if the token is not found in the environment or configuration files.
pub fn get_github_token() -> Result<String, Box<dyn Error>> {
    if let Ok(token) = env::var("GITHUB_TOKEN") {
        if env::var("CODESPACES").is_ok() {
            return Ok(token);
        }
    }
    let config_dir = get_config_path()?;
    let file_paths = vec![
        format!("{config_dir}/github-copilot/hosts.json"),
        format!("{config_dir}/github-copilot/apps.json"),
    ];
    for file_path in file_paths {
        if Path::new(&file_path).exists() {
            let content = fs::read_to_string(&file_path)?;
            let json_value: Value = serde_json::from_str(&content)?;
            if let Some(obj) = json_value.as_object() {
                for (key, value) in obj {
                    if key.contains("github.com") {
                        if let Some(oauth_token) = value.get("oauth_token") {
                            if let Some(token_str) = oauth_token.as_str() {
                                return Ok(token_str.to_string());
                            }
                        }
                    }
                }
            }
        }
    }
    Err("Failed to find GitHub token".into())
}

/// Returns the user's configuration directory.
///
/// On Unix systems, this is determined by the `XDG_CONFIG_HOME` environment variable or defaults
/// to `$HOME/.config`. On Windows, it uses `LOCALAPPDATA`.
///
/// # Errors
///
/// Returns an error if the configuration directory cannot be determined.
pub fn get_config_path() -> Result<String, Box<dyn Error>> {
    if let Ok(xdg) = env::var("XDG_CONFIG_HOME") {
        if !xdg.is_empty() {
            return Ok(xdg);
        }
    }
    if cfg!(target_os = "windows") {
        if let Ok(local) = env::var("LOCALAPPDATA") {
            if !local.is_empty() {
                return Ok(local);
            }
        }
    } else if let Ok(home) = env::var("HOME") {
        return Ok(format!("{home}/.config"));
    }
    Err("Failed to find config directory".into())
}