Skip to content

Commit 4ebad97

Browse files
committed
fix: gemini batch embed endpoint
1 parent c39a142 commit 4ebad97

File tree

3 files changed

+13
-8
lines changed

3 files changed

+13
-8
lines changed

src/ax/ai/google-gemini/api.ts

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ const safetySettings: AxAIGoogleGeminiSafetySettings = [
5656
export const axAIGoogleGeminiDefaultConfig = (): AxAIGoogleGeminiConfig =>
5757
structuredClone({
5858
model: AxAIGoogleGeminiModel.Gemini15Pro,
59-
embedModel: AxAIGoogleGeminiEmbedModel.Embedding001,
59+
embedModel: AxAIGoogleGeminiEmbedModel.TextEmbedding004,
6060
safetySettings,
6161
...axBaseAIDefaultConfig()
6262
});
@@ -65,7 +65,7 @@ export const axAIGoogleGeminiDefaultCreativeConfig =
6565
(): AxAIGoogleGeminiConfig =>
6666
structuredClone({
6767
model: AxAIGoogleGeminiModel.Gemini15Flash,
68-
embedModel: AxAIGoogleGeminiEmbedModel.Embedding001,
68+
embedModel: AxAIGoogleGeminiEmbedModel.TextEmbedding004,
6969
safetySettings,
7070
...axBaseAIDefaultCreativeConfig()
7171
});
@@ -345,11 +345,14 @@ export class AxAIGoogleGemini extends AxBaseAI<
345345
}
346346

347347
const apiConfig = {
348-
name: `/models/${model}:batchEmbedText?key=${this.apiKey}`
348+
name: `/models/${model}:batchEmbedContents?key=${this.apiKey}`
349349
};
350350

351351
const reqValue: AxAIGoogleGeminiBatchEmbedRequest = {
352-
requests: req.texts.map((text) => ({ model, text }))
352+
requests: req.texts.map((text) => ({
353+
model: 'models/' + model,
354+
content: { parts: [{ text }] }
355+
}))
353356
};
354357

355358
return [apiConfig, reqValue];

src/ax/ai/google-gemini/types.ts

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ export enum AxAIGoogleGeminiModel {
1010
}
1111

1212
export enum AxAIGoogleGeminiEmbedModel {
13-
Embedding001 = 'embedding-001'
13+
TextEmbedding004 = 'text-embedding-004'
1414
}
1515

1616
export enum AxAIGoogleGeminiSafetyCategory {
@@ -159,7 +159,9 @@ export type AxAIGoogleGeminiConfig = AxModelConfig & {
159159
export type AxAIGoogleGeminiBatchEmbedRequest = {
160160
requests: {
161161
model: string;
162-
text: string;
162+
content: {
163+
parts: { text: string }[];
164+
};
163165
}[];
164166
};
165167

src/ax/dsp/sig.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@ import { createHash } from 'crypto';
33
import type { AxFunctionJSONSchema } from '../ai/types.js';
44

55
import {
6-
parseSignature,
76
type InputParsedField,
87
type OutputParsedField,
9-
type ParsedSignature
8+
type ParsedSignature,
9+
parseSignature
1010
} from './parser.js';
1111

1212
export interface AxField {

0 commit comments

Comments
 (0)