Merge pull request #8 from hyperlink/fix-streaming-nil-output-token

fumito-ito · Apr 2, 2024 · 21996e3 · 21996e3
2 parents 90c9432 + 2bf95f7
commit 21996e3
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/Sources/AnthropicSwiftSDK/Entity/Streaming/MessageDelta.swift b/Sources/AnthropicSwiftSDK/Entity/Streaming/MessageDelta.swift
@@ -10,5 +10,4 @@ import Foundation
 public struct MessageDelta: Decodable {
     public let stopReason: StopReason
     public let stopSequence: String?
-    public let usage: TokenUsage?
 }
diff --git a/Sources/AnthropicSwiftSDK/Network/StreamingResponse.swift b/Sources/AnthropicSwiftSDK/Network/StreamingResponse.swift
@@ -23,6 +23,7 @@ public struct StreamingMessageStartResponse: StreamingResponse {
 public struct StreamingMessageDeltaResponse: StreamingResponse {
     public let type: StreamingEvent
     public let delta: MessageDelta
+    public let usage: TokenUsage
 }
 
 public struct StreamingMessageStopResponse: StreamingResponse {

diff --git a/Tests/AnthropicSwiftSDKTests/Network/StreamingParser/StreamingDataLineParserTests.swift b/Tests/AnthropicSwiftSDKTests/Network/StreamingParser/StreamingDataLineParserTests.swift
@@ -30,15 +30,15 @@ final class StreamingDataLineParserTests: XCTestCase {
 
     func testParseMessageDeltaDataLine() throws {
         let line = """
-        data: {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence":null, "usage":{"output_tokens": 15}}}
+        data: {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence":null}, "usage":{"output_tokens": 15}}
         """
         let result = try StreamingDataLineParser.parse(dataLine: line) as StreamingMessageDeltaResponse
 
         XCTAssertEqual(result.type, .messageDelta)
         XCTAssertEqual(result.delta.stopReason, .endTurn)
         XCTAssertNil(result.delta.stopSequence)
-        XCTAssertEqual(result.delta.usage?.outputTokens, 15)
-        XCTAssertNil(result.delta.usage?.inputTokens)
+        XCTAssertEqual(result.usage.outputTokens, 15)
+        XCTAssertNil(result.usage.inputTokens)
     }
 
     func testParseMessageStopDataLine() throws {