From f659dc6151f837d4b59cd14ab4ed8525b8c57c0c Mon Sep 17 00:00:00 2001
From: dosco <832235+dosco@users.noreply.github.com>
Date: Sun, 17 Nov 2024 17:29:25 -0800
Subject: [PATCH] feat: added google search retrieval for gemini

---
 .vscode/settings.json            |  2 +-
 src/ax/ai/google-gemini/api.ts   | 16 +++++++++++++++-
 src/ax/ai/google-gemini/types.ts |  6 ++++++
 src/ax/dsp/sig.ts                | 10 +++++++++-
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 1d732fb8..59155184 100644
--- 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
index 691a44a8..6da0885f 100644
--- 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
index 99449745..eb215284 100644
--- 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
index e71e8eaa..5bfe3de1 100644
--- 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 {