feat(litellm): Use default tracer and Image generation - TTS - STT

README.md

@@ -35,7 +35,7 @@ uv add scope3ai
 | Cohere      | ✅              |    |     |                  |           |                  |
 | OpenAI      | ✅              | ✅ | ✅  | ✅               | ✅        | Images/Audio     |
 | Huggingface | ✅              | ✅ | ✅  | ✅               | ✅        |                  |
-| LiteLLM     | ✅              |    |     |                  |           | Images/Audio     |
+| LiteLLM     | ✅              |  ✅  |✅     |       ✅           |           | Images/Audio     |
 | MistralAi   | ✅              |    |     |                  |           | Images           |
 
 Roadmap:

scope3ai/api/tracer.py

@@ -73,7 +73,7 @@ def _link_trace(self, trace) -> None:
             self.traces.append(trace)
 
     def _unlink_trace(self, trace) -> None:
-        if not self.keep_traces:
+        if self.keep_traces:
             return
         if trace in self.traces:
             self.traces.remove(trace)

scope3ai/api/typesgen.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  aiapi.yaml
-#   timestamp: 2025-01-16T00:04:03+00:00
+#   timestamp: 2025-01-21T22:57:07+00:00
 
 from __future__ import annotations
 
@@ -243,7 +243,7 @@ class Image(RootModel[str]):
     ]
 
 
-class Task(Enum):
+class Task(str, Enum):
     """
     Common types of AI/ML models and their primary functions:
     - Text-based models for natural language processing
@@ -287,7 +287,7 @@ class Task(Enum):
     clustering = "clustering"
 
 
-class Family(Enum):
+class Family(str, Enum):
     """
     Core AI model families from various organizations:
     - Commercial models from major AI companies
@@ -325,7 +325,7 @@ class Family(Enum):
     gpt_j = "gpt-j"
 
 
-class DataType(Enum):
+class DataType(str, Enum):
     fp8 = "fp8"
     fp8_e4m3 = "fp8-e4m3"
     fp8_e5m2 = "fp8-e5m2"

scope3ai/tracers/litellm/chat.py

@@ -71,10 +71,11 @@ def litellm_chat_wrapper_non_stream(
     kwargs: Any,
 ) -> ChatCompletion:
     timer_start = time.perf_counter()
-    with Scope3AI.get_instance().trace(keep_traces=True) as trace:
+    keep_traces = not kwargs.pop("use_always_litellm_tracer", False)
+    with Scope3AI.get_instance().trace(keep_traces=keep_traces) as tracer:
         response = wrapped(*args, **kwargs)
-        if trace.traces:
-            setattr(response, "scope3ai", trace.traces[0])
+        if tracer.traces:
+            setattr(response, "scope3ai", tracer.traces[0])
             return response
     request_latency = time.perf_counter() - timer_start
     model = response.model
@@ -113,7 +114,12 @@ async def litellm_async_chat_wrapper_base(
     kwargs: Any,
 ) -> ChatCompletion:
     timer_start = time.perf_counter()
-    response = await wrapped(*args, **kwargs)
+    keep_traces = not kwargs.pop("use_always_litellm_tracer", False)
+    with Scope3AI.get_instance().trace(keep_traces=keep_traces) as tracer:
+        response = await wrapped(*args, **kwargs)
+        if tracer.traces:
+            setattr(response, "scope3ai", tracer.traces[0])
+            return response
     request_latency = time.perf_counter() - timer_start
     model = response.model
     if model is None:

scope3ai/tracers/litellm/instrument.py

@@ -1,13 +1,25 @@
 import litellm
-from wrapt import wrap_function_wrapper  # type: ignore[import-untyped]
 
+from scope3ai.base_tracer import BaseTracer
 from scope3ai.tracers.litellm.chat import (
     litellm_chat_wrapper,
     litellm_async_chat_wrapper,
 )
