Merge pull request #2 from opensass/completions

feat: impl deez completions

Author: wiseaidev

completions

@@ -0,0 +1,138 @@
+//! Reference: https://docs.x.ai/api/endpoints#chat-completions
+
+use crate::error::XaiError;
+use crate::traits::ChatCompletionsFetcher;
+use crate::traits::ClientConfig;
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatCompletionRequest {
+    pub model: String,
+    pub messages: Vec<Message>,
+    pub temperature: Option<f32>,
+    pub max_tokens: Option<u32>,
+    pub frequency_penalty: Option<f32>,
+    pub presence_penalty: Option<f32>,
+    pub n: Option<u32>,
+    pub stop: Option<Vec<String>>,
+    pub stream: Option<bool>,
+    pub logprobs: Option<bool>,
+    pub top_p: Option<f32>,
+    pub top_logprobs: Option<u32>,
+    pub seed: Option<u32>,
+    pub user: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Message {
+    pub role: String,
+    pub content: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatCompletionResponse {
+    pub id: String,
+    pub object: String,
+    pub created: u64,
+    pub model: String,
+    pub choices: Vec<Choice>,
+    pub usage: Option<Usage>,
+    pub system_fingerprint: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Choice {
+    pub index: u32,
+    pub message: Message,
+    pub finish_reason: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Usage {
+    pub prompt_tokens: u32,
+    pub completion_tokens: u32,
+    pub total_tokens: u32,
+}
+
+#[derive(Debug, Clone)]
+pub struct ChatCompletionsRequestBuilder<T: ClientConfig + Clone + Send + Sync> {
+    client: T,
+    request: ChatCompletionRequest,
+}
+
+impl<T> ChatCompletionsRequestBuilder<T>
+where
+    T: ClientConfig + Clone + Send + Sync,
+{
+    pub fn new(client: T, model: String, messages: Vec<Message>) -> Self {
+        Self {
+            client,
+            request: ChatCompletionRequest {
+                model,
+                messages,
+                temperature: None,
+                max_tokens: None,
+                frequency_penalty: None,
+                presence_penalty: None,
+                n: None,
+                stop: None,
+                stream: None,
+                logprobs: None,
+                top_p: None,
+                top_logprobs: None,
+                seed: None,
+                user: None,
+            },
+        }
+    }
+
+    pub fn temperature(mut self, temperature: f32) -> Self {
+        self.request.temperature = Some(temperature);
+        self
+    }
+
+    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
+        self.request.max_tokens = Some(max_tokens);
+        self
+    }
+
+    pub fn n(mut self, n: u32) -> Self {
+        self.request.n = Some(n);
+        self
+    }
+
+    pub fn stop(mut self, stop: Vec<String>) -> Self {
+        self.request.stop = Some(stop);
+        self
+    }
+
+    pub fn build(self) -> Result<ChatCompletionRequest, XaiError> {
+        Ok(self.request)
+    }
+}
+
+impl<T> ChatCompletionsFetcher for ChatCompletionsRequestBuilder<T>
+where
+    T: ClientConfig + Clone + Send + Sync,
+{
+    async fn create_chat_completion(
+        &self,
+        request: ChatCompletionRequest,
+    ) -> Result<ChatCompletionResponse, XaiError> {
+        let response = self
+            .client
+            .request(reqwest::Method::POST, "chat/completions")?
+            .json(&request)
+            .send()
+            .await?;
+
+        if response.status().is_success() {
+            let chat_completion = response.json::<ChatCompletionResponse>().await?;
+            Ok(chat_completion)
+        } else {
+            Err(XaiError::Http(
+                response.error_for_status().unwrap_err().to_string(),
+            ))
+        }
+    }
+}

src/completions.rs

@@ -1,5 +1,6 @@
 pub mod api_key;
 pub mod client;
+pub mod completions;
 pub mod error;
 pub mod traits;
 

src/lib.rs

@@ -1,6 +1,8 @@
 #![allow(async_fn_in_trait)]
 
 use crate::api_key::ApiKeyInfo;
+use crate::completions::ChatCompletionRequest;
+use crate::completions::ChatCompletionResponse;
 use crate::error::XaiError;
 use reqwest::{Method, RequestBuilder};
 
@@ -13,3 +15,10 @@ pub trait ClientConfig {
 pub trait ApiKeyFetcher {
     async fn fetch_api_key_info(&self) -> Result<ApiKeyInfo, XaiError>;
 }
+
+pub trait ChatCompletionsFetcher {
+    async fn create_chat_completion(
+        &self,
+        request: ChatCompletionRequest,
+    ) -> Result<ChatCompletionResponse, XaiError>;
+}

src/traits.rs

@@ -0,0 +1,124 @@
+use mockito::{Matcher, Server};
+use reqwest::Method;
+use serde_json::json;
+use x_ai::client::XaiClient;
+use x_ai::traits::ClientConfig;
+
+#[tokio::test]
+async fn test_chat_completions() {
+    let mut server = Server::new_async().await;
+
+    let chat_completion_mock = server
+        .mock("POST", "/v1/chat/completions")
+        .match_header("Content-Type", "application/json")
+        .match_body(Matcher::JsonString(r#"
+            {
+                "messages": [
+                    {
+                        "role": "system",
+                        "content": "You are Grok, a chatbot inspired by the Hitchhikers Guide to the Galaxy."
+                    },
+                    {
+                        "role": "user",
+                        "content": "What is the answer to life and universe?"
+                    }
+                ],
+                "model": "grok-beta",
+                "stream": false,
+                "temperature": 0
+            }
+        "#.to_string()))
+        .with_status(200)
+        .with_header("Content-Type", "application/json")
+        .with_body(r#"
+            {
+              "id": "304e12ef-81f4-4e93-a41c-f5f57f6a2b56",
+              "object": "chat.completion",
+              "created": 1728511727,
+              "model": "grok-beta",
+              "choices": [
+                {
+                  "index": 0,
+                  "message": {
+                    "role": "assistant",
+                    "content": "The answer to the ultimate question of life, the universe, and everything is **42**, according to Douglas Adams science fiction series \"The Hitchhiker's Guide to the Galaxy.\" This number is often humorously referenced in discussions about the meaning of life. However, in the context of the story, the actual question to which 42 is the answer remains unknown, symbolizing the ongoing search for understanding the purpose or meaning of existence."
+                  },
+                  "finish_reason": "stop"
+                }
+              ],
+              "usage": {
+                "prompt_tokens": 24,
+                "completion_tokens": 91,
+                "total_tokens": 115
+              },
+              "system_fingerprint": "fp_3813298403"
+            }
+        "#)
+        .create_async()
+        .await;
+
+    let client = XaiClient::builder()
+        .base_url(&format!("{}/", server.url()))
+        .build()
+        .expect("Failed to build XaiClient");
+
+    client.set_api_key("test-api-key".to_string());
+
+    let body = json!({
+        "messages": [
+            {
+                "role": "system",
+                "content": "You are Grok, a chatbot inspired by the Hitchhikers Guide to the Galaxy."
+            },
+            {
+                "role": "user",
+                "content": "What is the answer to life and universe?"
+            }
+        ],
+        "model": "grok-beta",
+        "stream": false,
+        "temperature": 0
+    });
+
+    let result = client
+        .request(Method::POST, "/v1/chat/completions")
+        .expect("body")
+        .json(&body)
+        .send()
+        .await;
+
+    assert!(result.is_ok());
+    let response = result.unwrap();
+    assert_eq!(response.status(), 200);
+
+    let response_text = response.text().await.unwrap();
+    assert_eq!(
+        response_text,
+        r#"
+            {
+              "id": "304e12ef-81f4-4e93-a41c-f5f57f6a2b56",
+              "object": "chat.completion",
+              "created": 1728511727,
+              "model": "grok-beta",
+              "choices": [
+                {
+                  "index": 0,
+                  "message": {
+                    "role": "assistant",
+                    "content": "The answer to the ultimate question of life, the universe, and everything is **42**, according to Douglas Adams science fiction series \"The Hitchhiker's Guide to the Galaxy.\" This number is often humorously referenced in discussions about the meaning of life. However, in the context of the story, the actual question to which 42 is the answer remains unknown, symbolizing the ongoing search for understanding the purpose or meaning of existence."
+                  },
+                  "finish_reason": "stop"
+                }
+              ],
+              "usage": {
+                "prompt_tokens": 24,
+                "completion_tokens": 91,
+                "total_tokens": 115
+              },
+              "system_fingerprint": "fp_3813298403"
+            }
+        "#
+    );
+
+    chat_completion_mock.assert_async().await;
+}