feat(litellm): multimodal

scope3ai/api/tracer.py

@@ -7,11 +7,13 @@ class Tracer:
     def __init__(
         self,
         name: str = None,
+        keep_traces: bool = False,
     ) -> None:
         from scope3ai.lib import Scope3AI
 
         self.scope3ai = Scope3AI.get_instance()
         self.name = name
+        self.keep_traces = keep_traces
         self.children: List[Tracer] = []
         self.rows: List[ModeledRow] = []
         self.traces = []  # type: List[Scope3AIContext]
@@ -71,5 +73,7 @@ def _link_trace(self, trace) -> None:
             self.traces.append(trace)
 
     def _unlink_trace(self, trace) -> None:
+        if not self.keep_traces:
+            return
         if trace in self.traces:
             self.traces.remove(trace)

scope3ai/base_tracer.py

@@ -0,0 +1,13 @@
+from wrapt import wrap_function_wrapper
+
+
+class BaseTracer:
+    wrapper_methods = []
+
+    def instrument(self) -> None:
+        for wrapper in self.wrapped_methods:
+            wrap_function_wrapper(
+                wrapper["module"],
+                wrapper["name"],
+                wrapper["wrapper"],
+            )

scope3ai/lib.py

@@ -108,6 +108,7 @@ class Scope3AI:
     _tracer: ContextVar[List[Tracer]] = ContextVar("tracer", default=[])
     _worker: Optional[BackgroundWorker] = None
     _providers: List[str] = []
+    _keep_tracers: bool = False
 
     def __new__(cls, *args, **kwargs):
         if cls._instance is None:
@@ -276,8 +277,8 @@ def current_tracer(self):
         return tracers[-1] if tracers else None
 
     @contextmanager
-    def trace(self):
-        tracer = Tracer()
+    def trace(self, keep_traces=False):
+        tracer = Tracer(keep_traces=keep_traces)
         try:
             self._push_tracer(tracer)
             yield tracer

scope3ai/tracers/anthropic/instrument.py

@@ -1,13 +1,13 @@
-from wrapt import wrap_function_wrapper  # type: ignore[import-untyped]
 from .chat import (
     anthropic_chat_wrapper,
     anthropic_async_chat_wrapper,
     anthropic_stream_chat_wrapper,
     anthropic_async_stream_chat_wrapper,
 )
+from ...base_tracer import BaseTracer
 
 
-class AnthropicInstrumentor:
+class AnthropicInstrumentor(BaseTracer):
     def __init__(self) -> None:
         self.wrapped_methods = [
             {
@@ -31,9 +31,3 @@ def __init__(self) -> None:
                 "wrapper": anthropic_async_stream_chat_wrapper,
             },
         ]
-
-    def instrument(self) -> None:
-        for wrapper in self.wrapped_methods:
-            wrap_function_wrapper(
-                wrapper["module"], wrapper["name"], wrapper["wrapper"]
-            )

scope3ai/tracers/cohere/instrument.py

@@ -1,4 +1,6 @@
 from wrapt import wrap_function_wrapper  # type: ignore[import-untyped]
+
+from scope3ai.base_tracer import BaseTracer
 from .chat import (
     cohere_chat_wrapper,
     cohere_async_chat_wrapper,
@@ -13,7 +15,7 @@
 )
 
 
