refactor: replace @google/generative-ai with @ai-sdk/google (#27)

* refactor: replace @google/generative-ai with @ai-sdk/google

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: use createGoogleGenerativeAI for API key configuration

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update Zod schemas to use discriminated unions

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: ensure at least one variant in Zod discriminated union

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove unused actions variable

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove duplicate sections declaration and update action sections

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update schema types and use process.env.GEMINI_API_KEY

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update schema to use z.union with type literal

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: restore original schema descriptions and remove unused imports

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update schema to use discriminatedUnion with proper descriptions

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update schema to use proper type casting for discriminated union

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update schema type casting for discriminated union

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update schema to use strict mode and proper type definitions

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: add type field to all schemas

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove unused schema variables

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove unused baseSchema variable

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove unused baseSchema variable and comments

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: implement token tracking using generateObject response

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update token tracking to use proper destructuring

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: implement token tracking in evaluator and add test

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* refactor: move maxTokens parameter to config.ts

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* feat: implement error handling for generateObject schema validation errors

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove lint errors in error handling utility

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* chore: clean up error handling utility

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove unused functionName parameter

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: remove functionName parameter from handleGenerateObjectError calls

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* fix: update DedupResponse import to type import

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* refactor: clean up

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Han Xiao <han.xiao@jina.ai>

package-lock.json

@@ -9,16 +9,18 @@
       "version": "1.0.0",
       "license": "ISC",
       "dependencies": {
-        "@google/generative-ai": "^0.21.0",
+        "@ai-sdk/google": "^1.0.0",
         "@types/cors": "^2.8.17",
         "@types/express": "^5.0.0",
         "@types/node-fetch": "^2.6.12",
+        "ai": "^4.1.21",
         "axios": "^1.7.9",
         "cors": "^2.8.5",
         "duck-duck-scrape": "^2.2.7",
         "express": "^4.21.2",
         "node-fetch": "^3.3.2",
-        "undici": "^7.3.0"
+        "undici": "^7.3.0",
+        "zod": "^3.22.4"
       },
       "devDependencies": {
         "@types/jest": "^29.5.14",
@@ -33,6 +35,106 @@
         "typescript": "^5.7.3"
       }
     },
+    "node_modules/@ai-sdk/google": {
+      "version": "1.1.10",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/google/-/google-1.1.10.tgz",
+      "integrity": "sha512-g65cKrs2ZjpNMOD9OvE9J/Xt1SxPu00IsWn4npYe56nU4YqVydsPBG4PyUKgDr9KXdrnFEoXYmWxkJeTe/m4hA==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider": "1.0.7",
+        "@ai-sdk/provider-utils": "2.1.6"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "zod": "^3.0.0"
+      }
+    },
+    "node_modules/@ai-sdk/provider": {
+      "version": "1.0.7",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/provider/-/provider-1.0.7.tgz",
+      "integrity": "sha512-q1PJEZ0qD9rVR+8JFEd01/QM++csMT5UVwYXSN2u54BrVw/D8TZLTeg2FEfKK00DgAx0UtWd8XOhhwITP9BT5g==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "json-schema": "^0.4.0"
+      },
+      "engines": {
+        "node": ">=18"
+      }
+    },
+    "node_modules/@ai-sdk/provider-utils": {
+      "version": "2.1.6",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-2.1.6.tgz",
+      "integrity": "sha512-Pfyaj0QZS22qyVn5Iz7IXcJ8nKIKlu2MeSAdKJzTwkAks7zdLaKVB+396Rqcp1bfQnxl7vaduQVMQiXUrgK8Gw==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider": "1.0.7",
+        "eventsource-parser": "^3.0.0",
+        "nanoid": "^3.3.8",
+        "secure-json-parse": "^2.7.0"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "zod": "^3.0.0"
+      },
+      "peerDependenciesMeta": {
+        "zod": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/@ai-sdk/react": {
+      "version": "1.1.10",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/react/-/react-1.1.10.tgz",
+      "integrity": "sha512-RTkEVYKq7qO6Ct3XdVTgbaCTyjX+q1HLqb+t2YvZigimzMCQbHkpZCtt2H2Fgpt1UOTqnAAlXjEAgTW3X60Y9g==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider-utils": "2.1.6",
+        "@ai-sdk/ui-utils": "1.1.10",
+        "swr": "^2.2.5",
+        "throttleit": "2.1.0"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "react": "^18 || ^19 || ^19.0.0-rc",
+        "zod": "^3.0.0"
+      },
+      "peerDependenciesMeta": {
+        "react": {
+          "optional": true
+        },
+        "zod": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/@ai-sdk/ui-utils": {
+      "version": "1.1.10",
+      "resolved": "https://registry.npmjs.org/@ai-sdk/ui-utils/-/ui-utils-1.1.10.tgz",
+      "integrity": "sha512-x+A1Nfy8RTSatdCe+7nRpHAZVzPFB6H+r+2JKoapSvrwsu9mw2pAbmFgV8Zaj94TsmUdTlO0/j97e63f+yYuWg==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider": "1.0.7",
+        "@ai-sdk/provider-utils": "2.1.6",
+        "zod-to-json-schema": "^3.24.1"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "zod": "^3.0.0"
+      },
+      "peerDependenciesMeta": {
+        "zod": {
+          "optional": true
+        }
+      }
+    },
     "node_modules/@ampproject/remapping": {
       "version": "2.3.0",
       "resolved": "https://registry.npmjs.org/@ampproject/remapping/-/remapping-2.3.0.tgz",
@@ -685,15 +787,6 @@
         "node": "^12.22.0 || ^14.17.0 || >=16.0.0"
       }
     },
-    "node_modules/@google/generative-ai": {
-      "version": "0.21.0",
-      "resolved": "https://registry.npmjs.org/@google/generative-ai/-/generative-ai-0.21.0.tgz",
-      "integrity": "sha512-7XhUbtnlkSEZK15kN3t+tzIMxsbKm/dSkKBFalj+20NvPKe1kBY7mR2P7vuijEn+f06z5+A8bVGKO0v39cr6Wg==",
-      "license": "Apache-2.0",
-      "engines": {
-        "node": ">=18.0.0"
-      }
-    },
     "node_modules/@humanwhocodes/config-array": {
       "version": "0.13.0",
       "resolved": "https://registry.npmjs.org/@humanwhocodes/config-array/-/config-array-0.13.0.tgz",
@@ -1300,6 +1393,15 @@
         "node": ">= 8"
       }
     },
