feat: new gemini embedding

Author: dosco

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

@@ -391,6 +391,8 @@ class AxAIGoogleGeminiImpl
         requests: req.texts.map((text) => ({
           model: 'models/' + model,
           content: { parts: [{ text }] },
+          outputDimensionality: this.config.dimensions,
+          taskType: this.config.embedType,
         })),
       }
     }

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

@@ -15,6 +15,7 @@ export enum AxAIGoogleGeminiModel {
 }
 
 export enum AxAIGoogleGeminiEmbedModel {
+  GeminiEmbedding = 'gemini-embedding-exp-03-07',
   TextEmbedding004 = 'text-embedding-004',
   TextEmbedding005 = 'text-embedding-005',
 }
@@ -34,6 +35,17 @@ export enum AxAIGoogleGeminiSafetyThreshold {
   BlockDefault = 'HARM_BLOCK_THRESHOLD_UNSPECIFIED',
 }
 
+export enum AxAIGoogleGeminiEmbedTypes {
+  SemanticSimilarity = 'SEMANTIC_SIMILARITY',
+  Classification = 'CLASSIFICATION',
+  Clustering = 'CLUSTERING',
+  RetrievalDocument = 'RETRIEVAL_DOCUMENT',
+  RetrievalQuery = 'RETRIEVAL_QUERY',
+  QuestionAnswering = 'QUESTION_ANSWERING',
+  FactVerification = 'FACT_VERIFICATION',
+  CodeRetrievalQuery = 'CODE_RETRIEVAL_QUERY',
+}
+
 export type AxAIGoogleGeminiContent =
   | {
       role: 'user'
@@ -169,6 +181,8 @@ export type AxAIGoogleGeminiConfig = AxModelConfig & {
   model: AxAIGoogleGeminiModel
   embedModel?: AxAIGoogleGeminiEmbedModel
   safetySettings?: AxAIGoogleGeminiSafetySettings
+  embedType?: AxAIGoogleGeminiEmbedTypes
+  dimensions?: number
 }
 
 /**

src/ax/ai/openai/api.ts

@@ -102,8 +102,7 @@ class AxAIOpenAIImpl<TModel, TEmbedModel>
 {
   constructor(
     private readonly config: Readonly<AxAIOpenAIConfig<TModel, TEmbedModel>>,
-    private streamingUsage: boolean,
-    private dimensions?: number
+    private streamingUsage: boolean
   ) {}
 
   getModelConfig(): AxModelConfig {
@@ -212,7 +211,7 @@ class AxAIOpenAIImpl<TModel, TEmbedModel>
     const reqValue = {
       model: model,
       input: req.texts,
-      dimensions: this.dimensions,
+      dimensions: this.config.dimensions,
     }
 
     return [apiConfig, reqValue]
@@ -449,8 +448,7 @@ export class AxAIOpenAIBase<TModel, TEmbedModel> extends AxBaseAI<
 
     const aiImpl = new AxAIOpenAIImpl<TModel, TEmbedModel>(
       config,
-      options?.streamingUsage ?? true,
-      config.dimensions
+      options?.streamingUsage ?? true
     )
 
     super(aiImpl, {