Return usage in response

Author: intitni

Tool/Sources/Logger/Logger.swift

@@ -23,6 +23,7 @@ public final class Logger {
     public static let license = Logger(category: "License")
     public static let `extension` = Logger(category: "Extension")
     public static let communicationBridge = Logger(category: "CommunicationBridge")
+    public static let chatProxy = Logger(category: "ChatProxy")
     public static let debug = Logger(category: "Debug")
     #if DEBUG
     /// Use a temp logger to log something temporary. I won't be available in release builds.

Tool/Sources/OpenAIService/ChatGPTService.swift

@@ -68,6 +68,12 @@ public enum ChatGPTResponse: Equatable {
     case partialText(String)
     case partialReasoning(String)
     case toolCalls([ChatMessage.ToolCall])
+    case usage(
+        promptTokens: Int,
+        completionTokens: Int,
+        cachedTokens: Int,
+        otherUsage: [String: Int]
+    )
 }
 
 public typealias ChatGPTResponseStream = AsyncThrowingStream<ChatGPTResponse, any Error>
@@ -196,7 +202,7 @@ public class ChatGPTService: ChatGPTServiceType {
                                 switch content {
                                 case let .partialText(text):
                                     continuation.yield(ChatGPTResponse.partialText(text))
-                                    
+
                                 case let .partialReasoning(text):
                                     continuation.yield(ChatGPTResponse.partialReasoning(text))
 
@@ -222,6 +228,20 @@ public class ChatGPTService: ChatGPTServiceType {
                                             }
                                         }
                                     }
+                                case let .usage(
+                                    promptTokens,
+                                    completionTokens,
+                                    cachedTokens,
+                                    otherUsage
+                                ):
+                                    continuation.yield(
+                                        .usage(
+                                            promptTokens: promptTokens,
+                                            completionTokens: completionTokens,
+                                            cachedTokens: cachedTokens,
+                                            otherUsage: otherUsage
+                                        )
+                                    )
                                 }
                             }
 
@@ -257,6 +277,12 @@ extension ChatGPTService {
         case partialReasoning(String)
         case partialText(String)
         case partialToolCalls([Int: ChatMessage.ToolCall])
+        case usage(
+            promptTokens: Int,
+            completionTokens: Int,
+            cachedTokens: Int,
+            otherUsage: [String: Int]
+        )
     }
 
     enum FunctionCallResult {
@@ -345,14 +371,19 @@ extension ChatGPTService {
                         if let content = delta.content {
                             continuation.yield(.partialText(content))
                         }
-                        
+
                         if let reasoning = delta.reasoningContent {
                             continuation.yield(.partialReasoning(reasoning))
                         }
                     }
 
                     Logger.service.info("ChatGPT usage: \(usage)")
-
+                    continuation.yield(.usage(
+                        promptTokens: usage.promptTokens,
+                        completionTokens: usage.completionTokens,
+                        cachedTokens: usage.cachedTokens,
+                        otherUsage: usage.otherUsage
+                    ))
                     continuation.finish()
                 } catch let error as CancellationError {
                     continuation.finish(throwing: error)
@@ -554,6 +585,7 @@ extension ChatGPTService {
                     case .system: .system
                     case .user: .user
                     case .assistant: .assistant
+                    case .tool: .tool
                     }
                 }(),
                 content: chatMessage.content ?? "",

Tool/Sources/OpenAIService/Configuration/ChatGPTConfiguration.swift

@@ -39,7 +39,7 @@ public extension ChatGPTConfiguration {
     }
 }
 
-public class OverridingChatGPTConfiguration: ChatGPTConfiguration {
+public final class OverridingChatGPTConfiguration: ChatGPTConfiguration {
     public struct Overriding: Codable {
         public var temperature: Double?
         public var modelId: String?