Add Gemini as image comparison algorithm #473

related to Visual-Regression-Tracker/Visual-Regression-Tracker#473

Author: pashidlos

package-lock.json

@@ -30,6 +30,7 @@
   "dependencies": {
     "@aws-sdk/client-s3": "^3.922.0",
     "@aws-sdk/s3-request-presigner": "^3.922.0",
+    "@google/genai": "^1.34.0",
     "@nestjs/cache-manager": "^3.0.1",
     "@nestjs/common": "^11.1.8",
     "@nestjs/config": "^4.0.2",
@@ -65,7 +66,8 @@
     "rxjs": "^7.8.2",
     "swagger-ui-express": "^4.6.3",
     "uuid-apikey": "^1.5.3",
-    "zod": "^4.2.1"
+    "zod": "^4.2.1",
+    "zod-to-json-schema": "^3.25.0"
   },
   "devDependencies": {
     "@darraghor/eslint-plugin-nestjs-typed": "^6.9.3",

package.json

@@ -4,13 +4,14 @@ import { LookSameService } from './libs/looks-same/looks-same.service';
 import { OdiffService } from './libs/odiff/odiff.service';
 import { PixelmatchService } from './libs/pixelmatch/pixelmatch.service';
 import { VlmService } from './libs/vlm/vlm.service';
-import { OllamaController } from './libs/vlm/ollama.controller';
-import { OllamaService } from './libs/vlm/ollama.service';
+import { OllamaController } from './libs/vlm/providers/ollama/ollama.controller';
+import { OllamaService } from './libs/vlm/providers/ollama/ollama.service';
+import { GeminiService } from './libs/vlm/providers/gemini/gemini.service';
 import { StaticModule } from '../static/static.module';
 
 @Module({
   controllers: [OllamaController],
-  providers: [CompareService, PixelmatchService, LookSameService, OdiffService, VlmService, OllamaService],
+  providers: [CompareService, PixelmatchService, LookSameService, OdiffService, VlmService, OllamaService, GeminiService],
   imports: [StaticModule],
   exports: [CompareService],
 })

src/compare/compare.module.ts

@@ -6,7 +6,8 @@ import { LookSameService } from './libs/looks-same/looks-same.service';
 import { OdiffService } from './libs/odiff/odiff.service';
 import { PixelmatchService } from './libs/pixelmatch/pixelmatch.service';
 import { VlmService } from './libs/vlm/vlm.service';
-import { OllamaService } from './libs/vlm/ollama.service';
+import { OllamaService } from './libs/vlm/providers/ollama/ollama.service';
+import { GeminiService } from './libs/vlm/providers/gemini/gemini.service';
 import { StaticModule } from '../static/static.module';
 import { ImageComparison } from '@prisma/client';
 import * as utils from '../static/utils';
@@ -26,6 +27,7 @@ describe('CompareService', () => {
         LookSameService,
         VlmService,
         OllamaService,
+        GeminiService,
         PrismaService,
         {
           provide: ConfigService,

src/compare/compare.service.spec.ts

@@ -0,0 +1,128 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { GeminiService } from './gemini.service';
+import { GeminiVlmConfig } from '../../vlm.types';
+
+// Mock the @google/genai module
+const mockGenerateContent = jest.fn();
+
+jest.mock('@google/genai', () => {
+  return {
+    GoogleGenAI: jest.fn().mockImplementation(() => ({
+      models: {
+        generateContent: mockGenerateContent,
+      },
+    })),
+  };
+});
+
+describe('GeminiService', () => {
+  let service: GeminiService;
+
+  beforeEach(async () => {
+    jest.clearAllMocks();
+
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [GeminiService],
+    }).compile();
+
+    service = module.get<GeminiService>(GeminiService);
+  });
+
+  const createConfig = (overrides?: Partial<GeminiVlmConfig>): GeminiVlmConfig => ({
+    provider: 'gemini',
+    model: 'gemini-1.5-pro',
+    prompt: 'Test prompt',
+    temperature: 0.1,
+    apiKey: 'test-api-key',
+    ...overrides,
+  });
+
+  const createMockResponse = (text: string) => ({
+    text,
+  });
+
+  describe('generate', () => {
+    it('should call Gemini SDK with correct parameters and return VlmProviderResponse', async () => {
+      const config = createConfig();
+      const testBytes = new Uint8Array([1, 2, 3, 4]);
+      const mockResponse = createMockResponse('{"identical": true, "description": "No differences"}');
+      mockGenerateContent.mockResolvedValue(mockResponse);
+
+      const result = await service.generate(config, [testBytes]);
+
+      expect(mockGenerateContent).toHaveBeenCalledWith({
+        model: config.model,
+        contents: [
+          { text: config.prompt },
+          {
+            inlineData: {
+              data: expect.any(String),
+              mimeType: 'image/png',
+            },
+          },
+        ],
+        config: {
+          temperature: config.temperature,
+          responseMimeType: 'application/json',
+          responseJsonSchema: expect.any(Object),
+        },
+      });
+      expect(result.content).toBe('{"identical": true, "description": "No differences"}');
+    });
+
+    it.each([
+      ['single image', [new Uint8Array([137, 80, 78, 71])], 2],
+      ['multiple images', [new Uint8Array([1, 2, 3]), new Uint8Array([4, 5, 6]), new Uint8Array([7, 8, 9])], 4],
+    ])('should handle %s and convert to base64', async (_, images, expectedPartsCount) => {
+      const config = createConfig();
+      const mockResponse = createMockResponse('{"identical": true}');
+      mockGenerateContent.mockResolvedValue(mockResponse);
+
+      await service.generate(config, images);
+
+      const callArgs = mockGenerateContent.mock.calls[0][0];
+      expect(callArgs.contents.length).toBe(expectedPartsCount);
+      
+      // Verify first image is base64 encoded
+      if (images.length > 0) {
+        const imagePart = callArgs.contents[1];
+        expect(imagePart.inlineData.mimeType).toBe('image/png');
+        expect(imagePart.inlineData.data).toBe(Buffer.from(images[0]).toString('base64'));
+      }
+    });
+
+    it('should always include hardcoded JSON schema', async () => {
+      const config = createConfig();
+      const mockResponse = createMockResponse('{"identical": true}');
+      mockGenerateContent.mockResolvedValue(mockResponse);
+
+      await service.generate(config, []);
+
+      const callArgs = mockGenerateContent.mock.calls[0][0];
+      expect(callArgs.config.responseMimeType).toBe('application/json');
+      expect(callArgs.config.responseJsonSchema).toEqual(
+        expect.objectContaining({
+          type: 'object',
+          properties: expect.objectContaining({
+            identical: expect.any(Object),
+            description: expect.any(Object),
+          }),
+        })
+      );
+    });
+
+    it.each([
+      ['API key is missing', { apiKey: '' }, 'Gemini API key is required'],
+      ['SDK call fails', { apiKey: 'test-api-key' }, 'API Error'],
+    ])('should throw error when %s', async (_, overrides, expectedError) => {
+      const config = createConfig(overrides);
+      
+      if (expectedError === 'API Error') {
+        mockGenerateContent.mockRejectedValue(new Error(expectedError));
+      }
+      
+      await expect(service.generate(config, [])).rejects.toThrow(expectedError);
+    });
+  });
+});
+

src/compare/libs/vlm/providers/gemini/gemini.service.spec.ts

@@ -0,0 +1,65 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { GoogleGenAI } from '@google/genai';
+import { zodToJsonSchema } from 'zod-to-json-schema';
+import { VlmProvider, VlmProviderResponse } from '../../vlm-provider.interface';
+import { VlmConfig, GeminiVlmConfig } from '../../vlm.types';
+import { VlmComparisonResultSchema } from '../../vlm.service';
+
+@Injectable()
+export class GeminiService implements VlmProvider {
+  private readonly logger: Logger = new Logger(GeminiService.name);
+
+  private getGenAI(apiKey: string): GoogleGenAI {
+    if (!apiKey) {
+      throw new Error('Gemini API key is required');
+    }
+    return new GoogleGenAI({ apiKey });
+  }
+
+  private imageToBase64(imageBytes: Uint8Array): string {
+    const base64 = Buffer.from(imageBytes).toString('base64');
+    return base64;
+  }
+
+  async generate(config: VlmConfig, images: Uint8Array[]): Promise<VlmProviderResponse> {
+    // Type guard: ensure this is Gemini config
+    if (config.provider !== 'gemini') {
+      throw new Error(`GeminiService requires GeminiVlmConfig, got ${config.provider}`);
+    }
+
+    const geminiConfig: GeminiVlmConfig = config;
+
+    const genAI = this.getGenAI(geminiConfig.apiKey);
+
+    try {
+      const imageParts = images.map((img) => ({
+        inlineData: {
+          data: this.imageToBase64(img),
+          mimeType: 'image/png',
+        },
+      }));
+
+      const parts = [
+        { text: geminiConfig.prompt },
+        ...imageParts,
+      ];
+
+      const result = await genAI.models.generateContent({
+        model: geminiConfig.model,
+        contents: parts,
+        config: {
+          temperature: geminiConfig.temperature,
+          responseMimeType: 'application/json' as const,
+          responseJsonSchema: zodToJsonSchema(VlmComparisonResultSchema),
+        },
+      });
+
+      return {
+        content: result.text,
+      };
+    } catch (error) {
+      this.logger.error(`Gemini generate request failed: ${error.message}`, error.stack);
+      throw error;
+    }
+  }
+}

src/compare/libs/vlm/providers/gemini/gemini.service.ts

@@ -0,0 +1,64 @@
+import {
+  Controller,
+  Get,
+  Post,
+  Query,
+  HttpException,
+  HttpStatus,
+  UseInterceptors,
+  UploadedFiles,
+} from '@nestjs/common';
+import { FilesInterceptor } from '@nestjs/platform-express';
+import { ApiTags, ApiConsumes, ApiBody } from '@nestjs/swagger';
+import { OllamaService } from './ollama.service';
+import { VlmProviderResponse } from '../../vlm-provider.interface';
+import { OllamaVlmConfig } from '../../vlm.types';
+
+@ApiTags('Ollama')
+@Controller('ollama')
+export class OllamaController {
+  constructor(private readonly ollamaService: OllamaService) {}
+
+  @Get('models')
+  async listModels() {
+    return { models: await this.ollamaService.listModels() };
+  }
+
+  @Post('compare')
+  @ApiConsumes('multipart/form-data')
+  @ApiBody({
+    schema: {
+      type: 'object',
+      required: ['images'],
+      properties: {
+        images: {
+          type: 'array',
+          items: { type: 'string', format: 'binary' },
+          description: 'Two images to compare (baseline and comparison)',
+        },
+      },
+    },
+  })
+  @UseInterceptors(FilesInterceptor('images', 2))
+  async compareImages(
+    @UploadedFiles() files: Express.Multer.File[],
+    @Query('model') model: string,
+    @Query('prompt') prompt: string,
+    @Query('temperature') temperature: string
+  ): Promise<VlmProviderResponse> {
+    if (files?.length !== 2) {
+      throw new HttpException('Two images required', HttpStatus.BAD_REQUEST);
+    }
+
+    const config: OllamaVlmConfig = {
+      provider: 'ollama',
+      model,
+      prompt,
+      temperature: Number(temperature),
+    };
+    const images = files.map((f) => new Uint8Array(f.buffer));
+
+    return this.ollamaService.generate(config, images);
+  }
+}
+