diff --git a/Sources/OpenAI/Public/Models/ChatQuery.swift b/Sources/OpenAI/Public/Models/ChatQuery.swift index 1975ed5c..9b43c7ac 100644 --- a/Sources/OpenAI/Public/Models/ChatQuery.swift +++ b/Sources/OpenAI/Public/Models/ChatQuery.swift @@ -74,6 +74,8 @@ public struct ChatQuery: Equatable, Codable, Streamable { /// If set, partial message deltas will be sent, like in ChatGPT. Tokens will be sent as data-only server-sent events as they become available, with the stream terminated by a data: [DONE] message. /// https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format public var stream: Bool + /// Options for streaming response. Only set this when you set stream: true. + public var streamOptions: Self.StreamOptions? public init( messages: [Self.ChatCompletionMessageParam], @@ -95,7 +97,8 @@ public struct ChatQuery: Equatable, Codable, Streamable { topLogprobs: Int? = nil, topP: Double? = nil, user: String? = nil, - stream: Bool = false + stream: Bool = false, + streamOptions: StreamOptions? = nil ) { self.messages = messages self.model = model @@ -117,6 +120,7 @@ public struct ChatQuery: Equatable, Codable, Streamable { self.topP = topP self.user = user self.stream = stream + self.streamOptions = streamOptions } public enum ChatCompletionMessageParam: Codable, Equatable { @@ -1149,6 +1153,24 @@ public struct ChatQuery: Equatable, Codable, Streamable { case function } } + + public struct StreamOptions: Codable, Equatable { + + /// If set, an additional chunk will be streamed before the data: [DONE] message. + /// The usage field on this chunk shows the token usage statistics for the entire request, + /// and the choices field will always be an empty array. All other chunks will also + /// include a usage field, but with a null value. + public let includeUsage: Bool + + public init(includeUsage: Bool) { + self.includeUsage = includeUsage + } + + public enum CodingKeys: String, CodingKey { + case includeUsage = "include_usage" + } + + } public enum CodingKeys: String, CodingKey { case messages @@ -1171,6 +1193,7 @@ public struct ChatQuery: Equatable, Codable, Streamable { case topP = "top_p" case user case stream + case streamOptions = "stream_options" } } diff --git a/Sources/OpenAI/Public/Models/ChatStreamResult.swift b/Sources/OpenAI/Public/Models/ChatStreamResult.swift index 3457b089..c8819117 100644 --- a/Sources/OpenAI/Public/Models/ChatStreamResult.swift +++ b/Sources/OpenAI/Public/Models/ChatStreamResult.swift @@ -130,6 +130,8 @@ public struct ChatStreamResult: Codable, Equatable { public let choices: [Choice] /// This fingerprint represents the backend configuration that the model runs with. Can be used in conjunction with the `seed` request parameter to understand when backend changes have been made that might impact determinism. public let systemFingerprint: String? + /// Usage statistics for the completion request. + public let usage: ChatResult.CompletionUsage? public enum CodingKeys: String, CodingKey { case id @@ -138,5 +140,6 @@ public struct ChatStreamResult: Codable, Equatable { case model case choices case systemFingerprint = "system_fingerprint" + case usage } } diff --git a/Tests/OpenAITests/OpenAITestsDecoder.swift b/Tests/OpenAITests/OpenAITestsDecoder.swift index 681af44d..4006a4d9 100644 --- a/Tests/OpenAITests/OpenAITestsDecoder.swift +++ b/Tests/OpenAITests/OpenAITestsDecoder.swift @@ -156,6 +156,55 @@ class OpenAITestsDecoder: XCTestCase { XCTAssertEqual(chatQueryAsDict, expectedValueAsDict) } + + func testChatQueryWithStreamOptions() async throws { + let chatQuery = ChatQuery(messages: [ + .init(role: .user, content: "Who are you?")! + ], model: .gpt4, stream: true, streamOptions: .init(includeUsage: true)) + let expectedValue = """ + { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Who are you?" + } + ], + "stream": true, + "stream_options": { + "include_usage" : true + } + } + """ + + let chatQueryAsDict = try jsonDataAsNSDictionary(JSONEncoder().encode(chatQuery)) + let expectedValueAsDict = try jsonDataAsNSDictionary(expectedValue.data(using: .utf8)!) + + XCTAssertEqual(chatQueryAsDict, expectedValueAsDict) + } + + func testChatQueryWithoutStreamOptions() async throws { + let chatQuery = ChatQuery(messages: [ + .init(role: .user, content: "Who are you?")! + ], model: .gpt4, stream: true) + let expectedValue = """ + { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Who are you?" + } + ], + "stream": true + } + """ + + let chatQueryAsDict = try jsonDataAsNSDictionary(JSONEncoder().encode(chatQuery)) + let expectedValueAsDict = try jsonDataAsNSDictionary(expectedValue.data(using: .utf8)!) + + XCTAssertEqual(chatQueryAsDict, expectedValueAsDict) + } func testChatQueryWithFunctionCall() async throws { let chatQuery = ChatQuery(