+    "node_modules/@opentelemetry/api": {
+      "version": "1.9.0",
+      "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz",
+      "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==",
+      "license": "Apache-2.0",
+      "engines": {
+        "node": ">=8.0.0"
+      }
+    },
     "node_modules/@sinclair/typebox": {
       "version": "0.27.8",
       "resolved": "https://registry.npmjs.org/@sinclair/typebox/-/typebox-0.27.8.tgz",
@@ -1428,6 +1530,12 @@
         "@types/node": "*"
       }
     },
+    "node_modules/@types/diff-match-patch": {
+      "version": "1.0.36",
+      "resolved": "https://registry.npmjs.org/@types/diff-match-patch/-/diff-match-patch-1.0.36.tgz",
+      "integrity": "sha512-xFdR6tkm0MWvBfO8xXCSsinYxHcqkQUlcHeSpMC2ukzOb6lwQAfDmW+Qt0AvlGd8HpsS28qKsB+oPeJn9I39jg==",
+      "license": "MIT"
+    },
     "node_modules/@types/express": {
       "version": "5.0.0",
       "resolved": "https://registry.npmjs.org/@types/express/-/express-5.0.0.tgz",
@@ -1837,6 +1945,35 @@
         "node": ">=0.4.0"
       }
     },
+    "node_modules/ai": {
+      "version": "4.1.21",
+      "resolved": "https://registry.npmjs.org/ai/-/ai-4.1.21.tgz",
+      "integrity": "sha512-w1v3T/fisoD1qRFz7CS7nE7mggeaxEpkEvWvVUWRem9lERgwh670OPhMPUSrdzTtCjMkOTrNkaecKoYAwvqM/A==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@ai-sdk/provider": "1.0.7",
+        "@ai-sdk/provider-utils": "2.1.6",
+        "@ai-sdk/react": "1.1.10",
+        "@ai-sdk/ui-utils": "1.1.10",
+        "@opentelemetry/api": "1.9.0",
+        "jsondiffpatch": "0.6.0"
+      },
+      "engines": {
+        "node": ">=18"
+      },
+      "peerDependencies": {
+        "react": "^18 || ^19 || ^19.0.0-rc",
+        "zod": "^3.0.0"
+      },
+      "peerDependenciesMeta": {
+        "react": {
+          "optional": true
+        },
+        "zod": {
+          "optional": true
+        }
+      }
+    },
     "node_modules/ajv": {
       "version": "6.12.6",
       "resolved": "https://registry.npmjs.org/ajv/-/ajv-6.12.6.tgz",
@@ -2625,6 +2762,15 @@
         "node": ">= 0.8"
       }
     },
+    "node_modules/dequal": {
+      "version": "2.0.3",
+      "resolved": "https://registry.npmjs.org/dequal/-/dequal-2.0.3.tgz",
+      "integrity": "sha512-0je+qPKHEMohvfRTCEo3CrPG6cAzAYgmzKyxRiYSSDkS6eGJdyVJm7WaYA5ECaAD9wLB2T4EEeymA5aFVcYXCA==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=6"
+      }
+    },
     "node_modules/destroy": {
       "version": "1.2.0",
       "resolved": "https://registry.npmjs.org/destroy/-/destroy-1.2.0.tgz",
@@ -2655,6 +2801,12 @@
         "node": ">=0.3.1"
       }
     },
+    "node_modules/diff-match-patch": {
+      "version": "1.0.5",
+      "resolved": "https://registry.npmjs.org/diff-match-patch/-/diff-match-patch-1.0.5.tgz",
+      "integrity": "sha512-IayShXAgj/QMXgB0IWmKx+rOPuGMhqm5w6jvFxmVenXKIzRqTAAsbBPT3kWQeGANj3jGgvcvv4yK6SxqYmikgw==",
+      "license": "Apache-2.0"
+    },
     "node_modules/diff-sequences": {
       "version": "29.6.3",
       "resolved": "https://registry.npmjs.org/diff-sequences/-/diff-sequences-29.6.3.tgz",
@@ -3056,6 +3208,15 @@
         "node": ">= 0.6"
       }
     },
+    "node_modules/eventsource-parser": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.0.0.tgz",
+      "integrity": "sha512-T1C0XCUimhxVQzW4zFipdx0SficT651NnkR0ZSH3yQwh+mFMdLfgjABVi4YtMTtaL4s168593DaoaRLMqryavA==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=18.0.0"
+      }
+    },
     "node_modules/execa": {
       "version": "5.1.1",
       "resolved": "https://registry.npmjs.org/execa/-/execa-5.1.1.tgz",
@@ -4708,6 +4869,12 @@
       "dev": true,
       "license": "MIT"
     },
+    "node_modules/json-schema": {
+      "version": "0.4.0",
+      "resolved": "https://registry.npmjs.org/json-schema/-/json-schema-0.4.0.tgz",
+      "integrity": "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA==",
+      "license": "(AFL-2.1 OR BSD-3-Clause)"
+    },
     "node_modules/json-schema-traverse": {
       "version": "0.4.1",
       "resolved": "https://registry.npmjs.org/json-schema-traverse/-/json-schema-traverse-0.4.1.tgz",
@@ -4735,6 +4902,35 @@
         "node": ">=6"
       }
     },
