feat: transform raw image responses to json

[unsync]

Author Description

⚠️ This is a base for discussion/comments.

Summary By MatterAI

🔄 What Changed

This PR adds functionality to transform image content type responses from Segmind to JSON format. It introduces a new handler for converting image responses to base64-encoded JSON, specifically targeting the imageGenerate route. The implementation includes error handling to fall back to the original response if transformation fails.

🔍 Impact of the Change

This change improves the handling of image responses by standardizing them to JSON format, making it easier for clients to process the responses. It's designed as a non-breaking change that enhances the existing functionality without disrupting current behavior.

📁 Total Files Changed

5 files changed with 67 additions and 1 deletion:

• Added new file: src/handlers/imageToJsonHandler.ts
• Modified: src/handlers/responseHandlers.ts
• Modified: src/handlers/retryHandler.ts
• Modified: src/providers/segmind/imageGenerate.ts
• Modified: src/providers/segmind/index.ts

🧪 Test Added

No tests have been added in this PR.

🔒 Security Vulnerabilities

No security vulnerabilities detected.

Type of Change

• New feature (non-breaking change which adds functionality)
• Bug fix (non-breaking change which fixes an issue)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Refactoring (no functional changes)

How Has This Been Tested?

• Unit Tests
• Integration Tests
• Manual Testing

Screenshots (if applicable)

N/A

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Related Issues

N/A

Quality Recommendations

1. Add unit tests for the new image transformation functionality

2. Add more detailed error logging in imageToJsonHandler.ts with specific error types

3. Consider adding a content-type validation function to avoid repetitive checks

4. Add documentation for the new image transformation feature

Sequence Diagram

sequenceDiagram
    participant Client
    participant ResponseHandler
    participant RetryHandler
    participant ImageToJsonHandler
    participant SegmindProvider

    Client->>ResponseHandler: Request image generation
    ResponseHandler->>RetryHandler: Process request with retry logic
    RetryHandler->>SegmindProvider: Call Segmind API
    SegmindProvider-->>RetryHandler: Return image response
    RetryHandler-->>ResponseHandler: Return response
    
    Note over ResponseHandler: Check if response needs transformation
    alt responseTransformer === 'imageGenerate' && has imageToJson transformer
        ResponseHandler->>ImageToJsonHandler: getTransformedResponse(response, transformer)
        Note over ImageToJsonHandler: Convert image to base64
        ImageToJsonHandler-->>ResponseHandler: Return JSON response
    end
    
    ResponseHandler-->>Client: Return final response

Loading

[narengogi]

@unsync I have a few items in mind wrt this PR

1. we should be handling transforms only for specific routes and not based on response, so if segmind has a /v1/images/generations route, we can do these changes for that route
2. I'm not entirely keen we should load the entire image file into memory and convert it to json, but that's okay

overall I'm saying this PR could only include changes to provider/segmind/imageGenerate.ts and provider/segmind/api.ts and provider/segmind/index.ts

[unsync]

@narengogi if i get this correctly, you would prefer to scope the changes to the providers handlers.

Unless i'm missing something, that would not work without a change in the reponseHandler because of this:

gateway/src/handlers/responseHandlers.ts

Lines 137 to 139 in de327f6

	if (responseContentType?.startsWith(CONTENT_TYPES.GENERIC_IMAGE_PATTERN)) {
	return { response: handleImageResponse(response), responseJson: null };
	}

This prevents to reach the code that handles the transformations:

gateway/src/handlers/responseHandlers.ts

Lines 159 to 166 in de327f6

	const nonStreamingResponse = await handleNonStreamingMode(
	response,
	responseTransformerFunction,
	strictOpenAiCompliance,
	gatewayRequestUrl,
	gatewayRequest,
	areSyncHooksAvailable
	);

And before you hit handleNonStreamingMode, you need to have a JSON body, otherwise this fails:

gateway/src/handlers/streamHandler.ts

Lines 248 to 250 in de327f6

	const isJsonParsingRequired = responseTransformer || areSyncHooksAvailable;
	const originalResponseBodyJson: Record<string, any> | null =
	isJsonParsingRequired ? await response.json() : null;

Since the provider returns a content type image/ then we don't enter the transformer function of the providers (if we did, i definitely would have gone this road, that was my first idea ^^).

Let me know if i missed something here 🙏