+from scope3ai.tracers.litellm.speech_to_text import (
+    litellm_speech_to_text_wrapper,
+    litellm_speech_to_text_wrapper_async,
+)
+from scope3ai.tracers.litellm.text_to_image import (
+    litellm_image_generation_wrapper_async,
+    litellm_image_generation_wrapper,
+)
+from scope3ai.tracers.litellm.text_to_speech import (
+    litellm_speech_generation_wrapper,
+    litellm_speech_generation_wrapper_async,
+)
 
 
-class LiteLLMInstrumentor:
+class LiteLLMInstrumentor(BaseTracer):
     def __init__(self) -> None:
         self.wrapped_methods = [
             {
@@ -20,10 +32,34 @@ def __init__(self) -> None:
                 "name": "acompletion",
                 "wrapper": litellm_async_chat_wrapper,
             },
+            {
+                "module": litellm,
+                "name": "image_generation",
+                "wrapper": litellm_image_generation_wrapper,
+            },
+            {
+                "module": litellm,
+                "name": "aimage_generation",
+                "wrapper": litellm_image_generation_wrapper_async,
+            },
+            {
+                "module": litellm,
+                "name": "speech",
+                "wrapper": litellm_speech_generation_wrapper,
+            },
+            {
+                "module": litellm,
+                "name": "aspeech",
+                "wrapper": litellm_speech_generation_wrapper_async,
+            },
+            {
+                "module": litellm,
+                "name": "transcription",
+                "wrapper": litellm_speech_to_text_wrapper,
+            },
+            {
+                "module": litellm,
+                "name": "atranscription",
+                "wrapper": litellm_speech_to_text_wrapper_async,
+            },
         ]
-
-    def instrument(self) -> None:
-        for wrapper in self.wrapped_methods:
-            wrap_function_wrapper(
-                wrapper["module"], wrapper["name"], wrapper["wrapper"]
-            )

scope3ai/tracers/litellm/speech_to_text.py

@@ -0,0 +1,83 @@
+import time
+from typing import Any, Callable, Optional
+
+import tiktoken
+from litellm import Completions
+from litellm.types.utils import TranscriptionResponse as _TranscriptionResponse
+
+from scope3ai import Scope3AI
+from scope3ai.api.types import ImpactRow
+from scope3ai.api.types import Scope3AIContext
+from scope3ai.api.typesgen import Task
+from scope3ai.constants import PROVIDERS
+from scope3ai.tracers.utils.audio import _get_file_audio_duration
+
+PROVIDER = PROVIDERS.LITELLM.value
+
+
+class TranscriptionResponse(_TranscriptionResponse):
+    scope3ai: Optional[Scope3AIContext] = None
+
+
+def litellm_speech_to_text_get_impact_row(
+    timer_start: Any,
+    response: TranscriptionResponse,
+    args,
+    kwargs,
+) -> (TranscriptionResponse, ImpactRow):
+    request_latency = time.perf_counter() - timer_start
+    file = args[0] if len(args) > 0 else kwargs.get("file")
+    model = args[1] if len(args) > 1 else kwargs.get("model")
+    request_latency = getattr(response, "_response_ms", request_latency)
+    encoder = tiktoken.get_encoding("cl100k_base")
+    options = {}
+    duration = _get_file_audio_duration(file)
+    if duration is not None:
+        options["input_audio_seconds"] = int(duration)
+    output_tokens = len(encoder.encode(response.text))
+    scope3_row = ImpactRow(
+        model_id=model,
+        output_tokens=output_tokens,
+        request_duration_ms=float(request_latency) * 1000,
+        managed_service_id=PROVIDER,
+        task=Task.speech_to_text,
+        **options,
+    )
+    return scope3_row
+
+
+def litellm_speech_to_text_wrapper(
+    wrapped: Callable, instance: Completions, args: Any, kwargs: Any
+):
+    timer_start = time.perf_counter()
+    keep_traces = not kwargs.pop("use_always_litellm_tracer", False)
+    with Scope3AI.get_instance().trace(keep_traces=keep_traces) as tracer:
+        response = wrapped(*args, **kwargs)
+        if tracer.traces:
+            setattr(response, "scope3ai", tracer.traces[0])
+            return response
+
+    impact_row = litellm_speech_to_text_get_impact_row(
+        timer_start, response, args, kwargs
+    )
+    scope3_ctx = Scope3AI.get_instance().submit_impact(impact_row)
+    response.scope3ai = scope3_ctx
+    return response
+
+
+async def litellm_speech_to_text_wrapper_async(
+    wrapped: Callable, instance: Completions, args: Any, kwargs: Any
+):
+    timer_start = time.perf_counter()
+    keep_traces = not kwargs.pop("use_always_litellm_tracer", False)
+    with Scope3AI.get_instance().trace(keep_traces=keep_traces) as tracer:
+        response = await wrapped(*args, **kwargs)
+        if tracer.traces:
+            setattr(response, "scope3ai", tracer.traces[0])
+            return response
+    impact_row = litellm_speech_to_text_get_impact_row(
+        timer_start, response, args, kwargs
+    )
+    scope3_ctx = await Scope3AI.get_instance().asubmit_impact(impact_row)
+    response.scope3ai = scope3_ctx
+    return response

