Merge pull request #1 from opensass/api-key

feat: initial impl && test api key for the plebs

Author: wiseaidev

Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "x-ai"
+version = "0.0.1"
+edition = "2021"
+
+[dependencies]
+reqwest = { version = "0.12.9", features = ["json", "blocking"] }
+serde = { version = "1.0.215", features = ["derive"] }
+serde_json = "1.0.133"
+thiserror = "2.0.3"
+
+[dev-dependencies]
+mockito = "1.6.1"
+tokio = { version = "1.41.1", features = ["full"] }

src/api_key.rs

@@ -0,0 +1,57 @@
+//! Reference: https://docs.x.ai/api/endpoints#api-key
+
+use crate::traits::ApiKeyFetcher;
+use crate::{error::XaiError, traits::ClientConfig};
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct ApiKeyInfo {
+    pub acls: Vec<String>,
+    pub api_key_blocked: bool,
+    pub api_key_disabled: bool,
+    pub api_key_id: String,
+    pub create_time: String,
+    pub modified_by: String,
+    pub modify_time: String,
+    pub name: String,
+    pub redacted_api_key: String,
+    pub team_blocked: bool,
+    pub team_id: String,
+    pub user_id: String,
+}
+
+#[derive(Debug, Clone)]
+pub struct ApiKeyRequestBuilder<T: ClientConfig + Clone + Send + Sync> {
+    client: T,
+}
+
+impl<T> ApiKeyRequestBuilder<T>
+where
+    T: ClientConfig + Clone + Send + Sync,
+{
+    pub fn new(client: T) -> Self {
+        Self { client }
+    }
+}
+
+impl<T> ApiKeyFetcher for ApiKeyRequestBuilder<T>
+where
+    T: ClientConfig + Clone + Send + Sync,
+{
+    async fn fetch_api_key_info(&self) -> Result<ApiKeyInfo, XaiError> {
+        let response = self
+            .client
+            .request(reqwest::Method::GET, "api-key")?
+            .send()
+            .await?;
+
+        if response.status().is_success() {
+            let api_key_info = response.json::<ApiKeyInfo>().await?;
+            Ok(api_key_info)
+        } else {
+            Err(XaiError::Http(
+                response.error_for_status().unwrap_err().to_string(),
+            ))
+        }
+    }
+}

src/client.rs

@@ -0,0 +1,60 @@
+use crate::error::XaiError;
+use crate::traits::ClientConfig;
+use crate::XAI_V1_URL;
+use reqwest::{Client as HttpClient, Method, RequestBuilder};
+use std::sync::{Arc, RwLock};
+
+#[derive(Clone, Debug)]
+pub struct XaiClient {
+    http_client: Arc<HttpClient>,
+    api_key: Arc<RwLock<Option<String>>>,
+    base_url: String,
+}
+
+impl XaiClient {
+    pub fn builder() -> XaiClientBuilder {
+        XaiClientBuilder::default()
+    }
+}
+
+impl ClientConfig for XaiClient {
+    fn set_api_key(&self, api_key: String) {
+        let mut key = self.api_key.write().unwrap();
+        *key = Some(api_key);
+    }
+
+    fn get_api_key(&self) -> Option<String> {
+        self.api_key.read().unwrap().clone()
+    }
+
+    fn request(&self, method: Method, endpoint: &str) -> Result<RequestBuilder, XaiError> {
+        let api_key = self.get_api_key().ok_or(XaiError::MissingApiKey)?;
+
+        let url = format!("{}/{}", self.base_url, endpoint);
+        let builder = self
+            .http_client
+            .request(method, &url)
+            .header("Authorization", format!("Bearer {}", api_key));
+        Ok(builder)
+    }
+}
+
+#[derive(Default, Debug)]
+pub struct XaiClientBuilder {
+    base_url: Option<String>,
+}
+
+impl XaiClientBuilder {
+    pub fn base_url(mut self, url: &str) -> Self {
+        self.base_url = Some(url.to_string());
+        self
+    }
+
+    pub fn build(self) -> Result<XaiClient, XaiError> {
+        Ok(XaiClient {
+            http_client: Arc::new(HttpClient::new()),
+            api_key: Arc::new(RwLock::new(None)),
+            base_url: self.base_url.unwrap_or_else(|| XAI_V1_URL.to_string()),
+        })
+    }
+}