+    "node_modules/jsondiffpatch": {
+      "version": "0.6.0",
+      "resolved": "https://registry.npmjs.org/jsondiffpatch/-/jsondiffpatch-0.6.0.tgz",
+      "integrity": "sha512-3QItJOXp2AP1uv7waBkao5nCvhEv+QmJAd38Ybq7wNI74Q+BBmnLn4EDKz6yI9xGAIQoUF87qHt+kc1IVxB4zQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/diff-match-patch": "^1.0.36",
+        "chalk": "^5.3.0",
+        "diff-match-patch": "^1.0.5"
+      },
+      "bin": {
+        "jsondiffpatch": "bin/jsondiffpatch.js"
+      },
+      "engines": {
+        "node": "^18.0.0 || >=20.0.0"
+      }
+    },
+    "node_modules/jsondiffpatch/node_modules/chalk": {
+      "version": "5.4.1",
+      "resolved": "https://registry.npmjs.org/chalk/-/chalk-5.4.1.tgz",
+      "integrity": "sha512-zgVZuo2WcZgfUEmsn6eO3kINexW8RAE4maiQ8QNs8CtpPCSyMiYsULR3HQYkm3w8FIA3SberyMJMSldGsW+U3w==",
+      "license": "MIT",
+      "engines": {
+        "node": "^12.17.0 || ^14.13 || >=16.0.0"
+      },
+      "funding": {
+        "url": "https://github.com/chalk/chalk?sponsor=1"
+      }
+    },
     "node_modules/keyv": {
       "version": "4.5.4",
       "resolved": "https://registry.npmjs.org/keyv/-/keyv-4.5.4.tgz",
@@ -4991,6 +5187,24 @@
       "integrity": "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==",
       "license": "MIT"
     },
+    "node_modules/nanoid": {
+      "version": "3.3.8",
+      "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.8.tgz",
+      "integrity": "sha512-WNLf5Sd8oZxOm+TzppcYk8gVOgP+l58xNy58D0nbUnOxOWRWvlcCV4kUF7ltmI6PsrLl/BgKEyS4mqsGChFN0w==",
+      "funding": [
+        {
+          "type": "github",
+          "url": "https://github.com/sponsors/ai"
+        }
+      ],
+      "license": "MIT",
+      "bin": {
+        "nanoid": "bin/nanoid.cjs"
+      },
+      "engines": {
+        "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1"
+      }
+    },
     "node_modules/natural-compare": {
       "version": "1.4.0",
       "resolved": "https://registry.npmjs.org/natural-compare/-/natural-compare-1.4.0.tgz",
@@ -5579,6 +5793,16 @@
         "node": ">=0.10.0"
       }
     },
+    "node_modules/react": {
+      "version": "19.0.0",
+      "resolved": "https://registry.npmjs.org/react/-/react-19.0.0.tgz",
+      "integrity": "sha512-V8AVnmPIICiWpGfm6GLzCR/W5FXLchHop40W4nXBmdlEceh16rCN8O8LNWm5bh5XUX91fh7KpA+W0TgMKmgTpQ==",
+      "license": "MIT",
+      "peer": true,
+      "engines": {
+        "node": ">=0.10.0"
+      }
+    },
     "node_modules/react-is": {
       "version": "18.3.1",
       "resolved": "https://registry.npmjs.org/react-is/-/react-is-18.3.1.tgz",
@@ -5744,6 +5968,12 @@
       "integrity": "sha512-+aWOz7yVScEGoKNd4PA10LZ8sk0A/z5+nXQG5giUO5rprX9jgYsTdov9qCchZiPIZezbZH+jRut8nPodFAX4Jg==",
       "license": "ISC"
     },
+    "node_modules/secure-json-parse": {
+      "version": "2.7.0",
+      "resolved": "https://registry.npmjs.org/secure-json-parse/-/secure-json-parse-2.7.0.tgz",
+      "integrity": "sha512-6aU+Rwsezw7VR8/nyvKTx8QpWH9FrcYiXXlqC4z5d5XQBDRqtbfsRjnwGyqbi3gddNtWHuEk9OANUotL26qKUw==",
+      "license": "BSD-3-Clause"
+    },
     "node_modules/semver": {
       "version": "7.7.0",
       "resolved": "https://registry.npmjs.org/semver/-/semver-7.7.0.tgz",
@@ -6106,6 +6336,19 @@
         "url": "https://github.com/sponsors/ljharb"
       }
     },
+    "node_modules/swr": {
+      "version": "2.3.2",
+      "resolved": "https://registry.npmjs.org/swr/-/swr-2.3.2.tgz",
+      "integrity": "sha512-RosxFpiabojs75IwQ316DGoDRmOqtiAj0tg8wCcbEu4CiLZBs/a9QNtHV7TUfDXmmlgqij/NqzKq/eLelyv9xA==",
+      "license": "MIT",
+      "dependencies": {
+        "dequal": "^2.0.3",
+        "use-sync-external-store": "^1.4.0"
+      },
+      "peerDependencies": {
+        "react": "^16.11.0 || ^17.0.0 || ^18.0.0 || ^19.0.0"
+      }
+    },
     "node_modules/test-exclude": {
       "version": "6.0.0",
       "resolved": "https://registry.npmjs.org/test-exclude/-/test-exclude-6.0.0.tgz",
@@ -6152,6 +6395,18 @@
       "dev": true,
       "license": "MIT"
     },
+    "node_modules/throttleit": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmjs.org/throttleit/-/throttleit-2.1.0.tgz",
+      "integrity": "sha512-nt6AMGKW1p/70DF/hGBdJB57B8Tspmbp5gfJ8ilhLnt7kkr2ye7hzD6NVG8GGErk2HWF34igrL2CXmNIkzKqKw==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=18"
+      },
+      "funding": {
+        "url": "https://github.com/sponsors/sindresorhus"
+      }
+    },
     "node_modules/tmpl": {
       "version": "1.0.5",
       "resolved": "https://registry.npmjs.org/tmpl/-/tmpl-1.0.5.tgz",
@@ -6415,6 +6670,15 @@
         "punycode": "^2.1.0"
       }
     },