-class CohereInstrumentor:
+class CohereInstrumentor(BaseTracer):
     def __init__(self) -> None:
         self.wrapped_methods = [
             {

scope3ai/tracers/huggingface/instrument.py

@@ -1,5 +1,6 @@
 from wrapt import wrap_function_wrapper  # type: ignore[import-untyped]
 
+from scope3ai.base_tracer import BaseTracer
 from scope3ai.tracers.huggingface.chat import (
     huggingface_chat_wrapper,
     huggingface_async_chat_wrapper,
@@ -38,7 +39,7 @@
 )
 
 
-class HuggingfaceInstrumentor:
+class HuggingfaceInstrumentor(BaseTracer):
     def __init__(self) -> None:
         self.wrapped_methods = [
             {

scope3ai/tracers/litellm/chat.py

@@ -1,3 +1,4 @@
+import logging
 import time
 from typing import Any, Callable, Optional, Union
 
@@ -8,9 +9,12 @@
 from scope3ai import Scope3AI
 from scope3ai.api.types import Scope3AIContext, ImpactRow
 from scope3ai.constants import PROVIDERS
+from scope3ai.tracers.utils.multimodal import aggregate_multimodal
 
 PROVIDER = PROVIDERS.LITELLM.value
 
+logger = logging.getLogger("scope3ai.tracers.litellm.chat")
+
 
 class ChatCompletion(ModelResponse):
     scope3ai: Optional[Scope3AIContext] = None
@@ -67,7 +71,11 @@ def litellm_chat_wrapper_non_stream(
     kwargs: Any,
 ) -> ChatCompletion:
     timer_start = time.perf_counter()
-    response = wrapped(*args, **kwargs)
+    with Scope3AI.get_instance().trace(keep_traces=True) as trace:
+        response = wrapped(*args, **kwargs)
+        if trace.traces:
+            setattr(response, "scope3ai", trace.traces[0])
+            return response
     request_latency = time.perf_counter() - timer_start
     model = response.model
     if model is None:
@@ -79,6 +87,9 @@ def litellm_chat_wrapper_non_stream(
         request_duration_ms=float(request_latency) * 1000,
         managed_service_id=PROVIDER,
     )
+    messages = args[1] if len(args) > 1 else kwargs.get("messages")
+    for message in messages:
+        aggregate_multimodal(message, scope3_row, logger)
     scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row)
     if scope3ai_ctx is not None:
         return ChatCompletion(**response.model_dump(), scope3ai=scope3ai_ctx)

scope3ai/tracers/mistralai/instrument.py

@@ -1,4 +1,5 @@
 from wrapt import wrap_function_wrapper  # type: ignore[import-untyped]
+from scope3ai.base_tracer import BaseTracer
 
 from scope3ai.tracers.mistralai.chat import (
     mistralai_v1_chat_wrapper,
@@ -8,7 +9,7 @@
 )
 
 
-class MistralAIInstrumentor:
+class MistralAIInstrumentor(BaseTracer):
     def __init__(self) -> None:
         self.wrapped_methods = [
             {

scope3ai/tracers/openai/chat.py

@@ -1,115 +1,74 @@
-import base64
 import logging
 import time
-from io import BytesIO
 from typing import Any, Callable, Optional, Union
 
 from openai import AsyncStream, Stream
+from openai._legacy_response import LegacyAPIResponse as _LegacyAPIResponse
 from openai.resources.chat import AsyncCompletions, Completions
 from openai.types.chat import ChatCompletion as _ChatCompletion
 from openai.types.chat import ChatCompletionChunk as _ChatCompletionChunk
 
 from scope3ai.api.types import ImpactRow, Scope3AIContext
-from scope3ai.api.typesgen import Image as RootImage
 from scope3ai.constants import PROVIDERS
 from scope3ai.lib import Scope3AI
-
-from .utils import MUTAGEN_MAPPING, _get_audio_duration
+from scope3ai.tracers.utils.multimodal import aggregate_multimodal
 
 PROVIDER = PROVIDERS.OPENAI.value
+
 logger = logging.getLogger("scope3ai.tracers.openai.chat")
 
 
-class ChatCompletion(_ChatCompletion):
+class LegacyApiResponse(_LegacyAPIResponse):
     scope3ai: Optional[Scope3AIContext] = None
 
 
-class ChatCompletionChunk(_ChatCompletionChunk):
+class ChatCompletion(_ChatCompletion):
     scope3ai: Optional[Scope3AIContext] = None
 
 
-def _openai_aggregate_multimodal_image(content: dict, row: ImpactRow) -> None:
-    from PIL import Image
-
-    url = content["image_url"]["url"]
-    if url.startswith("data:"):
-        # extract content type, and data part
-        # example: data:image/jpeg;base64,....
-        content_type, data = url.split(",", 1)
-        image_data = BytesIO(base64.b64decode(data))
-        image = Image.open(image_data)
-        width, height = image.size
-        size = RootImage(root=f"{width}x{height}")
-
-        if row.input_images is None:
-            row.input_images = [size]
-        else:
-            row.input_images.append(size)
-
-    else:
-        # TODO: not supported yet.
-        # Should we actually download the file here just to have the size ??
-        pass
-
-
-def _openai_aggregate_multimodal_audio(content: dict, row: ImpactRow) -> None:
-    input_audio = content["input_audio"]
-    format = input_audio["format"]
-    b64data = input_audio["data"]
-    assert format in MUTAGEN_MAPPING
-
-    # decode the base64 data
-    audio_data = base64.b64decode(b64data)
-    # TODO: accept audio duration as float in AiApi
-    duration = int(_get_audio_duration(format, audio_data))
-
-    if row.input_audio_seconds is None:
-        row.input_audio_seconds = duration
-    else:
-        row.input_audio_seconds += duration
-
-
-def _openai_aggregate_multimodal_content(content: dict, row: ImpactRow) -> None:
-    try:
-        content_type = content.get("type")
-        if content_type == "image_url":
-            _openai_aggregate_multimodal_image(content, row)
-        elif content_type == "input_audio":
-            _openai_aggregate_multimodal_audio(content, row)
-    except Exception as e:
-        logger.error(f"Error processing multimodal content: {e}")
-
-
-def _openai_aggregate_multimodal(message: dict, row: ImpactRow) -> None:
-    # if the message content is not a tuple/list, it's just text.
-    # so there is nothing multimodal in it, we can just forget about it.
-    content = message.get("content", [])
-    if isinstance(content, (tuple, list)):
-        for item in content:
-            _openai_aggregate_multimodal_content(item, row)
+class ChatCompletionChunk(_ChatCompletionChunk):
+    scope3ai: Optional[Scope3AIContext] = None
 
 
 def _openai_chat_wrapper(
     response: Any, request_latency: float, kwargs: dict
 ) -> ChatCompletion:
     model_requested = kwargs["model"]
-    model_used = response.model
-
-    scope3_row = ImpactRow(
-        model_id=model_requested,
-        model_used_id=model_used,
-        input_tokens=response.usage.prompt_tokens,
-        output_tokens=response.usage.completion_tokens,
-        request_duration_ms=request_latency * 1000,
-        managed_service_id=PROVIDER,
-    )
+    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"),
+            request_duration_ms=request_latency * 1000,
+            managed_service_id=PROVIDER,
+        )
+        messages = kwargs.get("messages", [])
+        for message in messages:
+            aggregate_multimodal(message, scope3_row, logger)
+        Scope3AI.get_instance().submit_impact(scope3_row)
+        scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row)
+        setattr(response, "scope3ai", scope3ai_ctx)
+        return response
+    else:
+        model_used = response.model
+        scope3_row = ImpactRow(
+            model_id=model_requested,
+            model_used_id=model_used,
+            input_tokens=response.usage.prompt_tokens,
+            output_tokens=response.usage.completion_tokens,
+            request_duration_ms=request_latency * 1000,
+            managed_service_id=PROVIDER,
+        )
+        messages = kwargs.get("messages", [])
+        for message in messages:
+            aggregate_multimodal(message, scope3_row, logger)
+        scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row)
+        return ChatCompletion(**response.model_dump(), scope3ai=scope3ai_ctx)
 
     # analyse multimodal part
-    messages = kwargs.get("messages", [])
-    for message in messages:
-        _openai_aggregate_multimodal(message, scope3_row)
-    scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row)
-    return ChatCompletion(**response.model_dump(), scope3ai=scope3ai_ctx)
 
 
 def openai_chat_wrapper_non_stream(

scope3ai/tracers/openai/instrument.py

@@ -1,25 +1,24 @@
-from wrapt import wrap_function_wrapper
-
+from scope3ai.base_tracer import BaseTracer
 from .chat import openai_chat_wrapper, openai_async_chat_wrapper
-from .text_to_speech import (
-    openai_text_to_speech_wrapper,
-    openai_async_text_to_speech_wrapper,
+from .speech_to_text import (
+    openai_async_speech_to_text_wrapper,
+    openai_speech_to_text_wrapper,
 )
 from .text_to_image import (
     openai_image_wrapper,
     openai_async_image_wrapper,
 )
-from .speech_to_text import (
-    openai_async_speech_to_text_wrapper,
-    openai_speech_to_text_wrapper,
+from .text_to_speech import (
+    openai_text_to_speech_wrapper,
+    openai_async_text_to_speech_wrapper,
 )
 from .translation import (
     openai_translation_wrapper,
     openai_async_translation_wrapper,
 )
 
 
-class OpenAIInstrumentor:
+class OpenAIInstrumentor(BaseTracer):
     def __init__(self) -> None:
         self.wrapped_methods = [
             {
@@ -93,9 +92,3 @@ def __init__(self) -> None:
                 "wrapper": openai_async_image_wrapper,
             },
         ]
-
-    def instrument(self) -> None:
-        for wrapper in self.wrapped_methods:
-            wrap_function_wrapper(
-                wrapper["module"], wrapper["name"], wrapper["wrapper"]
-            )

scope3ai/tracers/openai/speech_to_text.py

@@ -12,8 +12,7 @@
 from scope3ai.api.types import ImpactRow, Scope3AIContext, Task
 from scope3ai.constants import PROVIDERS
 from scope3ai.lib import Scope3AI
-
-from .utils import _get_file_audio_duration
+from scope3ai.tracers.utils.audio import _get_file_audio_duration
 
 PROVIDER = PROVIDERS.OPENAI.value
 

scope3ai/tracers/openai/text_to_speech.py

@@ -8,8 +8,7 @@
 from scope3ai.api.types import ImpactRow, Scope3AIContext, Task
 from scope3ai.constants import PROVIDERS
 from scope3ai.lib import Scope3AI
-
-from .utils import _get_audio_duration
+from scope3ai.tracers.utils.audio import _get_audio_duration
 
 PROVIDER = PROVIDERS.OPENAI.value
 

scope3ai/tracers/openai/translation.py

@@ -12,8 +12,7 @@
 from scope3ai.api.types import ImpactRow, Scope3AIContext, Task
 from scope3ai.constants import PROVIDERS
 from scope3ai.lib import Scope3AI
-
-from .utils import _get_file_audio_duration
+from scope3ai.tracers.utils.audio import _get_file_audio_duration
 
 PROVIDER = PROVIDERS.OPENAI.value