feat(openai): add support for image generation tracking

[coletebou]

Summary

Adds automatic tracing for OpenAI image generation API calls:

• images.generate()
• images.edit()
• images.create_variation()

Changes

langfuse/openai.py

• Added 6 new OpenAiDefinition entries for image methods (sync and async variants)
• Extended _get_langfuse_data_from_kwargs() to extract image generation parameters (size, quality, style, response_format, n)
• Extended _get_langfuse_data_from_default_response() to parse image responses (URLs and base64 data wrapped in LangfuseMedia)
• Added handling in _wrap() and _wrap_async() to track image count as usage metrics

tests/test_openai.py

• Added tests for sync and async images.generate()
• Added skipped test stubs for images.edit() and images.create_variation() (require DALL-E 2 + image files)

Tracked Data

Field	Value
name	OpenAI-image (default)
input	Prompt text (generate/edit) or None (create_variation)
model	dall-e-2, dall-e-3, gpt-image-1, etc.
model_parameters	size, quality, style, response_format, n
output	URL or base64 image wrapped in LangfuseMedia
usage	Image count ({output: N, total: N, unit: IMAGES})

Notes

• create_variation() does not accept a prompt parameter, so input will be None for those calls. This is expected behavior.
• Base64 responses are wrapped in LangfuseMedia for proper handling in the Langfuse UI.
• Tests for edit() and create_variation() are skipped as they require DALL-E 2 API access and PNG image files with alpha channels.

---

Important

Adds support for tracking OpenAI image generation API calls in Langfuse, including handling image-specific parameters and responses.

• Behavior:
  • Adds support for tracking OpenAI image generation API calls in langfuse/openai.py, including images.generate(), images.edit(), and images.create_variation().
  • Extends _get_langfuse_data_from_kwargs() to handle image parameters like size, quality, style, response_format, and n.
  • Extends _get_langfuse_data_from_default_response() to parse image responses, wrapping base64 data in LangfuseMedia.
  • Updates _wrap() and _wrap_async() to track image count as usage metrics.
• Tests:
  • Adds tests in tests/test_openai.py for sync and async images.generate().
  • Adds skipped test stubs for images.edit() and images.create_variation() due to DALL-E 2 requirements.
• Misc:
  • create_variation() does not use a prompt, so input is None for these calls.

^{This description was created by for 7154a72. You can customize this summary. It will automatically update as commits are pushed.}

Disclaimer: Experimental PR review

Greptile Overview

Greptile Summary

This PR adds comprehensive tracking support for OpenAI image generation APIs (images.generate(), images.edit(), and images.create_variation()).

Key changes:

• Added 6 new OpenAiDefinition entries for image methods (sync/async variants)
• Extended parameter extraction to capture image-specific parameters (size, quality, style, response_format, n)
• Added response parsing to handle both URL and base64 image responses, wrapping base64 data in LangfuseMedia
• Implemented custom usage tracking with image count metrics
• Added comprehensive tests for generate() with appropriate skipped stubs for edit() and create_variation()

Implementation quality:
The implementation follows existing patterns in the codebase and handles edge cases well. The code properly differentiates between methods that accept prompts (generate, edit) vs those that don't (create_variation). Response handling correctly unwraps single-image responses while preserving lists for multiple images. Usage tracking accurately counts images regardless of response format.

Confidence Score: 4/5

• This PR is safe to merge with minimal risk
• The implementation is well-structured and follows existing patterns in the codebase. The code properly handles different image generation methods, parameter extraction, and response parsing. Tests are comprehensive for the testable portions. No security issues or breaking changes detected.
• No files require special attention

Important Files Changed

Filename	Overview
langfuse/openai.py	Added image generation tracking support with proper parameter extraction and usage metrics - implementation is solid with good separation of concerns
tests/test_openai.py	Added comprehensive tests for image generation with proper assertions - sync and async tests included, edit/variation tests appropriately skipped

Sequence Diagram

sequenceDiagram
    participant App as Application
    participant LF as Langfuse Wrapper
    participant OAI as OpenAI API
    participant DB as Langfuse Backend

    App->>LF: images.generate(prompt, model, size, ...)
    activate LF
    LF->>LF: _get_langfuse_data_from_kwargs()
    Note over LF: Extract prompt, model, params<br/>(size, quality, style, n)
    LF->>DB: start_observation(type="generation")
    LF->>OAI: images.generate()
    activate OAI
    OAI-->>LF: ImagesResponse (url or b64_json)
    deactivate OAI
    LF->>LF: _get_langfuse_data_from_default_response()
    Note over LF: Parse response, wrap base64 in<br/>LangfuseMedia, extract revised_prompt
    LF->>LF: Calculate usage (image_count)
    LF->>DB: update(output, usage_details, model)
    LF->>DB: end()
    LF-->>App: Return OpenAI response
    deactivate LF

Loading

_{(5/5) You can turn off certain types of comments like style here!}

[greptile-apps]

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}