feat: multimodal output audio for OpenAi and Litellm

README.md

@@ -33,9 +33,9 @@ uv add scope3ai
 |-------------|-----------------|----|-----|------------------|-----------|------------------|-------------------|
 | Anthropic   | ✅              |    |     |                  |           |                  |
 | Cohere      | ✅              |    |     |                  |           |                  |
-| OpenAI      | ✅              | ✅ | ✅  | ✅               | ✅        | Images/Audio     |
+| OpenAI      | ✅              | ✅ | ✅  | ✅               | ✅        | Images/Audio     | Audio
 | Huggingface | ✅              | ✅ | ✅  | ✅               | ✅        |                  |
-| LiteLLM     | ✅              |  ✅  |✅     |       ✅           |           | Images/Audio     |
+| LiteLLM     | ✅              |  ✅  |✅     |       ✅           |           | Images/Audio     | Audio
 | MistralAi   | ✅              |    |     |                  |           | Images           |
 
 Roadmap:

scope3ai/tracers/litellm/chat.py

@@ -9,7 +9,10 @@
 from scope3ai import Scope3AI
 from scope3ai.api.types import Scope3AIContext, ImpactRow
 from scope3ai.constants import PROVIDERS
-from scope3ai.tracers.utils.multimodal import aggregate_multimodal
+from scope3ai.tracers.utils.multimodal import (
+    aggregate_multimodal,
+    aggregate_multimodal_audio_content_output,
+)
 
 PROVIDER = PROVIDERS.LITELLM.value
 
@@ -72,6 +75,7 @@ def litellm_chat_wrapper_non_stream(
 ) -> ChatCompletion:
     timer_start = time.perf_counter()
     keep_traces = not kwargs.pop("use_always_litellm_tracer", False)
+    modalities = kwargs.get("modalities", [])
     with Scope3AI.get_instance().trace(keep_traces=keep_traces) as tracer:
         response = wrapped(*args, **kwargs)
         if tracer.traces:
@@ -88,6 +92,15 @@ def litellm_chat_wrapper_non_stream(
         request_duration_ms=float(request_latency) * 1000,
         managed_service_id=PROVIDER,
     )
+    if "audio" in modalities:
+        audio_format = kwargs.get("audio", {}).get("format", "mp3")
+        for choice in response.choices:
+            audio_data = getattr(choice.message, "audio")
+            if audio_data:
+                audio_content = audio_data.data
+                aggregate_multimodal_audio_content_output(
+                    audio_content, audio_format, scope3_row
+                )
     messages = args[1] if len(args) > 1 else kwargs.get("messages")
     for message in messages:
         aggregate_multimodal(message, scope3_row, logger)

scope3ai/tracers/openai/chat.py

@@ -11,7 +11,10 @@
 from scope3ai.api.types import ImpactRow, Scope3AIContext
 from scope3ai.constants import PROVIDERS
 from scope3ai.lib import Scope3AI
-from scope3ai.tracers.utils.multimodal import aggregate_multimodal
+from scope3ai.tracers.utils.multimodal import (
+    aggregate_multimodal,
+    aggregate_multimodal_audio_content_output,
+)
 
 PROVIDER = PROVIDERS.OPENAI.value
 
@@ -33,18 +36,29 @@ class ChatCompletionChunk(_ChatCompletionChunk):
 def _openai_chat_wrapper(
     response: Any, request_latency: float, kwargs: dict
 ) -> ChatCompletion:
-    model_requested = kwargs["model"]
+    model_requested = kwargs.get("model")
+    modalities = kwargs.get("modalities", [])
     if type(response) is _LegacyAPIResponse:
         http_response = response.http_response.json()
         model_used = http_response.get("model")
         scope3_row = ImpactRow(
             model_id=model_requested,
             model_used_id=model_used,
-            input_tokens=http_response.get("usage").get("prompt_tokens"),
-            output_tokens=http_response.get("usage").get("completion_tokens"),
+            input_tokens=http_response.get("usage", {}).get("prompt_tokens"),
+            output_tokens=http_response.get("usage", {}).get("completion_tokens"),
             request_duration_ms=request_latency * 1000,
             managed_service_id=PROVIDER,
         )
+        if "audio" in modalities:
+            audio_format = kwargs.get("audio", {}).get("format", "mp3")
+            for choice in http_response.get("choices", []):
+                audio_data = choice.get("message", {}).get("audio", {})
+                if audio_data:
+                    audio_content = audio_data.get("data")
+                    aggregate_multimodal_audio_content_output(
+                        audio_content, audio_format, scope3_row
+                    )
+
         messages = kwargs.get("messages", [])
         for message in messages:
             aggregate_multimodal(message, scope3_row, logger)
@@ -53,15 +67,24 @@ def _openai_chat_wrapper(
         setattr(response, "scope3ai", scope3ai_ctx)
         return response
     else:
-        model_used = response.model
         scope3_row = ImpactRow(
             model_id=model_requested,
-            model_used_id=model_used,
+            model_used_id=response.model,
             input_tokens=response.usage.prompt_tokens,
             output_tokens=response.usage.completion_tokens,
             request_duration_ms=request_latency * 1000,
             managed_service_id=PROVIDER,
         )
+        if "audio" in modalities:
+            audio_format = kwargs.get("audio", {}).get("format", "mp3")
+            for choice in response.choices:
+                audio_data = getattr(choice.message, "audio")
+                if audio_data:
+                    audio_content = audio_data.data
+                    aggregate_multimodal_audio_content_output(
+                        audio_content, audio_format, scope3_row
+                    )
+
         messages = kwargs.get("messages", [])
         for message in messages:
             aggregate_multimodal(message, scope3_row, logger)

scope3ai/tracers/utils/multimodal.py

@@ -49,6 +49,20 @@ def aggregate_multimodal_audio(content: dict, row: ImpactRow) -> None:
             row.input_audio_seconds += duration
 
 
+def aggregate_multimodal_audio_content_output(
+    content: str, audio_format: str, row: ImpactRow
+) -> None:
+    assert audio_format in MUTAGEN_MAPPING
+
+    audio_data = base64.b64decode(content)
+    duration = _get_audio_duration(audio_format, audio_data)
+    if duration:
+        if row.output_audio_seconds is None:
+            row.output_audio_seconds = duration
+        else:
+            row.output_audio_seconds += duration
+
+
 def aggregate_multimodal_content(
     content: dict, row: ImpactRow, logger: logging.Logger
 ) -> None: