Skip to content

Commit

Permalink
Added finishReason and usage info to the ChatCompletionResponse.
Browse files Browse the repository at this point in the history
  • Loading branch information
btfranklin committed Aug 8, 2023
1 parent 7a04e9b commit 533b3fd
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 12 deletions.
32 changes: 29 additions & 3 deletions Sources/CleverBird/chat/ChatCompletionResponse.swift
Expand Up @@ -3,36 +3,62 @@
struct ChatCompletionResponse: Codable, Identifiable {

struct Choice: Codable {

enum FinishReason: String, Codable {
case stop
case maxTokens
}

let message: ChatMessage
let finishReason: FinishReason
let functionCall: FunctionCall?

enum CodingKeys: String, CodingKey {
case message
case finishReason
case functionCall
}
}

struct Usage: Codable {
let promptTokens: Int
let completionTokens: Int
let totalTokens: Int
}

let choices: [Choice]
let usage: Usage
let id: String

enum CodingKeys: String, CodingKey {
case choices
case usage
case id
}

init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
id = try container.decode(String.self, forKey: .id)

self.id = try container.decode(String.self, forKey: .id)

var choicesContainer = try container.nestedUnkeyedContainer(forKey: .choices)
var choices: [Choice] = []
while !choicesContainer.isAtEnd {
let choiceContainer = try choicesContainer.nestedContainer(keyedBy: Choice.CodingKeys.self)

var message = try choiceContainer.decode(ChatMessage.self, forKey: .message)
let functionCall = try choiceContainer.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
message.id = id
let choice = Choice(message: message, functionCall: functionCall)

let finishReason = try choiceContainer.decode(Choice.FinishReason.self, forKey: .finishReason)

let functionCall = try choiceContainer.decodeIfPresent(FunctionCall.self, forKey: .functionCall)

let choice = Choice(message: message, finishReason: finishReason, functionCall: functionCall)
choices.append(choice)
}
self.choices = choices

self.usage = try container.decode(Usage.self, forKey: .usage)
}

}
18 changes: 9 additions & 9 deletions Sources/CleverBird/chat/ChatMessage.swift
Expand Up @@ -69,19 +69,19 @@ public struct ChatMessage: Codable, Identifiable {

public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
role = try container.decode(Role.self, forKey: .role)
content = try container.decodeIfPresent(String.self, forKey: .content)
functionCall = try container.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
name = try container.decodeIfPresent(String.self, forKey: .name)
id = "pending"
self.role = try container.decode(Role.self, forKey: .role)
self.content = try container.decodeIfPresent(String.self, forKey: .content)
self.functionCall = try container.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
self.name = try container.decodeIfPresent(String.self, forKey: .name)
self.id = "pending"
}

public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
try container.encode(role, forKey: .role)
try container.encodeIfPresent(content, forKey: .content)
try container.encodeIfPresent(functionCall, forKey: .functionCall)
try container.encodeIfPresent(name, forKey: .name)
try container.encode(self.role, forKey: .role)
try container.encodeIfPresent(self.content, forKey: .content)
try container.encodeIfPresent(self.functionCall, forKey: .functionCall)
try container.encodeIfPresent(self.name, forKey: .name)
}
}

Expand Down

0 comments on commit 533b3fd

Please sign in to comment.