fix: consolidate API response on streams (#26)

Author: maschad

src/atoma_sdk/confidential_chat.py

@@ -173,7 +173,7 @@ def create(
                     encrypted_message=encrypted_response
                 )
                 return utils.unmarshal_json(
-                    decrypted_response.decode('utf-8'), 
+                    decrypted_response.decode('utf-8'),
                     models.ChatCompletionResponse
                 )
             except Exception as e:
@@ -362,7 +362,7 @@ async def create_async(
                     encrypted_message=encrypted_response
                 )
                 return utils.unmarshal_json(
-                    decrypted_response.decode('utf-8'), 
+                    decrypted_response.decode('utf-8'),
                     models.ChatCompletionResponse
                 )
             except Exception as e:
@@ -554,7 +554,8 @@ def decrypt_chunk(raw_chunk):
                     # Skip chunks with empty choices
                     if not decrypted_json.get('choices'):
                         return None
-                    return models.ChatCompletionChunk.model_validate(decrypted_json)
+                    # Wrap the chunk in a StreamResponse to maintain consistent API
+                    return models.ChatCompletionStreamResponse(data=models.ChatCompletionChunk.model_validate(decrypted_json))
                 except Exception as e:
                     raise models.APIError(f"Failed to decrypt stream chunk: {str(e)}", 500, str(e), None)
 
@@ -740,7 +741,9 @@ async def decrypt_chunk(raw_chunk):
                         salt=salt,
                         encrypted_message=encrypted_chunk.data
                     )
-                    return utils.unmarshal_json(decrypted_chunk.decode('utf-8'), models.ChatCompletionStreamResponse)
+                    decrypted_json = json.loads(decrypted_chunk.decode('utf-8'))
+                    # Wrap the chunk in a StreamResponse to maintain consistent API
+                    return models.ChatCompletionStreamResponse(data=models.ChatCompletionChunk.model_validate(decrypted_json))
                 except Exception as e:
                     raise models.APIError(f"Failed to decrypt stream chunk: {str(e)}", 500, str(e), None)
 

test/test_confidential_chat.py

@@ -26,11 +26,11 @@ def test_chat_completion_basic(client):
         ]
     )
 
-    print(completion.choices[0].message)
-    
+    print(completion.data.choices[0].message)
+
     assert completion is not None
-    assert len(completion.choices) > 0
-    assert completion.choices[0].message.content is not None
+    assert len(completion.data.choices) > 0
+    assert completion.data.choices[0].message.content is not None
 
 def test_chat_completion_with_system_message(client):
     completion = client.confidential_chat.create(
@@ -40,10 +40,10 @@ 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
+    assert len(completion.data.choices) > 0
+    assert completion.data.choices[0].message.content is not None
 
 @pytest.mark.asyncio
 async def test_chat_completion_async(client):
@@ -54,10 +54,10 @@ 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
+    assert len(completion.data.choices) > 0
+    assert completion.data.choices[0].message.content is not None
 
 def test_chat_completion_stream(client):
     completion = client.confidential_chat.create_stream(
@@ -70,15 +70,15 @@ 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 is not None
-        assert len(chunk.choices) > 0
-        assert chunk.choices[0].delta is not None
+        assert chunk.data is not None
+        assert len(chunk.data.choices) > 0
+        assert chunk.data.choices[0].delta is not None
         chunk_count += 1
-    
+
     # Verify we got multiple chunks
     assert chunk_count > 0