diff --git a/sdks/node/src/__tests__/integrations/openai2.test.ts b/sdks/node/src/__tests__/integrations/openai2.test.ts index 9b8f7a86444..1e2f679c9db 100644 --- a/sdks/node/src/__tests__/integrations/openai2.test.ts +++ b/sdks/node/src/__tests__/integrations/openai2.test.ts @@ -1,5 +1,10 @@ import {Api as TraceServerApi} from '../../generated/traceServerApi'; -import {openAIStreamReducer, wrapOpenAI} from '../../integrations/openai'; +import { + makeOpenAIImagesGenerateOp, + openAIStreamReducer, + wrapOpenAI, +} from '../../integrations/openai'; +import {isWeaveImage} from '../../media'; import {WandbServerApi} from '../../wandb/wandbServerApi'; import {WeaveClient} from '../../weaveClient'; @@ -201,3 +206,36 @@ describe('OpenAI Integration', () => { }); }); }); + +describe('makeOpenAIImagesGenerateOp', () => { + it('converts b64_json images to WeaveImage objects and preserves other items', async () => { + const mockGenerate = jest.fn().mockResolvedValue({ + data: [ + { + url: 'https://example.com/image.png', + }, + { + b64_json: + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==', + }, + ], + }); + + const wrappedGenerate = makeOpenAIImagesGenerateOp(mockGenerate); + const result = await wrappedGenerate({prompt: 'draw a picture'}); + + // Verify the result structure + expect(result.data).toHaveLength(2); + + // First item should remain unchanged + expect(result.data[0]).toEqual({url: 'https://example.com/image.png'}); + + // Second item should be converted to WeaveImage + expect(isWeaveImage(result.data[1])).toBe(true); + expect(result.data[1].imageType).toBe('png'); + expect(Buffer.isBuffer(result.data[1].data)).toBe(true); + + // Verify the original function was called with correct args + expect(mockGenerate).toHaveBeenCalledWith({prompt: 'draw a picture'}); + }); +});