feat: add chatkit api + custom headers support (#471)

* add chatkit api

* beta() apis

* simplify

* update readme

* add chatkit example

Author: 64bit

async-openai/README.md

@@ -23,24 +23,28 @@
 
 - It's based on [OpenAI OpenAPI spec](https://github.com/openai/openai-openapi)
 - Current features:
-  - [x] Assistants (v2)
+  - [x] Administration (partially implemented)
+  - [x] Assistants (beta)
   - [x] Audio
   - [x] Batch
   - [x] Chat
-  - [x] Completions (Legacy)
+  - [x] ChatKit (beta)
+  - [x] Completions (legacy)
   - [x] Conversations
-  - [x] Containers | Container Files
+  - [x] Containers
   - [x] Embeddings
+  - [x] Evals
   - [x] Files
   - [x] Fine-Tuning
   - [x] Images
   - [x] Models
   - [x] Moderations
-  - [x] Organizations | Administration (partially implemented)
-  - [x] Realtime GA (partially implemented)
+  - [x] Realtime (partially implemented)
   - [x] Responses
   - [x] Uploads
+  - [x] Vector Stores
   - [x] Videos
+  - [x] Webhooks
 - Bring your own custom types for Request or Response objects.
 - SSE streaming on available APIs
 - Requests (except SSE streaming) including form submissions are retried with exponential backoff when [rate limited](https://platform.openai.com/docs/guides/rate-limits).

async-openai/src/chatkit.rs

@@ -0,0 +1,116 @@
+use serde::Serialize;
+
+use crate::{
+    config::Config,
+    error::OpenAIError,
+    types::chatkit::{
+        ChatSessionResource, CreateChatSessionBody, DeletedThreadResource, ThreadItemListResource,
+        ThreadListResource, ThreadResource,
+    },
+    Client,
+};
+
+/// ChatKit API for managing sessions and threads.
+///
+/// Related guide: [ChatKit](https://platform.openai.com/docs/api-reference/chatkit)
+pub struct Chatkit<'c, C: Config> {
+    client: &'c Client<C>,
+}
+
+impl<'c, C: Config> Chatkit<'c, C> {
+    pub fn new(client: &'c Client<C>) -> Self {
+        Self { client }
+    }
+
+    /// Access sessions API.
+    pub fn sessions(&self) -> ChatkitSessions<'_, C> {
+        ChatkitSessions::new(self.client)
+    }
+
+    /// Access threads API.
+    pub fn threads(&self) -> ChatkitThreads<'_, C> {
+        ChatkitThreads::new(self.client)
+    }
+}
+
+/// ChatKit sessions API.
+pub struct ChatkitSessions<'c, C: Config> {
+    client: &'c Client<C>,
+}
+
+impl<'c, C: Config> ChatkitSessions<'c, C> {
+    pub fn new(client: &'c Client<C>) -> Self {
+        Self { client }
+    }
+
+    /// Create a ChatKit session.
+    #[crate::byot(T0 = serde::Serialize, R = serde::de::DeserializeOwned)]
+    pub async fn create(
+        &self,
+        request: CreateChatSessionBody,
+    ) -> Result<ChatSessionResource, OpenAIError> {
+        self.client.post("/chatkit/sessions", request).await
+    }
+
+    /// Cancel a ChatKit session.
+    #[crate::byot(T0 = std::fmt::Display, R = serde::de::DeserializeOwned)]
+    pub async fn cancel(&self, session_id: &str) -> Result<ChatSessionResource, OpenAIError> {
+        self.client
+            .post(
+                &format!("/chatkit/sessions/{session_id}/cancel"),
+                serde_json::json!({}),
+            )
+            .await
+    }
+}
+
+/// ChatKit threads API.
+pub struct ChatkitThreads<'c, C: Config> {
+    client: &'c Client<C>,
+}
+
+impl<'c, C: Config> ChatkitThreads<'c, C> {
+    pub fn new(client: &'c Client<C>) -> Self {
+        Self { client }
+    }
+
+    /// List ChatKit threads.
+    #[crate::byot(T0 = serde::Serialize, R = serde::de::DeserializeOwned)]
+    pub async fn list<Q>(&self, query: &Q) -> Result<ThreadListResource, OpenAIError>
+    where
+        Q: Serialize + ?Sized,
+    {
+        self.client.get_with_query("/chatkit/threads", &query).await
+    }
+
+    /// Retrieve a ChatKit thread.
+    #[crate::byot(T0 = std::fmt::Display, R = serde::de::DeserializeOwned)]
+    pub async fn retrieve(&self, thread_id: &str) -> Result<ThreadResource, OpenAIError> {
+        self.client
+            .get(&format!("/chatkit/threads/{thread_id}"))
+            .await
+    }
+
+    /// Delete a ChatKit thread.
+    #[crate::byot(T0 = std::fmt::Display, R = serde::de::DeserializeOwned)]
+    pub async fn delete(&self, thread_id: &str) -> Result<DeletedThreadResource, OpenAIError> {
+        self.client
+            .delete(&format!("/chatkit/threads/{thread_id}"))
+            .await
+    }
+
+    /// List ChatKit thread items.
+    #[crate::byot(T0 = std::fmt::Display, T1 = serde::Serialize, R = serde::de::DeserializeOwned)]
+    pub async fn list_items<Q>(
+        &self,
+        thread_id: &str,
+        query: &Q,
+    ) -> Result<ThreadItemListResource, OpenAIError>
+    where
+        Q: Serialize + ?Sized,
+    {
+        self.client
+            .get_with_query(&format!("/chatkit/threads/{thread_id}/items"), &query)
+            .await
+    }
+}