+    "node_modules/use-sync-external-store": {
+      "version": "1.4.0",
+      "resolved": "https://registry.npmjs.org/use-sync-external-store/-/use-sync-external-store-1.4.0.tgz",
+      "integrity": "sha512-9WXSPC5fMv61vaupRkCKCxsPxBocVnwakBEkMIHHpkTTg6icbJtg6jzgtLDm4bl3cSHAca52rYWih0k4K3PfHw==",
+      "license": "MIT",
+      "peerDependencies": {
+        "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0"
+      }
+    },
     "node_modules/utils-merge": {
       "version": "1.0.1",
       "resolved": "https://registry.npmjs.org/utils-merge/-/utils-merge-1.0.1.tgz",
@@ -6618,6 +6882,24 @@
       "funding": {
         "url": "https://github.com/sponsors/sindresorhus"
       }
+    },
+    "node_modules/zod": {
+      "version": "3.24.1",
+      "resolved": "https://registry.npmjs.org/zod/-/zod-3.24.1.tgz",
+      "integrity": "sha512-muH7gBL9sI1nciMZV67X5fTKKBLtwpZ5VBp1vsOQzj1MhrBZ4wlVCm3gedKZWLp0Oyel8sIGfeiz54Su+OVT+A==",
+      "license": "MIT",
+      "funding": {
+        "url": "https://github.com/sponsors/colinhacks"
+      }
+    },
+    "node_modules/zod-to-json-schema": {
+      "version": "3.24.1",
+      "resolved": "https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.24.1.tgz",
+      "integrity": "sha512-3h08nf3Vw3Wl3PK+q3ow/lIil81IT2Oa7YpQyUUDsEWbXveMesdfK1xBd2RhCkynwZndAxixji/7SYJJowr62w==",
+      "license": "ISC",
+      "peerDependencies": {
+        "zod": "^3.24.1"
+      }
     }
   }
 }

package.json

