feat: added google search retrieval for gemini

ax-llm · Nov 18, 2024 · f659dc6 · f659dc6
1 parent 4e03050
commit f659dc6
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 3 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -4,7 +4,7 @@
   "editor.formatOnSave": true,
   "typescript.tsdk": "./node_modules/typescript/lib",
   "typescript.enablePromptUseWorkspaceTsdk": true,
-  "typescript.format.enable": false,
+  "typescript.format.enable": true,
   "typescript.tsserver.watchOptions": {},
   "js/ts.implicitProjectConfig.target": "ESNext"
 }
diff --git a/src/ax/ai/google-gemini/api.ts b/src/ax/ai/google-gemini/api.ts
@@ -70,13 +70,21 @@ export const axAIGoogleGeminiDefaultCreativeConfig =
       ...axBaseAIDefaultCreativeConfig()
     });
 
+export interface AxAIGoogleGeminiOptionsTools {
+  codeExecution?: boolean;
+  googleSearchRetrieval?: {
+    mode?: 'MODE_DYNAMIC';
+    dynamicThreshold?: number;
+  };
+}
+
 export interface AxAIGoogleGeminiArgs {
   name: 'google-gemini';
   apiKey: string;
   projectId?: string;
   region?: string;
   config?: Readonly<Partial<AxAIGoogleGeminiConfig>>;
-  options?: Readonly<AxAIServiceOptions & { codeExecution?: boolean }>;
+  options?: Readonly<AxAIServiceOptions & AxAIGoogleGeminiOptionsTools>;
   modelMap?: Record<
     string,
     AxAIGoogleGeminiModel | AxAIGoogleGeminiEmbedModel | string
@@ -285,6 +293,12 @@ export class AxAIGoogleGemini extends AxBaseAI<
       tools.push({ code_execution: {} });
     }
 
+    if (this.options?.googleSearchRetrieval) {
+      tools.push({
+        google_search_retrieval: this.options.googleSearchRetrieval
+      });
+    }
+
     if (tools.length === 0) {
       tools = undefined;
     }

diff --git a/src/ax/ai/google-gemini/types.ts b/src/ax/ai/google-gemini/types.ts
@@ -78,9 +78,15 @@ export type AxAIGoogleGeminiToolFunctionDeclaration = {
   parameters?: object;
 };
 
+export type AxAIGoogleGeminiToolGoogleSearchRetrieval = {
+  mode?: 'MODE_DYNAMIC';
+  dynamic_threshold?: number;
+};
+
 export type AxAIGoogleGeminiTool = {
   function_declarations?: AxAIGoogleGeminiToolFunctionDeclaration[];
   code_execution?: object;
+  google_search_retrieval?: AxAIGoogleGeminiToolGoogleSearchRetrieval;
 };
 
 export type AxAIGoogleGeminiToolConfig = {

diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts
@@ -39,7 +39,15 @@ export class AxSignature {
   private sigHash: string;
   private sigString: string;
 
-  constructor(signature: Readonly<AxSignature | string>) {
+  constructor(signature?: Readonly<AxSignature | string>) {
+    if (!signature) {
+      this.inputFields = [];
+      this.outputFields = [];
+      this.sigHash = '';
+      this.sigString = '';
+      return;
+    }
+
     if (typeof signature === 'string') {
       let sig: ParsedSignature;
       try {