Assistants API support for Azure (#41)

* Azure assistants API * Threads * Runs * Run step * stream * Vector store * Changes on parameters * Vector store file
jamesrochabrun · May 29, 2024 · 7821552 · 7821552
1 parent 352dc41
commit 7821552
Show file tree

Hide file tree

Showing 10 changed files with 590 additions and 79 deletions.
diff --git a/Sources/OpenAI/AIProxy/AIProxyService.swift b/Sources/OpenAI/AIProxy/AIProxyService.swift
@@ -653,10 +653,11 @@ struct AIProxyService: OpenAIService {
  }
 
  func modifyVectorStore(
+ parameters: VectorStoreParameter,
  id: String)
  async throws -> VectorStoreObject
  {
- let request = try await OpenAIAPI.vectorStore(.modify(vectorStoreID: id)).request(aiproxyPartialKey: partialKey, organizationID: organizationID, method: .post, betaHeaderField: Self.assistantsBetaV2, deviceCheckBypass: deviceCheckBypass)
+ let request = try await OpenAIAPI.vectorStore(.modify(vectorStoreID: id)).request(aiproxyPartialKey: partialKey, organizationID: organizationID, method: .post, params: parameters, betaHeaderField: Self.assistantsBetaV2, deviceCheckBypass: deviceCheckBypass)
  return try await fetch(type: VectorStoreObject.self, with: request)
  }
 

diff --git a/Sources/OpenAI/Azure/AzureOpenAIAPI.swift b/Sources/OpenAI/Azure/AzureOpenAIAPI.swift
@@ -13,20 +13,127 @@ enum AzureOpenAIAPI {
 
  static var azureOpenAIResource: String = ""
 
- case chat(deploymentID: String) // https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference?tabs=python
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/assistant
+ case assistant(AssistantCategory)
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions
+ case chat(deploymentID: String)
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-messages?tabs=python
+ case message(MessageCategory)
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-runs?tabs=python
+ case run(RunCategory)
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-runs?tabs=python#list-run-steps
+ case runStep(RunStepCategory)
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-threads?tabs=python#create-a-thread
+ case thread(ThreadCategory)
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/file-search?tabs=python#vector-stores
+ case vectorStore(VectorStoreCategory)
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/file-search?tabs=python#vector-stores
+ case vectorStoreFile(VectorStoreFileCategory)
+
+ enum AssistantCategory {
+ case create
+ case list
+ case retrieve(assistantID: String)
+ case modify(assistantID: String)
+ case delete(assistantID: String)
+ }
+
+ enum MessageCategory {
+ case create(threadID: String)
+ case retrieve(threadID: String, messageID: String)
+ case modify(threadID: String, messageID: String)
+ case list(threadID: String)
+ }
+
+ enum RunCategory {
+ case create(threadID: String)
+ case retrieve(threadID: String, runID: String)
+ case modify(threadID: String, runID: String)
+ case list(threadID: String)
+ case cancel(threadID: String, runID: String)
+ case submitToolOutput(threadID: String, runID: String)
+ case createThreadAndRun
+ }
+
+ enum RunStepCategory {
+ case retrieve(threadID: String, runID: String, stepID: String)
+ case list(threadID: String, runID: String)
+ }
+
+ enum ThreadCategory {
+ case create
+ case retrieve(threadID: String)
+ case modify(threadID: String)
+ case delete(threadID: String)
+ }
+
+ enum VectorStoreCategory {
+ case create
+ case list
+ case retrieve(vectorStoreID: String)
+ case modify(vectorStoreID: String)
+ case delete(vectorStoreID: String)
+ }
+
+ /// https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/file-search?tabs=python#file-search-support
+ enum VectorStoreFileCategory {
+ case create(vectorStoreID: String)
+ case list(vectorStoreID: String)
+ case retrieve(vectorStoreID: String, fileID: String)
+ case delete(vectorStoreID: String, fileID: String)
+ }
 }
 
 // MARK: Endpoint
 
 extension AzureOpenAIAPI: Endpoint {
 
  var base: String {
- "https://\(Self.azureOpenAIResource).openai.azure.com"
+ "https://\(Self.azureOpenAIResource)/openai.azure.com"
  }
 
  var path: String {
  switch self {
- case .chat(let deploymentID): "/openai/deployments/\(deploymentID)/chat/completions"
+ case .chat(let deploymentID): return "/openai/deployments/\(deploymentID)/chat/completions"
+ case .assistant(let category):
+ switch category {
+ case .create, .list: return "/openai/assistants"
+ case .retrieve(let assistantID), .modify(let assistantID), .delete(let assistantID): return "/openai/assistants/\(assistantID)"
+ }
+ case .message(let category):
+ switch category {
+ case .create(let threadID), .list(let threadID): return "/openai/threads/\(threadID)/messages"
+ case .retrieve(let threadID, let messageID), .modify(let threadID, let messageID): return "/openai/threads/\(threadID)/messages/\(messageID)"
+ }
+ case .run(let category):
+ switch category {
+ case .create(let threadID), .list(let threadID): return "/openai/threads/\(threadID)/runs"
+ case .retrieve(let threadID, let runID), .modify(let threadID, let runID): return "/openai/threads/\(threadID)/runs/\(runID)"
+ case .cancel(let threadID, let runID): return "/openai/threads/\(threadID)/runs/\(runID)/cancel"
+ case .submitToolOutput(let threadID, let runID): return "/openai/threads/\(threadID)/runs/\(runID)//submit_tool_outputs"
+ case .createThreadAndRun: return "/openai/threads/runs"
+ }
+ case .runStep(let category):
+ switch category {
+ case .retrieve(let threadID, let runID, let stepID): return "/openai/threads/\(threadID)/runs/\(runID)/steps/\(stepID)"
+ case .list(let threadID, let runID): return "/openai/threads/\(threadID)/runs/\(runID)/steps"
+ }
+ case .thread(let category):
+ switch category {
+ case .create: return "/openai/threads"
+ case .retrieve(let threadID), .modify(let threadID), .delete(let threadID): return "/openai/threads/\(threadID)"
+ }
+ case .vectorStore(let category):
+ switch category {
+ case .create, .list: return "/openai/vector_stores"
+ case .retrieve(let vectorStoreID), .modify(let vectorStoreID), .delete(let vectorStoreID): return "/openai/vector_stores/\(vectorStoreID)"
+ }
+ case .vectorStoreFile(let category):
+ switch category {
+ case .create(let vectorStoreID), .list(let vectorStoreID): return "/openai/vector_stores/\(vectorStoreID)/files"
+ case .retrieve(let vectorStoreID, let fileID), .delete(let vectorStoreID, let fileID): return "/openai/vector_stores/\(vectorStoreID)/files/\(fileID)"
+ }
  }
  }
 }
diff --git a/Sources/OpenAI/Azure/AzureOpenAIConfiguration.swift b/Sources/OpenAI/Azure/AzureOpenAIConfiguration.swift
@@ -19,13 +19,18 @@ public struct AzureOpenAIConfiguration {
  /// The API version to use for this operation. This follows the YYYY-MM-DD format.
  let apiVersion: String
 
+ /// Azure configuration extra headers for a request.
+ let extraHeaders: [String: String]?
+
  public init(
  resourceName: String,
  openAIAPIKey: Authorization,
- apiVersion: String)
+ apiVersion: String,
+ extraHeaders: [String: String]? = nil)
  {
  self.resourceName = resourceName
  self.openAIAPIKey = openAIAPIKey
  self.apiVersion = apiVersion
+ self.extraHeaders = extraHeaders
  }
 }