src/error.rs

@@ -0,0 +1,22 @@
+use thiserror::Error;
+
+#[derive(Error, Debug, Clone)]
+pub enum XaiError {
+    #[error("HTTP error: {0}")]
+    Http(String),
+
+    #[error("Missing API key. Please set an API key before making requests.")]
+    MissingApiKey,
+
+    #[error("Unexpected response format.")]
+    UnexpectedResponseFormat,
+
+    #[error("Other error: {0}")]
+    Other(String),
+}
+
+impl From<reqwest::Error> for XaiError {
+    fn from(err: reqwest::Error) -> Self {
+        XaiError::Http(err.to_string())
+    }
+}

src/lib.rs

@@ -0,0 +1,6 @@
+pub mod api_key;
+pub mod client;
+pub mod error;
+pub mod traits;
+
+pub const XAI_V1_URL: &str = "https://api.x.ai/v1";

src/traits.rs

@@ -0,0 +1,15 @@
+#![allow(async_fn_in_trait)]
+
+use crate::api_key::ApiKeyInfo;
+use crate::error::XaiError;
+use reqwest::{Method, RequestBuilder};
+
+pub trait ClientConfig {
+    fn set_api_key(&self, api_key: String);
+    fn get_api_key(&self) -> Option<String>;
+    fn request(&self, method: Method, endpoint: &str) -> Result<RequestBuilder, XaiError>;
+}
+
+pub trait ApiKeyFetcher {
+    async fn fetch_api_key_info(&self) -> Result<ApiKeyInfo, XaiError>;
+}

tests/api_key.rs

@@ -0,0 +1,60 @@
+use mockito::{Matcher, Server};
+use x_ai::api_key::ApiKeyRequestBuilder;
+use x_ai::client::XaiClient;
+use x_ai::traits::ApiKeyFetcher;
+use x_ai::traits::ClientConfig;
+
+#[tokio::test]
+async fn test_fetch_api_key_info() {
+    let mut server = Server::new_async().await;
+
+    let mock_response = r#"
+    {
+        "acls": ["api-key:model:*", "api-key:endpoint:*"],
+        "api_key_blocked": false,
+        "api_key_disabled": false,
+        "api_key_id": "ae1e1841-4326-a8a9-8a1a7237db11",
+        "create_time": "2024-01-01T12:55:18.139305Z",
+        "modified_by": "3d38b4dc-4eb7-4785-ae26-c3fa8997ffc7",
+        "modify_time": "2024-08-28T17:20:12.343321Z",
+        "name": "My API Key",
+        "redacted_api_key": "xG1k...b14o",
+        "team_blocked": false,
+        "team_id": "5ea6f6bd-7815-4b8a-9135-28b2d7ba6722",
+        "user_id": "59fbe5f2-040b-46d5-8325-868bb8f23eb2"
+    }
+    "#;
+
+    let _mock = server
+        .mock("GET", "/api-key")
+        .match_header("Authorization", Matcher::Regex(r"^Bearer .*".to_string()))
+        .with_status(200)
+        .with_header("Content-Type", "application/json")
+        .with_body(mock_response)
+        .create_async()
+        .await;
+
+    let client = XaiClient::builder()
+        .base_url(&server.url())
+        .build()
+        .expect("Failed to build XaiClient");
+
+    client.set_api_key("test-api-key".to_string());
+
+    let request_builder = ApiKeyRequestBuilder::new(client);
+
+    let result = request_builder.fetch_api_key_info().await;
+
+    assert!(result.is_ok());
+
+    let api_key_info = result.unwrap();
+    assert_eq!(api_key_info.api_key_id, "ae1e1841-4326-a8a9-8a1a7237db11");
+    assert_eq!(api_key_info.name, "My API Key");
+    assert!(!api_key_info.api_key_blocked);
+    assert_eq!(
+        api_key_info.acls,
+        vec!["api-key:model:*", "api-key:endpoint:*"]
+    );
+
+    server.reset();
+}