scope3ai/tracers/litellm/text_to_image.py

@@ -0,0 +1,85 @@
+import time
+from typing import Any, Callable, Optional
+
+import tiktoken
+from litellm import Completions
+from litellm.utils import ImageResponse as _ImageResponse
+
+from scope3ai import Scope3AI
+from scope3ai.api.types import ImpactRow, Scope3AIContext
+from scope3ai.api.typesgen import Image as RootImage, Task
+from scope3ai.constants import PROVIDERS
+
+PROVIDER = PROVIDERS.LITELLM.value
+DEFAULT_MODEL = "dall-e-2"
+DEFAULT_SIZE = "1024x1024"
+DEFAULT_N = 1
+
+
+class ImageResponse(_ImageResponse):
+    scope3ai: Optional[Scope3AIContext] = None
+
+
+def litellm_image_generation_get_impact_row(
+    timer_start: Any,
+    response: ImageResponse,
+    args,
+    kwargs,
+) -> ImpactRow:
+    request_latency = time.perf_counter() - timer_start
+    prompt = args[0] if len(args) > 0 else kwargs.get("prompt")
+    model = args[1] if len(args) > 1 else kwargs.get("model")
+    request_latency = getattr(response, "_response_ms", request_latency)
+
+    encoder = tiktoken.get_encoding("cl100k_base")
+    input_tokens = len(encoder.encode(prompt))
+    n = kwargs.get("n", DEFAULT_N)
+    size = RootImage(root=kwargs.get("size", DEFAULT_SIZE))
+
+    scope3_row = ImpactRow(
+        model_id=model or DEFAULT_MODEL,
+        task=Task.text_to_image,
+        request_duration_ms=float(request_latency) * 1000,
+        managed_service_id=PROVIDER,
+        output_images=[size] * n,
+        input_tokens=input_tokens,
+    )
+    return scope3_row
+
+
+def litellm_image_generation_wrapper(
+    wrapped: Callable, instance: Completions, args: Any, kwargs: Any
+):
+    timer_start = time.perf_counter()
+    keep_traces = not kwargs.pop("use_always_litellm_tracer", False)
+    with Scope3AI.get_instance().trace(keep_traces=keep_traces) as tracer:
+        response = wrapped(*args, **kwargs)
+        if tracer.traces:
+            setattr(response, "scope3ai", tracer.traces[0])
+            return response
+
+    impact_row = litellm_image_generation_get_impact_row(
+        timer_start, response, args, kwargs
+    )
+    scope3_ctx = Scope3AI.get_instance().submit_impact(impact_row)
+    response.scope3ai = scope3_ctx
+    return response
+
+
+async def litellm_image_generation_wrapper_async(
+    wrapped: Callable, instance: Completions, args: Any, kwargs: Any
+):
+    timer_start = time.perf_counter()
+    keep_traces = not kwargs.pop("use_always_litellm_tracer", False)
+    with Scope3AI.get_instance().trace(keep_traces=keep_traces) as tracer:
+        response = await wrapped(*args, **kwargs)
+        if tracer.traces:
+            setattr(response, "scope3ai", tracer.traces[0])
+            return response
+
+    impact_row = litellm_image_generation_get_impact_row(
+        timer_start, response, args, kwargs
+    )
+    scope3_ctx = await Scope3AI.get_instance().asubmit_impact(impact_row)
+    response.scope3ai = scope3_ctx
+    return response