async-openai/src/client.rs

@@ -7,6 +7,7 @@ use reqwest_eventsource::{Error as EventSourceError, Event, EventSource, Request
 use serde::{de::DeserializeOwned, Serialize};
 
 use crate::{
+    chatkit::Chatkit,
     config::{Config, OpenAIConfig},
     error::{map_deserialization_error, ApiError, OpenAIError, StreamError, WrappedError},
     file::Files,
@@ -188,6 +189,10 @@ impl<C: Config> Client<C> {
         Evals::new(self)
     }
 
+    pub fn chatkit(&self) -> Chatkit<'_, C> {
+        Chatkit::new(self)
+    }
+
     pub fn config(&self) -> &C {
         &self.config
     }

async-openai/src/config.rs

@@ -3,6 +3,8 @@ use reqwest::header::{HeaderMap, AUTHORIZATION};
 use secrecy::{ExposeSecret, SecretString};
 use serde::Deserialize;
 
+use crate::error::OpenAIError;
+
 /// Default v1 API base url
 pub const OPENAI_API_BASE: &str = "https://api.openai.com/v1";
 /// Organization header
@@ -59,6 +61,8 @@ pub struct OpenAIConfig {
     api_key: SecretString,
     org_id: String,
     project_id: String,
+    #[serde(skip)]
+    custom_headers: HeaderMap,
 }
 
 impl Default for OpenAIConfig {
@@ -70,6 +74,7 @@ impl Default for OpenAIConfig {
                 .into(),
             org_id: Default::default(),
             project_id: Default::default(),
+            custom_headers: HeaderMap::new(),
         }
     }
 }
@@ -104,6 +109,21 @@ impl OpenAIConfig {
         self
     }
 
+    /// Add a custom header that will be included in all requests.
+    /// Headers are merged with existing headers, with custom headers taking precedence.
+    pub fn with_header<K, V>(mut self, key: K, value: V) -> Result<Self, OpenAIError>
+    where
+        K: reqwest::header::IntoHeaderName,
+        V: TryInto<reqwest::header::HeaderValue>,
+        V::Error: Into<reqwest::header::InvalidHeaderValue>,
+    {
+        let header_value = value.try_into().map_err(|e| {
+            OpenAIError::InvalidArgument(format!("Invalid header value: {}", e.into()))
+        })?;
+        self.custom_headers.insert(key, header_value);
+        Ok(self)
+    }
+
     pub fn org_id(&self) -> &str {
         &self.org_id
     }
@@ -134,9 +154,10 @@ impl Config for OpenAIConfig {
                 .unwrap(),
         );
 
-        // hack for Assistants APIs
-        // Calls to the Assistants API require that you pass a Beta header
-        // headers.insert(OPENAI_BETA_HEADER, "assistants=v2".parse().unwrap());
+        // Merge custom headers, with custom headers taking precedence
+        for (key, value) in self.custom_headers.iter() {
+            headers.insert(key, value.clone());
+        }
 
         headers
     }

async-openai/src/lib.rs

@@ -145,6 +145,7 @@ mod audio;
 mod audit_logs;
 mod batches;
 mod chat;
+mod chatkit;
 mod client;
 mod completion;
 pub mod config;
@@ -193,6 +194,7 @@ pub use audio::Audio;
 pub use audit_logs::AuditLogs;
 pub use batches::Batches;
 pub use chat::Chat;
+pub use chatkit::Chatkit;
 pub use client::Client;
 pub use completion::Completions;
 pub use container_files::ContainerFiles;

async-openai/src/types/chatkit/mod.rs

@@ -0,0 +1,5 @@
+mod session;
+mod thread;
+
+pub use session::*;
+pub use thread::*;