diff --git a/package-lock.json b/package-lock.json index c96cd7f..0f0c1ce 100644 --- a/package-lock.json +++ b/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" + } } } } diff --git a/package.json b/package.json index fa6cf53..aa25be8 100644 --- a/package.json +++ b/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": "", - "license": "ISC", + "author": "Jina AI", + "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", diff --git a/src/agent.ts b/src/agent.ts index 2b0616c..ef11abd 100644 --- a/src/agent.ts +++ b/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 = { - action: { - type: SchemaType.STRING, - 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" - } + const properties: Record = { + action: z.enum(['placeholder']), // Will update later with actual actions + think: z.string().describe("Explain why choose this action, what's the thought process behind choosing this action") }; if (allowSearch) { actions.push("search"); - properties.searchQuery = { - type: SchemaType.STRING, - description: "Only required when choosing 'search' action, must be a short, keyword-based query that BM25, tf-idf based search engines can understand." - }; + properties.searchQuery = z.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(); } if (allowAnswer) { actions.push("answer"); - properties.answer = { - type: SchemaType.STRING, - description: "Only required when choosing 'answer' action, must be the final answer in natural language" - }; - properties.references = { - type: SchemaType.ARRAY, - items: { - type: SchemaType.OBJECT, - 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" - }; + properties.answer = z.string() + .describe("Only required when choosing 'answer' action, must be the final answer in natural language").optional(); + properties.references = z.array( + z.object({ + exactQuote: z.string().describe("Exact relevant quote from the document"), + url: z.string().describe("URL of the document; must be directly from the context") + }).required() + ).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(); } if (allowReflect) { actions.push("reflect"); - properties.questionsToAnswer = { - type: SchemaType.ARRAY, - items: { - type: SchemaType.STRING, - 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 - }; + properties.questionsToAnswer = z.array( + z.string().describe("each question must be a single line, concise and clear. not composite or compound, less than 20 words.") + ).max(2) + .describe("List of most important questions to fill the knowledge gaps of finding the answer to the original question").optional(); } if (allowRead) { actions.push("visit"); - properties.URLTargets = { - type: SchemaType.ARRAY, - items: { - type: SchemaType.STRING - }, - maxItems: 2, - description: "Must be an array of URLs, choose up the most relevant 2 URLs to visit" - }; + properties.URLTargets = z.array(z.string()) + .max(2) + .describe("Must be an array of URLs, choose up the most relevant 2 URLs to visit").optional(); } // 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} ${k.answer} -${k.references.length > 0 ? ` +${k.references?.length > 0 ? ` ${JSON.stringify(k.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(` - Visit any URLs from below to gather external knowledge, choose the most relevant URLs that might contain the answer @@ -222,7 +191,7 @@ ${urlList} } if (allowSearch) { - actions.push(` + actionSections.push(` - 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(` - Provide final response only when 100% certain - Responses must be definitive (no ambiguity, uncertainty, or disclaimers)${allowReflect ? '\n- If doubts remain, use instead' : ''} @@ -241,7 +210,7 @@ ${urlList} } if (beastMode) { - actions.push(` + actionSections.push(` - 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(` - 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.join('\n\n')} +${actionSections.join('\n\n')} `); @@ -356,22 +325,25 @@ export async function getResponse(question: string, tokenBudget: number = 1_000_ false ); - const model = genAI.getGenerativeModel({ - model: modelConfigs.agent.model, - generationConfig: { - temperature: modelConfigs.agent.temperature, - responseMimeType: "application/json", - responseSchema: getSchema(allowReflect, allowRead, allowAnswer, allowSearch) - } - }); - - const result = await model.generateContent(prompt); - const response = await result.response; - const usage = response.usageMetadata; - context.tokenTracker.trackUsage('agent', usage?.totalTokenCount || 0); - - - thisStep = JSON.parse(response.text()); + const model = createGoogleGenerativeAI({apiKey: process.env.GEMINI_API_KEY})(modelConfigs.agent.model); + let object; + let totalTokens = 0; + try { + const result = await generateObject({ + model, + schema: getSchema(allowReflect, allowRead, allowAnswer, allowSearch), + prompt, + maxTokens: modelConfigs.agent.maxTokens + }); + object = result.object; + totalTokens = result.usage?.totalTokens || 0; + } catch (error) { + const result = await handleGenerateObjectError(error); + object = result.object; + totalTokens = result.totalTokens; + } + 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}]`); @@ -683,8 +655,8 @@ You decided to think out of the box or cut from a completely different angle.`); } else { console.log('Enter Beast mode!!!') // any answer is better than no answer, humanity last resort - step ++; - totalStep ++; + step++; + totalStep++; const prompt = getPrompt( question, diaryContext, @@ -699,22 +671,27 @@ You decided to think out of the box or cut from a completely different angle.`); true ); - const model = genAI.getGenerativeModel({ - model: modelConfigs.agentBeastMode.model, - generationConfig: { - temperature: modelConfigs.agentBeastMode.temperature, - responseMimeType: "application/json", - responseSchema: getSchema(false, false, allowAnswer, false) - } - }); - - const result = await model.generateContent(prompt); - const response = await result.response; - const usage = response.usageMetadata; - context.tokenTracker.trackUsage('agent', usage?.totalTokenCount || 0); + const model = createGoogleGenerativeAI({apiKey: process.env.GEMINI_API_KEY})(modelConfigs.agentBeastMode.model); + let object; + let totalTokens = 0; + try { + const result = await generateObject({ + model, + schema: getSchema(false, false, allowAnswer, false), + prompt, + maxTokens: modelConfigs.agentBeastMode.maxTokens + }); + object = result.object; + totalTokens = result.usage?.totalTokens || 0; + } catch (error) { + const result = await handleGenerateObjectError(error); + 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] || ""; diff --git a/src/config.ts b/src/config.ts index 3543b5f..da7bd11 100644 --- a/src/config.ts +++ b/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 = { diff --git a/src/tools/__tests__/dedup.test.ts b/src/tools/__tests__/dedup.test.ts index ed615a7..321cfb6 100644 --- a/src/tools/__tests__/dedup.test.ts +++ b/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); diff --git a/src/tools/__tests__/evaluator.test.ts b/src/tools/__tests__/evaluator.test.ts index 391b4bb..4335e02 100644 --- a/src/tools/__tests__/evaluator.test.ts +++ b/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); + }); }); diff --git a/src/tools/dedup.ts b/src/tools/dedup.ts index c48f5ce..a23c967 100644 --- a/src/tools/dedup.ts +++ b/src/tools/dedup.ts @@ -1,38 +1,20 @@ -import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai"; -import { GEMINI_API_KEY, modelConfigs } from "../config"; +import { createGoogleGenerativeAI } from '@ai-sdk/google'; +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 = { - type: SchemaType.OBJECT, - properties: { - think: { - 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 responseSchema = z.object({ + think: z.string().describe('Strategic reasoning about the overall deduplication approach'), + unique_queries: z.array(z.string().describe('Unique query that passed the deduplication process, must be less than 30 characters')) + .describe('Array of semantically unique queries').max(3) }); +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); - const response = await result.response; - const usage = response.usageMetadata; - const json = JSON.parse(response.text()) as DedupResponse; - console.log('Dedup:', json.unique_queries); - const tokens = usage?.totalTokenCount || 0; + let object; + let tokens = 0; + try { + const result = await generateObject({ + model, + schema: responseSchema, + prompt, + maxTokens: modelConfigs.dedup.maxTokens + }); + object = result.object; + tokens = result.usage?.totalTokens || 0; + } catch (error) { + const result = await handleGenerateObjectError(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; diff --git a/src/tools/error-analyzer.ts b/src/tools/error-analyzer.ts index 7f55b94..7eaf8a3 100644 --- a/src/tools/error-analyzer.ts +++ b/src/tools/error-analyzer.ts @@ -1,38 +1,19 @@ -import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai"; -import { GEMINI_API_KEY, modelConfigs } from "../config"; +import { createGoogleGenerativeAI } from '@ai-sdk/google'; +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 = { - type: SchemaType.OBJECT, - properties: { - recap: { - 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 responseSchema = z.object({ + recap: z.string().describe('Recap of the actions taken and the steps conducted'), + blame: z.string().describe('Which action or the step was the root cause of the answer rejection'), + improvement: z.string().describe('Suggested key improvement for the next iteration, do not use bullet points, be concise and hot-take vibe.') }); +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); - const response = await result.response; - const usage = response.usageMetadata; - const json = JSON.parse(response.text()) as ErrorAnalysisResponse; + let object; + let tokens = 0; + try { + const result = await generateObject({ + model, + schema: responseSchema, + prompt, + maxTokens: modelConfigs.errorAnalyzer.maxTokens + }); + object = result.object; + tokens = result.usage?.totalTokens || 0; + } catch (error) { + const result = await handleGenerateObjectError(error); + object = result.object; + tokens = result.totalTokens; + } console.log('Error analysis:', { - is_valid: !json.blame, - reason: json.blame || 'No issues found' + 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; diff --git a/src/tools/evaluator.ts b/src/tools/evaluator.ts index 222e0d2..ce14f96 100644 --- a/src/tools/evaluator.ts +++ b/src/tools/evaluator.ts @@ -1,34 +1,18 @@ -import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai"; -import { GEMINI_API_KEY, modelConfigs } from "../config"; +import { createGoogleGenerativeAI } from '@ai-sdk/google'; +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 = { - type: SchemaType.OBJECT, - properties: { - 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 responseSchema = z.object({ + is_definitive: z.boolean().describe('Whether the answer provides a definitive response without uncertainty or \'I don\'t know\' type statements'), + reasoning: z.string().describe('Explanation of why the answer is or isn\'t definitive') }); +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); - const response = await result.response; - const usage = response.usageMetadata; - const json = JSON.parse(response.text()) as EvaluationResponse; + let object; + let totalTokens = 0; + try { + const result = await generateObject({ + model, + schema: responseSchema, + prompt, + maxTokens: modelConfigs.evaluator.maxTokens + }); + object = result.object; + totalTokens = result.usage?.totalTokens || 0; + } catch (error) { + const result = await handleGenerateObjectError(error); + object = result.object; + totalTokens = result.totalTokens; + } console.log('Evaluation:', { - definitive: json.is_definitive, - reason: json.reasoning + definitive: object.is_definitive, + reason: object.reasoning }); - const tokens = usage?.totalTokenCount || 0; - (tracker || new TokenTracker()).trackUsage('evaluator', tokens); - return { response: json, tokens }; + (tracker || new TokenTracker()).trackUsage('evaluator', totalTokens); + return { response: object, tokens: totalTokens }; } catch (error) { console.error('Error in answer evaluation:', error); throw error; diff --git a/src/tools/query-rewriter.ts b/src/tools/query-rewriter.ts index 8410413..431715a 100644 --- a/src/tools/query-rewriter.ts +++ b/src/tools/query-rewriter.ts @@ -1,41 +1,21 @@ -import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai"; -import { GEMINI_API_KEY, modelConfigs } from "../config"; +import { createGoogleGenerativeAI } from '@ai-sdk/google'; +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 = { - type: SchemaType.OBJECT, - properties: { - think: { - 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 responseSchema = z.object({ + think: z.string().describe('Strategic reasoning about query complexity and search approach'), + queries: z.array(z.string().describe('Search query, must be less than 30 characters')) + .min(1) + .max(3) + .describe('Array of search queries, orthogonal to each other') }); +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,18 +95,27 @@ 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); - const response = await result.response; - const usage = response.usageMetadata; - const json = JSON.parse(response.text()) as KeywordsResponse; - - console.log('Query rewriter:', json.queries); - const tokens = usage?.totalTokenCount || 0; + let object; + let tokens = 0; + try { + const result = await generateObject({ + model, + schema: responseSchema, + prompt, + maxTokens: modelConfigs.queryRewriter.maxTokens + }); + object = result.object; + tokens = result.usage?.totalTokens || 0; + } catch (error) { + const result = await handleGenerateObjectError(error); + object = result.object; + tokens = result.totalTokens; + } + console.log('Query rewriter:', object.queries); (tracker || new TokenTracker()).trackUsage('query-rewriter', tokens); - - return { queries: json.queries, tokens }; + return { queries: object.queries, tokens }; } catch (error) { console.error('Error in query rewriting:', error); throw error; } -} \ No newline at end of file +} diff --git a/src/types.ts b/src/types.ts index b15c296..0e7db89 100644 --- a/src/types.ts +++ b/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; - required?: string[]; - }; - properties?: Record; - required?: string[]; - maxItems?: number; -}; - -export type ResponseSchema = { - type: SchemaType; - properties: Record; - required: string[]; -}; - export interface StreamMessage { type: 'progress' | 'answer' | 'error'; data: string | StepAction; diff --git a/src/utils/error-handling.ts b/src/utils/error-handling.ts new file mode 100644 index 0000000..aed77e0 --- /dev/null +++ b/src/utils/error-handling.ts @@ -0,0 +1,22 @@ +import {NoObjectGeneratedError} from "ai"; + +export interface GenerateObjectResult { + object: T; + totalTokens: number; +} + +export async function handleGenerateObjectError(error: unknown): Promise> { + 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; +}