feat(openai): support for image generation

README.md

@@ -33,7 +33,7 @@ uv add scope3ai
 |-------------|-----------------|----|-----|------------------|-----------|
 | Anthropic   | ✅              |    |     |                  |           |
 | Cohere      | ✅              |    |     |                  |           |
-| OpenAI      | ✅              | ✅ | ✅  |                  |           |
+| OpenAI      | ✅              | ✅ | ✅  | ✅               |           |
 | Huggingface | ✅              | ✅ | ✅  | ✅               | ✅        |
 | LiteLLM     | ✅              |    |     |                  |           |
 | MistralAi   | ✅              |    |     |                  |           |

examples/openai-async-image.py

@@ -0,0 +1,27 @@
+import asyncio
+from scope3ai import Scope3AI
+from openai import AsyncOpenAI
+
+
+async def main():
+    client = AsyncOpenAI()
+    scope3 = Scope3AI.init()
+
+    with scope3.trace() as tracer:
+        response = await client.images.generate(
+            model="dall-e-2",
+            prompt="A beautiful landscape",
+            n=1,
+            size="512x512",
+        )
+        print(response.data[0].url)
+
+        impact = tracer.impact()
+        print(impact)
+        print(f"Total Energy Wh: {impact.total_energy_wh}")
+        print(f"Total GCO2e: {impact.total_gco2e}")
+        print(f"Total MLH2O: {impact.total_mlh2o}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/openai-image.py

@@ -0,0 +1,26 @@
+from scope3ai import Scope3AI
+from openai import OpenAI
+
+
+def main():
+    client = OpenAI()
+    scope3 = Scope3AI.init()
+
+    with scope3.trace() as tracer:
+        response = client.images.generate(
+            model="dall-e-2",
+            prompt="A beautiful landscape",
+            n=1,
+            size="512x512",
+        )
+        print(response.data[0].url)
+
+        impact = tracer.impact()
+        print(impact)
+        print(f"Total Energy Wh: {impact.total_energy_wh}")
+        print(f"Total GCO2e: {impact.total_gco2e}")
+        print(f"Total MLH2O: {impact.total_mlh2o}")
+
+
+if __name__ == "__main__":
+    main()

scope3ai/tracers/openai/instrument.py

@@ -5,6 +5,10 @@
     openai_text_to_speech_wrapper,
     openai_async_text_to_speech_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,
@@ -44,6 +48,36 @@ def __init__(self) -> None:
                 "name": "AsyncTranscriptions.create",
                 "wrapper": openai_async_speech_to_text_wrapper,
             },
+            {
+                "module": "openai.resources.images",
+                "name": "Images.create_variation",
+                "wrapper": openai_image_wrapper,
+            },
+            {
+                "module": "openai.resources.images",
+                "name": "Images.edit",
+                "wrapper": openai_image_wrapper,
+            },
+            {
+                "module": "openai.resources.images",
+                "name": "Images.generate",
+                "wrapper": openai_image_wrapper,
+            },
+            {
+                "module": "openai.resources.images",
+                "name": "AsyncImages.create_variation",
+                "wrapper": openai_async_image_wrapper,
+            },
+            {
+                "module": "openai.resources.images",
+                "name": "AsyncImages.edit",
+                "wrapper": openai_async_image_wrapper,
+            },
+            {
+                "module": "openai.resources.images",
+                "name": "AsyncImages.generate",
+                "wrapper": openai_async_image_wrapper,
+            },
         ]
 
     def instrument(self) -> None:

scope3ai/tracers/openai/text_to_image.py

@@ -0,0 +1,56 @@
+import time
+from typing import Any, Callable, Optional
+
+from openai.resources.images import AsyncImages, Images
+from openai.types.images_response import ImagesResponse as _ImageResponse
+
+from scope3ai.api.types import ImpactRow, Model, Scope3AIContext, Task
+from scope3ai.lib import Scope3AI
+
+PROVIDER = "openai"
+DEFAULT_MODEL = "dall-e-2"
+DEFAULT_SIZE = "1024x1024"
+DEFAULT_N = 1
+
+
+class ImageResponse(_ImageResponse):
+    scope3ai: Optional[Scope3AIContext] = None
+
+
+def _openai_image_wrapper(
+    response: _ImageResponse, request_latency: float, **kwargs: Any
+) -> ImageResponse:
+    model = kwargs.get("model", DEFAULT_MODEL)
+    size = kwargs.get("size", DEFAULT_SIZE)
+    n = kwargs.get("n", DEFAULT_N)
+
+    scope3_row = ImpactRow(
+        model=Model(id=model),
+        task=Task.text_to_image,
+        output_images=[size] * n,
+        request_duration_ms=request_latency * 1000,
+        managed_service_id=PROVIDER,
+    )
+
+    scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row)
+    result = ImageResponse.model_construct(**response.model_dump())
+    result.scope3ai = scope3ai_ctx
+    return result
+
+
+def openai_image_wrapper(
+    wrapped: Callable, instance: Images, args: Any, kwargs: Any
+) -> ImageResponse:
+    timer_start = time.perf_counter()
+    response = wrapped(*args, **kwargs)
+    request_latency = time.perf_counter() - timer_start
+    return _openai_image_wrapper(response, request_latency, **kwargs)
+
+
+async def openai_async_image_wrapper(
+    wrapped: Callable, instance: AsyncImages, args: Any, kwargs: Any
+) -> ImageResponse:
+    timer_start = time.perf_counter()
+    response = await wrapped(*args, **kwargs)
+    request_latency = time.perf_counter() - timer_start
+    return _openai_image_wrapper(response, request_latency, **kwargs)