And if that doesn't makes it's way to main no hard feelings, i'm using this on a fork, i just wanted for others to benefit from it 👍

[narengogi]

I spent some time on this, I was initially misreading what the PR is doing, this fix makes sense for the problem,

one change we could do is to convert to json only if routes are in the unified routes like v1/images/generations

I'm copying @VisargD on this for a review of the approach

[unsync]

I spent some time on this, I was initially misreading what the PR is doing, this fix makes sense for the problem,

one change we could do is to convert to json only if routes are in the unified routes like v1/images/generations

I'm copying @VisargD on this for a review of the approach

@narengogi @VisargD i made a last commit to the PR to, imho, limit the scope of the behaviour in a more elegant way:

  // if the original response is an image, convert it to JSON if the provider has a transformer
  if (
    responseTransformer === 'imageGenerate' && // check that we are on the imageGenerate route
    responseTransformerFunction && // check that we have a transformer for this provider
    providerTransformers?.[`imageToJson`] && // check that we have a 'imageToJson" transformer for this provider
    response.headers
      ?.get('content-type')
      ?.startsWith(CONTENT_TYPES.GENERIC_IMAGE_PATTERN) // check that the original response content type is an image
  ) {
    // transformers are async, because we read the body as an array buffer
    response = await providerTransformers?.[`imageToJson`](response);
  }

This ensures that :

• we only apply it when the response is initially a image/***
• we are looking for a responseTransformer of value imageGenerate (ensuring to do this only on image generation)
• the provider has the responseTransformerFunction (no point transforming the response if the provider doesn't provide a transformer)
• obviously, checks if the provider has a imageToJson transformer

I also added a try/catch the the transformer wrapper, so that if the transformation fails we return the original response to avoid any failures due to this logic

[matter-code-review]

Summary By MatterAI

🔄 What Changed

This PR adds functionality to transform raw image responses from providers (specifically Segmind) into JSON format. It introduces a new handler for converting image responses to base64-encoded JSON, enhancing the API's response consistency. The implementation includes fallback mechanisms to return the original response if transformation fails.

🔍 Impact of the Change

This change improves client experience by providing consistent JSON responses instead of raw binary image data. It makes it easier for clients to handle image generation responses without needing special handling for binary data. The transformation is specifically targeted at the imageGenerate route and only applies when appropriate conditions are met.

📁 Total Files Changed

• Added: src/handlers/imageToJsonHandler.ts (35 lines)
• Modified: src/handlers/responseHandlers.ts (16 additions, 1 deletion)
• Modified: src/handlers/retryHandler.ts (2 additions)
• Modified: src/providers/segmind/imageGenerate.ts (12 additions)
• Modified: src/providers/segmind/index.ts (2 additions)

🧪 Test Added

N/A - No tests were added in this PR.

🔒Security Vulnerabilities

N/A - No security vulnerabilities were identified.

Type of Change

• New feature (non-breaking change which adds functionality)
• Bug fix (non-breaking change which fixes an issue)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Refactoring (no functional changes)

How Has This Been Tested?

• Unit Tests
• Manual Testing
• Integration Tests

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have added tests that prove my fix is effective or that my feature works
• I have made corresponding changes to the documentation
• My changes generate no new warnings

Sequence Diagram

sequenceDiagram
    participant Client
    participant Gateway
    participant ResponseHandler
    participant ImageToJsonHandler
    participant SegmindProvider
    participant ExternalAPI

    Client->>Gateway: Request image generation
    Gateway->>ExternalAPI: Forward request to provider API
    ExternalAPI-->>Gateway: Return raw image response
    Gateway->>ResponseHandler: Process response
    Note over ResponseHandler: Check if response is for imageGenerate
    Note over ResponseHandler: Check if content-type is image
    Note over ResponseHandler: Check if provider has imageToJson transformer
    ResponseHandler->>SegmindProvider: Call imageToJson transformer
    SegmindProvider->>ImageToJsonHandler: Call getTransformedResponse()
    Note over ImageToJsonHandler: Convert image to base64
    Note over ImageToJsonHandler: Transform to JSON format
    ImageToJsonHandler-->>SegmindProvider: Return JSON Response
    SegmindProvider-->>ResponseHandler: Return transformed response
    ResponseHandler-->>Gateway: Return final response
    Gateway-->>Client: Return JSON response instead of raw image
    Note over Gateway: If transformation fails, original response is returned

Loading