Merge pull request #28 from maschad/mc/fix/chat-completion-chunk

fix: add missing fields to chat completion chunk

Author: maschad

src/atoma_sdk/models/chatcompletionchunk.py

@@ -5,6 +5,7 @@
     ChatCompletionChunkChoice,
     ChatCompletionChunkChoiceTypedDict,
 )
+from .usage import Usage, UsageTypedDict
 from atoma_sdk.types import BaseModel
 from typing import List
 from typing_extensions import TypedDict
@@ -19,6 +20,14 @@ class ChatCompletionChunkTypedDict(TypedDict):
     r"""A unique identifier for the chat completion chunk."""
     model: str
     r"""The model used for the chat completion."""
+    id: str
+    r"""The unique identifier for the chat completion."""
+    system_fingerprint: str | None = None
+    r"""The system fingerprint for the chat completion."""
+    usage: UsageTypedDict | None = None
+    r"""The usage information for the chat completion."""
+    service_tier: str | None = None
+    r"""The service tier for the chat completion."""
 
 
 class ChatCompletionChunk(BaseModel):
@@ -33,3 +42,13 @@ class ChatCompletionChunk(BaseModel):
 
     model: str
     r"""The model used for the chat completion."""
+
+    system_fingerprint: str | None = None
+    r"""The system fingerprint for the chat completion."""
+
+    usage: Usage | None = None
+    r"""The usage information for the chat completion."""
+
+    service_tier: str | None = None
+    r"""The service tier for the chat completion."""
+

src/atoma_sdk/models/chatcompletionresponse.py

@@ -21,6 +21,11 @@ class ChatCompletionResponseTypedDict(TypedDict):
     system_fingerprint: NotRequired[Nullable[str]]
     r"""The system fingerprint for the completion, if applicable."""
     usage: NotRequired[Nullable[CompletionUsageTypedDict]]
+    r"""The usage information for the chat completion."""
+    service_tier: NotRequired[Nullable[str]]
+    r"""The service tier for the chat completion."""
+    object: NotRequired[Nullable[str]]
+    r"""The object type for the chat completion."""
 
 
 class ChatCompletionResponse(BaseModel):
@@ -40,6 +45,13 @@ class ChatCompletionResponse(BaseModel):
     r"""The system fingerprint for the completion, if applicable."""
 
     usage: OptionalNullable[CompletionUsage] = UNSET
+    r"""The usage information for the chat completion."""
+
+    service_tier: OptionalNullable[str] = UNSET
+    r"""The service tier for the chat completion."""
+
+    object: OptionalNullable[str] = UNSET
+    r"""The object type for the chat completion."""
 
     @model_serializer(mode="wrap")
     def serialize_model(self, handler):

test/test_chat_completions.py

@@ -27,7 +27,7 @@ def test_chat_completion_basic(client):
     )
 
     print(completion.choices[0].message)
-    
+
     assert completion is not None
     assert len(completion.choices) > 0
     assert completion.choices[0].message.content is not None
@@ -40,7 +40,7 @@ def test_chat_completion_with_system_message(client):
             {"role": "user", "content": "Hello!"}
         ]
     )
-    
+
     assert completion is not None
     assert len(completion.choices) > 0
     assert completion.choices[0].message.content is not None
@@ -54,7 +54,7 @@ async def test_chat_completion_async(client):
             {"role": "user", "content": "Hello!"}
         ]
     )
-    
+
     assert completion is not None
     assert len(completion.choices) > 0
     assert completion.choices[0].message.content is not None
@@ -70,15 +70,19 @@ def test_chat_completion_stream(client):
 
     # Verify we get a valid stream
     assert completion is not None
-    
+
     # Process the stream and verify chunks
     chunk_count = 0
     for chunk in completion:
         assert chunk is not None
         assert chunk.data is not None
         assert len(chunk.data.choices) > 0
         assert chunk.data.choices[0].delta is not None
+        assert chunk.data.system_fingerprint is not None
+        assert chunk.data.usage is not None
+        assert chunk.data.service_tier is not None
         chunk_count += 1
-    
+
     # Verify we got multiple chunks
     assert chunk_count > 0
+