From ca39c28447f22ee886659cb3933f82d6b1fd5f9c Mon Sep 17 00:00:00 2001
From: Alexander Chen <alexander.qili.chen@gmail.com>
Date: Tue, 15 Apr 2025 14:35:44 -0400
Subject: [PATCH] Correctly handled response headers in speak rest client for
 aura 2

---
 deepgram/clients/speak/v1/rest/client.py | 67 +++++++++++++++++-------
 1 file changed, 47 insertions(+), 20 deletions(-)

diff --git a/deepgram/clients/speak/v1/rest/client.py b/deepgram/clients/speak/v1/rest/client.py
index e17877e4..b86eb0cf 100644
--- a/deepgram/clients/speak/v1/rest/client.py
+++ b/deepgram/clients/speak/v1/rest/client.py
@@ -160,15 +160,28 @@ def stream_memory(
         self._logger.info("addons: %s", addons)
         self._logger.info("headers: %s", headers)
 
-        return_vals = [
-            "content-type",
-            "request-id",
-            "model-uuid",
-            "model-name",
-            "char-count",
-            "transfer-encoding",
-            "date",
-        ]
+        is_aura_2_model = options["model"].split("-")[1] == "2"
+
+        if is_aura_2_model:
+            return_vals = [
+                "content-type",
+                "request-id",
+                "model-name",
+                "characters",
+                "transfer-encoding",
+                "date",
+            ]
+        else:
+            return_vals = [
+                "content-type",
+                "request-id",
+                "model-uuid",
+                "model-name",
+                "char-count",
+                "transfer-encoding",
+                "date",
+            ]
+
         result = self.post_memory(
             url,
             options=options,
@@ -181,17 +194,31 @@ def stream_memory(
         )
 
         self._logger.info("result: %s", result)
-        resp = SpeakRESTResponse(
-            content_type=str(result["content-type"]),
-            request_id=str(result["request-id"]),
-            model_uuid=str(result["model-uuid"]),
-            model_name=str(result["model-name"]),
-            characters=int(str(result["char-count"])),
-            transfer_encoding=str(result["transfer-encoding"]),
-            date=str(result["date"]),
-            stream=cast(io.BytesIO, result["stream"]),
-            stream_memory=cast(io.BytesIO, result["stream"]),
-        )
+
+        if is_aura_2_model:
+            resp = SpeakRESTResponse(
+                content_type=str(result["content-type"]),
+                request_id=str(result["request-id"]),
+                model_name=str(result["model-name"]),
+                characters=int(str(result["characters"])),
+                transfer_encoding=str(result["transfer-encoding"]),
+                date=str(result["date"]),
+                stream=cast(io.BytesIO, result["stream"]),
+                stream_memory=cast(io.BytesIO, result["stream"]),
+            )
+        else:
+            resp = SpeakRESTResponse(
+                content_type=str(result["content-type"]),
+                request_id=str(result["request-id"]),
+                model_uuid=str(result["model-uuid"]),
+                model_name=str(result["model-name"]),
+                characters=int(str(result["char-count"])),
+                transfer_encoding=str(result["transfer-encoding"]),
+                date=str(result["date"]),
+                stream=cast(io.BytesIO, result["stream"]),
+                stream_memory=cast(io.BytesIO, result["stream"]),
+            )
+
         self._logger.verbose("resp Object: %s", resp)
         self._logger.notice("speak succeeded")
         self._logger.debug("SpeakClient.stream LEAVE")