fix: gemini batch embed endpoint

ax-llm · Oct 16, 2024 · 4ebad97 · 4ebad97
1 parent c39a142
commit 4ebad97
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git a/src/ax/ai/google-gemini/api.ts b/src/ax/ai/google-gemini/api.ts
@@ -56,7 +56,7 @@ const safetySettings: AxAIGoogleGeminiSafetySettings = [
 export const axAIGoogleGeminiDefaultConfig = (): AxAIGoogleGeminiConfig =>
   structuredClone({
     model: AxAIGoogleGeminiModel.Gemini15Pro,
-    embedModel: AxAIGoogleGeminiEmbedModel.Embedding001,
+    embedModel: AxAIGoogleGeminiEmbedModel.TextEmbedding004,
     safetySettings,
     ...axBaseAIDefaultConfig()
   });
@@ -65,7 +65,7 @@ export const axAIGoogleGeminiDefaultCreativeConfig =
   (): AxAIGoogleGeminiConfig =>
     structuredClone({
       model: AxAIGoogleGeminiModel.Gemini15Flash,
-      embedModel: AxAIGoogleGeminiEmbedModel.Embedding001,
+      embedModel: AxAIGoogleGeminiEmbedModel.TextEmbedding004,
       safetySettings,
       ...axBaseAIDefaultCreativeConfig()
     });
@@ -345,11 +345,14 @@ export class AxAIGoogleGemini extends AxBaseAI<
     }
 
     const apiConfig = {
-      name: `/models/${model}:batchEmbedText?key=${this.apiKey}`
+      name: `/models/${model}:batchEmbedContents?key=${this.apiKey}`
     };
 
     const reqValue: AxAIGoogleGeminiBatchEmbedRequest = {
-      requests: req.texts.map((text) => ({ model, text }))
+      requests: req.texts.map((text) => ({
+        model: 'models/' + model,
+        content: { parts: [{ text }] }
+      }))
     };
 
     return [apiConfig, reqValue];

diff --git a/src/ax/ai/google-gemini/types.ts b/src/ax/ai/google-gemini/types.ts
@@ -10,7 +10,7 @@ export enum AxAIGoogleGeminiModel {
 }
 
 export enum AxAIGoogleGeminiEmbedModel {
-  Embedding001 = 'embedding-001'
+  TextEmbedding004 = 'text-embedding-004'
 }
 
 export enum AxAIGoogleGeminiSafetyCategory {
@@ -159,7 +159,9 @@ export type AxAIGoogleGeminiConfig = AxModelConfig & {
 export type AxAIGoogleGeminiBatchEmbedRequest = {
   requests: {
     model: string;
-    text: string;
+    content: {
+      parts: { text: string }[];
+    };
   }[];
 };
 

diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts
@@ -3,10 +3,10 @@ import { createHash } from 'crypto';
 import type { AxFunctionJSONSchema } from '../ai/types.js';
 
 import {
-  parseSignature,
   type InputParsedField,
   type OutputParsedField,
-  type ParsedSignature
+  type ParsedSignature,
+  parseSignature
 } from './parser.js';
 
 export interface AxField {