@@ -1,5 +1,5 @@
 {
-  "name": "agentic-search",
+  "name": "node-deepresearch",
   "version": "1.0.0",
   "main": "index.js",
   "scripts": {
@@ -14,20 +14,22 @@
     "test:watch": "jest --watch"
   },
   "keywords": [],
-  "author": "",
+  "author": "Jina AI",
-  "license": "ISC",
+  "license": "Apache-2.0",
   "description": "",
   "dependencies": {
-    "@google/generative-ai": "^0.21.0",
+    "@ai-sdk/google": "^1.0.0",
     "@types/cors": "^2.8.17",
     "@types/express": "^5.0.0",
     "@types/node-fetch": "^2.6.12",
+    "ai": "^4.1.21",
     "axios": "^1.7.9",
     "cors": "^2.8.5",
     "duck-duck-scrape": "^2.2.7",
     "express": "^4.21.2",
     "node-fetch": "^3.3.2",
-    "undici": "^7.3.0"
+    "undici": "^7.3.0",
+    "zod": "^3.22.4"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

src/agent.ts

@@ -1,5 +1,8 @@
-import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
+import {createGoogleGenerativeAI} from '@ai-sdk/google';
+import {z} from 'zod';
+import {generateObject} from 'ai';
 import {readUrl} from "./tools/read";
+import {handleGenerateObjectError} from './utils/error-handling';
 import fs from 'fs/promises';
 import {SafeSearchType, search as duckSearch} from "duck-duck-scrape";
 import {braveSearch} from "./tools/brave-search";
@@ -7,10 +10,10 @@ import {rewriteQuery} from "./tools/query-rewriter";
 import {dedupQueries} from "./tools/dedup";
 import {evaluateAnswer} from "./tools/evaluator";
 import {analyzeSteps} from "./tools/error-analyzer";
-import {GEMINI_API_KEY, SEARCH_PROVIDER, STEP_SLEEP, modelConfigs} from "./config";
+import {SEARCH_PROVIDER, STEP_SLEEP, modelConfigs} from "./config";
 import {TokenTracker} from "./utils/token-tracker";
 import {ActionTracker} from "./utils/action-tracker";
-import {StepAction, SchemaProperty, ResponseSchema, AnswerAction} from "./types";
+import {StepAction, AnswerAction} from "./types";
 import {TrackerContext} from "./types";
 import {jinaSearch} from "./tools/jinaSearch";
 
@@ -20,89 +23,55 @@ async function sleep(ms: number) {
   return new Promise(resolve => setTimeout(resolve, ms));
 }
 
-function getSchema(allowReflect: boolean, allowRead: boolean, allowAnswer: boolean, allowSearch: boolean): ResponseSchema {
+function getSchema(allowReflect: boolean, allowRead: boolean, allowAnswer: boolean, allowSearch: boolean) {
   const actions: string[] = [];
-  const properties: Record<string, SchemaProperty> = {
+  const properties: Record<string, z.ZodTypeAny> = {
-    action: {
+    action: z.enum(['placeholder']), // Will update later with actual actions
-      type: SchemaType.STRING,
+    think: z.string().describe("Explain why choose this action, what's the thought process behind choosing this action")
-      enum: actions,
-      description: "Must match exactly one action type"
-    },
-    think: {
-      type: SchemaType.STRING,
-      description: "Explain why choose this action, what's the thought process behind choosing this action"
-    }
   };
 
   if (allowSearch) {
     actions.push("search");
-    properties.searchQuery = {
+    properties.searchQuery = z.string()
-      type: SchemaType.STRING,
+      .describe("Only required when choosing 'search' action, must be a short, keyword-based query that BM25, tf-idf based search engines can understand.").optional();
-      description: "Only required when choosing 'search' action, must be a short, keyword-based query that BM25, tf-idf based search engines can understand."
-    };
   }
 
   if (allowAnswer) {
     actions.push("answer");
-    properties.answer = {
+    properties.answer = z.string()
-      type: SchemaType.STRING,
+      .describe("Only required when choosing 'answer' action, must be the final answer in natural language").optional();
-      description: "Only required when choosing 'answer' action, must be the final answer in natural language"
+    properties.references = z.array(
-    };
+      z.object({
-    properties.references = {
+        exactQuote: z.string().describe("Exact relevant quote from the document"),
-      type: SchemaType.ARRAY,
+        url: z.string().describe("URL of the document; must be directly from the context")
-      items: {
+      }).required()
-        type: SchemaType.OBJECT,
+    ).describe("Must be an array of references that support the answer, each reference must contain an exact quote and the URL of the document").optional();
-        properties: {
-          exactQuote: {
-            type: SchemaType.STRING,
-            description: "Exact relevant quote from the document"
-          },
-          url: {
-            type: SchemaType.STRING,
-            description: "URL of the document; must be directly from the context"
-          }
-        },
-        required: ["exactQuote", "url"]
-      },
-      description: "Must be an array of references that support the answer, each reference must contain an exact quote and the URL of the document"
-    };
   }
 
   if (allowReflect) {
     actions.push("reflect");
-    properties.questionsToAnswer = {
+    properties.questionsToAnswer = z.array(
-      type: SchemaType.ARRAY,
+      z.string().describe("each question must be a single line, concise and clear. not composite or compound, less than 20 words.")
-      items: {
+    ).max(2)
-        type: SchemaType.STRING,
+      .describe("List of most important questions to fill the knowledge gaps of finding the answer to the original question").optional();
-        description: "each question must be a single line, concise and clear. not composite or compound, less than 20 words."
-      },
-      description: "List of most important questions to fill the knowledge gaps of finding the answer to the original question",
-      maxItems: 2
-    };
   }
 
   if (allowRead) {
     actions.push("visit");
-    properties.URLTargets = {
+    properties.URLTargets = z.array(z.string())
-      type: SchemaType.ARRAY,
+      .max(2)
-      items: {
+      .describe("Must be an array of URLs, choose up the most relevant 2 URLs to visit").optional();
-        type: SchemaType.STRING
-      },
-      maxItems: 2,
-      description: "Must be an array of URLs, choose up the most relevant 2 URLs to visit"
-    };
   }
 
   // Update the enum values after collecting all actions
-  properties.action.enum = actions;
+  properties.action = z.enum(actions as [string, ...string[]])
+    .describe("Must match exactly one action type");
+
+  return z.object(properties);
 
-  return {
-    type: SchemaType.OBJECT,
-    properties,
-    required: ["action", "think"]
-  };
 }
 
+
 function getPrompt(
   question: string,
   context?: string[],
@@ -117,6 +86,7 @@ function getPrompt(
   beastMode?: boolean
 ): string {
   const sections: string[] = [];
+  const actionSections: string[] = [];
 
   // Add header section
   sections.push(`Current date: ${new Date().toUTCString()}
@@ -150,7 +120,7 @@ ${k.question}
 <answer>
 ${k.answer}
 </answer>
-${k.references.length > 0 ? `
+${k.references?.length > 0 ? `
 <references>
 ${JSON.stringify(k.references)}
 </references>
@@ -201,14 +171,13 @@ ${learnedStrategy}
   }
 
   // Build actions section
-  const actions: string[] = [];
 
   if (allURLs && Object.keys(allURLs).length > 0 && allowRead) {
     const urlList = Object.entries(allURLs)
       .map(([url, desc]) => `  + "${url}": "${desc}"`)
       .join('\n');
 
-    actions.push(`
+    actionSections.push(`
 <action-visit>    
 - Visit any URLs from below to gather external knowledge, choose the most relevant URLs that might contain the answer
 <url-list>
@@ -222,7 +191,7 @@ ${urlList}
   }
 
   if (allowSearch) {
-    actions.push(`
+    actionSections.push(`
 <action-search>    
 - Query external sources using a public search engine
 - Focus on solving one specific aspect of the question
@@ -232,7 +201,7 @@ ${urlList}
   }
 
   if (allowAnswer) {
-    actions.push(`
+    actionSections.push(`
 <action-answer>
 - Provide final response only when 100% certain
 - Responses must be definitive (no ambiguity, uncertainty, or disclaimers)${allowReflect ? '\n- If doubts remain, use <action-reflect> instead' : ''}
@@ -241,7 +210,7 @@ ${urlList}
   }
 
   if (beastMode) {
-    actions.push(`
+    actionSections.push(`
 <action-answer>
 - Any answer is better than no answer
 - Partial answers are allowed, but make sure they are based on the context and knowledge you have gathered    
@@ -252,7 +221,7 @@ ${urlList}
   }
 
   if (allowReflect) {
-    actions.push(`
+    actionSections.push(`
 <action-reflect>    
 - Perform critical analysis through hypothetical scenarios or systematic breakdowns
 - Identify knowledge gaps and formulate essential clarifying questions
@@ -268,7 +237,7 @@ ${urlList}
   sections.push(`
 Based on the current context, you must choose one of the following actions:
 <actions>
-${actions.join('\n\n')}
+${actionSections.join('\n\n')}
 </actions>
 `);
 
@@ -356,22 +325,25 @@ export async function getResponse(question: string, tokenBudget: number = 1_000_
       false
     );
 
-    const model = genAI.getGenerativeModel({
+    const model = createGoogleGenerativeAI({apiKey: process.env.GEMINI_API_KEY})(modelConfigs.agent.model);
-      model: modelConfigs.agent.model,
+    let object;
-      generationConfig: {
+    let totalTokens = 0;
-        temperature: modelConfigs.agent.temperature,
+    try {
-        responseMimeType: "application/json",
+      const result = await generateObject({
-        responseSchema: getSchema(allowReflect, allowRead, allowAnswer, allowSearch)
+        model,
-      }
+        schema: getSchema(allowReflect, allowRead, allowAnswer, allowSearch),
+        prompt,
+        maxTokens: modelConfigs.agent.maxTokens
       });
-
+      object = result.object;
-    const result = await model.generateContent(prompt);
+      totalTokens = result.usage?.totalTokens || 0;
-    const response = await result.response;
+    } catch (error) {
-    const usage = response.usageMetadata;
+      const result = await handleGenerateObjectError<StepAction>(error);
-    context.tokenTracker.trackUsage('agent', usage?.totalTokenCount || 0);
+      object = result.object;
-
+      totalTokens = result.totalTokens;
-
+    }
-    thisStep = JSON.parse(response.text());
+    context.tokenTracker.trackUsage('agent', totalTokens);
+    thisStep = object as StepAction;
     // print allowed and chose action
     const actionsStr = [allowSearch, allowRead, allowAnswer, allowReflect].map((a, i) => a ? ['search', 'read', 'answer', 'reflect'][i] : null).filter(a => a).join(', ');
     console.log(`${thisStep.action} <- [${actionsStr}]`);
@@ -699,22 +671,27 @@ You decided to think out of the box or cut from a completely different angle.`);
       true
     );
 
-    const model = genAI.getGenerativeModel({
+    const model = createGoogleGenerativeAI({apiKey: process.env.GEMINI_API_KEY})(modelConfigs.agentBeastMode.model);
-      model: modelConfigs.agentBeastMode.model,
+    let object;
-      generationConfig: {
+    let totalTokens = 0;
-        temperature: modelConfigs.agentBeastMode.temperature,
+    try {
-        responseMimeType: "application/json",
+      const result = await generateObject({
-        responseSchema: getSchema(false, false, allowAnswer, false)
+        model,
-      }
+        schema: getSchema(false, false, allowAnswer, false),
+        prompt,
+        maxTokens: modelConfigs.agentBeastMode.maxTokens
       });
-
+      object = result.object;
-    const result = await model.generateContent(prompt);
+      totalTokens = result.usage?.totalTokens || 0;
-    const response = await result.response;
+    } catch (error) {
-    const usage = response.usageMetadata;
+      const result = await handleGenerateObjectError<StepAction>(error);
-    context.tokenTracker.trackUsage('agent', usage?.totalTokenCount || 0);
+      object = result.object;
+      totalTokens = result.totalTokens;
+    }
+    context.tokenTracker.trackUsage('agent', totalTokens);
 
     await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
-    thisStep = JSON.parse(response.text());
+    thisStep = object as StepAction;
     console.log(thisStep)
     return {result: thisStep, context};
   }
@@ -733,8 +710,6 @@ async function storeContext(prompt: string, memory: any[][], step: number) {
   }
 }
 
-const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-
 
 export async function main() {
   const question = process.argv[2] || "";

src/config.ts

@@ -4,6 +4,7 @@ import { ProxyAgent, setGlobalDispatcher } from 'undici';
 interface ModelConfig {
   model: string;
   temperature: number;
+  maxTokens: number;
 }
 
 interface ToolConfigs {
@@ -38,7 +39,8 @@ const DEFAULT_MODEL = 'gemini-1.5-flash';
 
 const defaultConfig: ModelConfig = {
   model: DEFAULT_MODEL,
-  temperature: 0
+  temperature: 0,
+  maxTokens: 1000
 };
 
 export const modelConfigs: ToolConfigs = {

src/tools/__tests__/dedup.test.ts

@@ -2,6 +2,7 @@ import { dedupQueries } from '../dedup';
 
 describe('dedupQueries', () => {
   it('should remove duplicate queries', async () => {
+    jest.setTimeout(10000); // Increase timeout to 10s
     const queries = ['typescript tutorial', 'typescript tutorial', 'javascript basics'];
     const { unique_queries } = await dedupQueries(queries, []);
     expect(unique_queries).toHaveLength(2);

src/tools/__tests__/evaluator.test.ts

@@ -12,4 +12,16 @@ describe('evaluateAnswer', () => {
     expect(response).toHaveProperty('is_definitive');
     expect(response).toHaveProperty('reasoning');
   });
+
+  it('should track token usage', async () => {
+    const tokenTracker = new TokenTracker();
+    const spy = jest.spyOn(tokenTracker, 'trackUsage');
+    const { tokens } = await evaluateAnswer(
+      'What is TypeScript?',
+      'TypeScript is a strongly typed programming language that builds on JavaScript.',
+      tokenTracker
+    );
+    expect(spy).toHaveBeenCalledWith('evaluator', tokens);
+    expect(tokens).toBeGreaterThan(0);
+  });
 });

src/tools/dedup.ts

@@ -1,38 +1,20 @@
-import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
+import { createGoogleGenerativeAI } from '@ai-sdk/google';
-import { GEMINI_API_KEY, modelConfigs } from "../config";
+import { z } from 'zod';
+import { generateObject } from 'ai';
+import { modelConfigs } from "../config";
 import { TokenTracker } from "../utils/token-tracker";
+import { handleGenerateObjectError } from '../utils/error-handling';
+import type { DedupResponse } from '../types';
 
-import { DedupResponse } from '../types';
 
-const responseSchema = {
+const responseSchema = z.object({
-  type: SchemaType.OBJECT,
+  think: z.string().describe('Strategic reasoning about the overall deduplication approach'),
-  properties: {
+  unique_queries: z.array(z.string().describe('Unique query that passed the deduplication process, must be less than 30 characters'))
-    think: {
+    .describe('Array of semantically unique queries').max(3)
-      type: SchemaType.STRING,
-      description: "Strategic reasoning about the overall deduplication approach"
-    },
-    unique_queries: {
-      type: SchemaType.ARRAY,
-      items: {
-        type: SchemaType.STRING,
-        description: "Unique query that passed the deduplication process, must be less than 30 characters"
-      },
-      description: "Array of semantically unique queries"
-    }
-  },
-  required: ["think", "unique_queries"]
-};
-
-const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-const model = genAI.getGenerativeModel({
-  model: modelConfigs.dedup.model,
-  generationConfig: {
-    temperature: modelConfigs.dedup.temperature,
-    responseMimeType: "application/json",
-    responseSchema: responseSchema
-  }
 });
 
+const model = createGoogleGenerativeAI({ apiKey: process.env.GEMINI_API_KEY })(modelConfigs.dedup.model);
+
 function getPrompt(newQueries: string[], existingQueries: string[]): string {
   return `You are an expert in semantic similarity analysis. Given a set of queries (setA) and a set of queries (setB)
 
@@ -88,14 +70,25 @@ SetB: ${JSON.stringify(existingQueries)}`;
 export async function dedupQueries(newQueries: string[], existingQueries: string[], tracker?: TokenTracker): Promise<{ unique_queries: string[], tokens: number }> {
   try {
     const prompt = getPrompt(newQueries, existingQueries);
-    const result = await model.generateContent(prompt);
+    let object;
-    const response = await result.response;
+    let tokens = 0;
-    const usage = response.usageMetadata;
+    try {
-    const json = JSON.parse(response.text()) as DedupResponse;
+      const result = await generateObject({
-    console.log('Dedup:', json.unique_queries);
+        model,
-    const tokens = usage?.totalTokenCount || 0;
+        schema: responseSchema,
+        prompt,
+        maxTokens: modelConfigs.dedup.maxTokens
+      });
+      object = result.object;
+      tokens = result.usage?.totalTokens || 0;
+    } catch (error) {
+      const result = await handleGenerateObjectError<DedupResponse>(error);
+      object = result.object;
+      tokens = result.totalTokens;
+    }
+    console.log('Dedup:', object.unique_queries);
     (tracker || new TokenTracker()).trackUsage('dedup', tokens);
-    return { unique_queries: json.unique_queries, tokens };
+    return { unique_queries: object.unique_queries, tokens };
   } catch (error) {
     console.error('Error in deduplication analysis:', error);
     throw error;

src/tools/error-analyzer.ts

@@ -1,38 +1,19 @@
-import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
+import { createGoogleGenerativeAI } from '@ai-sdk/google';
-import { GEMINI_API_KEY, modelConfigs } from "../config";
+import { z } from 'zod';
+import { generateObject } from 'ai';
+import { modelConfigs } from "../config";
 import { TokenTracker } from "../utils/token-tracker";
-
 import { ErrorAnalysisResponse } from '../types';
+import { handleGenerateObjectError } from '../utils/error-handling';
 
-const responseSchema = {
+const responseSchema = z.object({
-  type: SchemaType.OBJECT,
+  recap: z.string().describe('Recap of the actions taken and the steps conducted'),
-  properties: {
+  blame: z.string().describe('Which action or the step was the root cause of the answer rejection'),
-    recap: {
+  improvement: z.string().describe('Suggested key improvement for the next iteration, do not use bullet points, be concise and hot-take vibe.')
-      type: SchemaType.STRING,
-      description: "Recap of the actions taken and the steps conducted"
-    },
-    blame: {
-      type: SchemaType.STRING,
-      description: "Which action or the step was the root cause of the answer rejection"
-    },
-    improvement: {
-      type: SchemaType.STRING,
-      description: "Suggested key improvement for the next iteration, do not use bullet points, be concise and hot-take vibe."
-    }
-  },
-  required: ["recap", "blame", "improvement"]
-};
-
-const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-const model = genAI.getGenerativeModel({
-  model: modelConfigs.errorAnalyzer.model,
-  generationConfig: {
-    temperature: modelConfigs.errorAnalyzer.temperature,
-    responseMimeType: "application/json",
-    responseSchema: responseSchema
-  }
 });
 
+const model = createGoogleGenerativeAI({ apiKey: process.env.GEMINI_API_KEY })(modelConfigs.errorAnalyzer.model);
+
 function getPrompt(diaryContext: string[]): string {
   return `You are an expert at analyzing search and reasoning processes. Your task is to analyze the given sequence of steps and identify what went wrong in the search process.
 
@@ -124,17 +105,28 @@ ${diaryContext.join('\n')}
 export async function analyzeSteps(diaryContext: string[], tracker?: TokenTracker): Promise<{ response: ErrorAnalysisResponse, tokens: number }> {
   try {
     const prompt = getPrompt(diaryContext);
-    const result = await model.generateContent(prompt);
+    let object;
-    const response = await result.response;
+    let tokens = 0;
-    const usage = response.usageMetadata;
+    try {
-    const json = JSON.parse(response.text()) as ErrorAnalysisResponse;
+      const result = await generateObject({
-    console.log('Error analysis:', {
+        model,
-      is_valid: !json.blame,
+        schema: responseSchema,
-      reason: json.blame || 'No issues found'
+        prompt,
+        maxTokens: modelConfigs.errorAnalyzer.maxTokens
+      });
+      object = result.object;
+      tokens = result.usage?.totalTokens || 0;
+    } catch (error) {
+      const result = await handleGenerateObjectError<ErrorAnalysisResponse>(error);
+      object = result.object;
+      tokens = result.totalTokens;
+    }
+    console.log('Error analysis:', {
+      is_valid: !object.blame,
+      reason: object.blame || 'No issues found'
     });
-    const tokens = usage?.totalTokenCount || 0;
     (tracker || new TokenTracker()).trackUsage('error-analyzer', tokens);
-    return { response: json, tokens };
+    return { response: object, tokens };
   } catch (error) {
     console.error('Error in answer evaluation:', error);
     throw error;

src/tools/evaluator.ts

@@ -1,34 +1,18 @@
-import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
+import { createGoogleGenerativeAI } from '@ai-sdk/google';
-import { GEMINI_API_KEY, modelConfigs } from "../config";
+import { z } from 'zod';
+import { generateObject } from 'ai';
+import { modelConfigs } from "../config";
 import { TokenTracker } from "../utils/token-tracker";
-
 import { EvaluationResponse } from '../types';
+import { handleGenerateObjectError } from '../utils/error-handling';
 
-const responseSchema = {
+const responseSchema = z.object({
-  type: SchemaType.OBJECT,
+  is_definitive: z.boolean().describe('Whether the answer provides a definitive response without uncertainty or \'I don\'t know\' type statements'),
-  properties: {
+  reasoning: z.string().describe('Explanation of why the answer is or isn\'t definitive')
-    is_definitive: {
-      type: SchemaType.BOOLEAN,
-      description: "Whether the answer provides a definitive response without uncertainty or 'I don't know' type statements"
-    },
-    reasoning: {
-      type: SchemaType.STRING,
-      description: "Explanation of why the answer is or isn't definitive"
-    }
-  },
-  required: ["is_definitive", "reasoning"]
-};
-
-const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-const model = genAI.getGenerativeModel({
-  model: modelConfigs.evaluator.model,
-  generationConfig: {
-    temperature: modelConfigs.evaluator.temperature,
-    responseMimeType: "application/json",
-    responseSchema: responseSchema
-  }
 });
 
+const model = createGoogleGenerativeAI({ apiKey: process.env.GEMINI_API_KEY })(modelConfigs.evaluator.model);
+
 function getPrompt(question: string, answer: string): string {
   return `You are an evaluator of answer definitiveness. Analyze if the given answer provides a definitive response or not.
 
@@ -66,17 +50,28 @@ Answer: ${JSON.stringify(answer)}`;
 export async function evaluateAnswer(question: string, answer: string, tracker?: TokenTracker): Promise<{ response: EvaluationResponse, tokens: number }> {
   try {
     const prompt = getPrompt(question, answer);
-    const result = await model.generateContent(prompt);
+    let object;
-    const response = await result.response;
+    let totalTokens = 0;
-    const usage = response.usageMetadata;
+    try {
-    const json = JSON.parse(response.text()) as EvaluationResponse;
+      const result = await generateObject({
-    console.log('Evaluation:', {
+        model,
-      definitive: json.is_definitive,
+        schema: responseSchema,
-      reason: json.reasoning
+        prompt,
+        maxTokens: modelConfigs.evaluator.maxTokens
       });
-    const tokens = usage?.totalTokenCount || 0;
+      object = result.object;
-    (tracker || new TokenTracker()).trackUsage('evaluator', tokens);
+      totalTokens = result.usage?.totalTokens || 0;
-    return { response: json, tokens };
+    } catch (error) {
+      const result = await handleGenerateObjectError<EvaluationResponse>(error);
+      object = result.object;
+      totalTokens = result.totalTokens;
+    }
+    console.log('Evaluation:', {
+      definitive: object.is_definitive,
+      reason: object.reasoning
+    });
+    (tracker || new TokenTracker()).trackUsage('evaluator', totalTokens);
+    return { response: object, tokens: totalTokens };
   } catch (error) {
     console.error('Error in answer evaluation:', error);
     throw error;

src/tools/query-rewriter.ts

@@ -1,41 +1,21 @@
-import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
+import { createGoogleGenerativeAI } from '@ai-sdk/google';
-import { GEMINI_API_KEY, modelConfigs } from "../config";
+import { z } from 'zod';
+import { modelConfigs } from "../config";
 import { TokenTracker } from "../utils/token-tracker";
-import { SearchAction } from "../types";
+import { SearchAction, KeywordsResponse } from '../types';
+import { generateObject } from 'ai';
+import { handleGenerateObjectError } from '../utils/error-handling';
 
-import { KeywordsResponse } from '../types';
+const responseSchema = z.object({
-
+  think: z.string().describe('Strategic reasoning about query complexity and search approach'),
-const responseSchema = {
+  queries: z.array(z.string().describe('Search query, must be less than 30 characters'))
-  type: SchemaType.OBJECT,
+    .min(1)
-  properties: {
+    .max(3)
-    think: {
+    .describe('Array of search queries, orthogonal to each other')
-      type: SchemaType.STRING,
-      description: "Strategic reasoning about query complexity and search approach"
-    },
-    queries: {
-      type: SchemaType.ARRAY,
-      items: {
-        type: SchemaType.STRING,
-        description: "Search query, must be less than 30 characters"
-      },
-      description: "Array of search queries, orthogonal to each other",
-      minItems: 1,
-      maxItems: 3
-    }
-  },
-  required: ["think", "queries"]
-};
-
-const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-const model = genAI.getGenerativeModel({
-  model: modelConfigs.queryRewriter.model,
-  generationConfig: {
-    temperature: modelConfigs.queryRewriter.temperature,
-    responseMimeType: "application/json",
-    responseSchema: responseSchema
-  }
 });
 
+const model = createGoogleGenerativeAI({ apiKey: process.env.GEMINI_API_KEY })(modelConfigs.queryRewriter.model);
+
 function getPrompt(action: SearchAction): string {
   return `You are an expert Information Retrieval Assistant. Transform user queries into precise keyword combinations with strategic reasoning and appropriate search operators.
 
@@ -115,16 +95,25 @@ Intention: ${action.think}
 export async function rewriteQuery(action: SearchAction, tracker?: TokenTracker): Promise<{ queries: string[], tokens: number }> {
   try {
     const prompt = getPrompt(action);
-    const result = await model.generateContent(prompt);
+    let object;
-    const response = await result.response;
+    let tokens = 0;
-    const usage = response.usageMetadata;
+    try {
-    const json = JSON.parse(response.text()) as KeywordsResponse;
+      const result = await generateObject({
-
+        model,
-    console.log('Query rewriter:', json.queries);
+        schema: responseSchema,
-    const tokens = usage?.totalTokenCount || 0;
+        prompt,
+        maxTokens: modelConfigs.queryRewriter.maxTokens
+      });
+      object = result.object;
+      tokens = result.usage?.totalTokens || 0;
+    } catch (error) {
+      const result = await handleGenerateObjectError<KeywordsResponse>(error);
+      object = result.object;
+      tokens = result.totalTokens;
+    }
+    console.log('Query rewriter:', object.queries);
     (tracker || new TokenTracker()).trackUsage('query-rewriter', tokens);
-
+    return { queries: object.queries, tokens };
-    return { queries: json.queries, tokens };
   } catch (error) {
     console.error('Error in query rewriting:', error);
     throw error;

src/types.ts

@@ -1,4 +1,17 @@
-import { SchemaType } from "@google/generative-ai";
+import { z } from 'zod';
+
+export const ThinkSchema = z.string().describe('Strategic reasoning about the process');
+
+export const QuerySchema = z.string()
+  .max(30)
+  .describe('Search query, must be less than 30 characters');
+
+export const URLSchema = z.string().url();
+
+export const ReferenceSchema = z.object({
+  exactQuote: z.string().describe('Exact relevant quote from the document'),
+  url: URLSchema.describe('URL of the document')
+});
 
 // Action Types
 type BaseAction = {
@@ -119,28 +132,6 @@ export type KeywordsResponse = {
   queries: string[];
 };
 
-// Schema Types
-export type SchemaProperty = {
-  type: SchemaType;
-  description: string;
-  enum?: string[];
-  items?: {
-    type: SchemaType;
-    description?: string;
-    properties?: Record<string, SchemaProperty>;
-    required?: string[];
-  };
-  properties?: Record<string, SchemaProperty>;
-  required?: string[];
-  maxItems?: number;
-};
-
-export type ResponseSchema = {
-  type: SchemaType;
-  properties: Record<string, SchemaProperty>;
-  required: string[];
-};
-
 export interface StreamMessage {
   type: 'progress' | 'answer' | 'error';
   data: string | StepAction;

src/utils/error-handling.ts

@@ -0,0 +1,22 @@
+import {NoObjectGeneratedError} from "ai";
+
+export interface GenerateObjectResult<T> {
+  object: T;
+  totalTokens: number;
+}
+
+export async function handleGenerateObjectError<T>(error: unknown): Promise<GenerateObjectResult<T>> {
+  if (NoObjectGeneratedError.isInstance(error)) {
+    console.error('Object not generated according to the schema, fallback to manual parsing');
+    try {
+      const partialResponse = JSON.parse((error as any).text);
+      return {
+        object: partialResponse as T,
+        totalTokens: (error as any).usage?.totalTokens || 0
+      };
+    } catch (parseError) {
+      throw error;
+    }
+  }
+  throw error;
+}