├── .dockerignore ├── .eslintrc.js ├── .github ├── visuals │ ├── chatbox.gif │ ├── demo.gif │ ├── demo2.gif │ ├── demo3.gif │ ├── demo4.gif │ └── demo6.gif └── workflows │ ├── cd.yml │ ├── npm-publish.yml │ └── test.yml ├── .gitignore ├── .vscode └── launch.json ├── Dockerfile ├── LICENSE ├── README.md ├── config.json ├── docker-compose.yml ├── jest.config.js ├── jest.setup.js ├── jina-ai ├── .dockerignore ├── Dockerfile ├── config.json ├── package-lock.json ├── package.json ├── src │ ├── dto │ │ └── jina-embeddings-auth.ts │ ├── lib │ │ ├── async-context.ts │ │ ├── billing.ts │ │ ├── env-config.ts │ │ ├── errors.ts │ │ ├── firestore.ts │ │ ├── logger.ts │ │ └── registry.ts │ ├── patch-express.ts │ ├── rate-limit.ts │ └── server.ts └── tsconfig.json ├── package-lock.json ├── package.json ├── src ├── __tests__ │ ├── agent.test.ts │ ├── docker.test.ts │ └── server.test.ts ├── agent.ts ├── app.ts ├── cli.ts ├── config.ts ├── evals │ ├── batch-evals.ts │ └── ego-questions.json ├── server.ts ├── tools │ ├── __tests__ │ │ ├── error-analyzer.test.ts │ │ ├── evaluator.test.ts │ │ ├── read.test.ts │ │ └── search.test.ts │ ├── brave-search.ts │ ├── broken-ch-fixer.ts │ ├── build-ref.ts │ ├── code-sandbox.ts │ ├── cosine.ts │ ├── dedup.ts │ ├── embeddings.ts │ ├── error-analyzer.ts │ ├── evaluator.ts │ ├── grounding.ts │ ├── jina-classify-spam.ts │ ├── jina-dedup.ts │ ├── jina-latechunk.ts │ ├── jina-rerank.ts │ ├── jina-search.ts │ ├── md-fixer.ts │ ├── query-rewriter.ts │ ├── read.ts │ ├── segment.ts │ └── serper-search.ts ├── types.ts └── utils │ ├── action-tracker.ts │ ├── axios-client.ts │ ├── date-tools.ts │ ├── i18n.json │ ├── safe-generator.ts │ ├── schemas.ts │ ├── text-tools.ts │ ├── token-tracker.ts │ └── url-tools.ts └── tsconfig.json /.dockerignore: -------------------------------------------------------------------------------- 1 | node_modules -------------------------------------------------------------------------------- /.eslintrc.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | parser: '@typescript-eslint/parser', 3 | plugins: ['@typescript-eslint'], 4 | extends: [ 5 | 'eslint:recommended', 6 | 'plugin:@typescript-eslint/recommended' 7 | ], 8 | env: { 9 | node: true, 10 | es6: true 11 | }, 12 | parserOptions: { 13 | ecmaVersion: 2020, 14 | sourceType: 'module' 15 | }, 16 | rules: { 17 | 'no-console': ['error', { allow: ['log', 'error'] }], 18 | '@typescript-eslint/no-var-requires': 'off', 19 | '@typescript-eslint/no-explicit-any': 'off' 20 | }, 21 | ignorePatterns: ["jina-ai/**/*"] 22 | }; 23 | -------------------------------------------------------------------------------- /.github/visuals/chatbox.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jina-ai/node-DeepResearch/f1ef6c5cd02426cad38b5df45b7e4f53ea7435e5/.github/visuals/chatbox.gif -------------------------------------------------------------------------------- /.github/visuals/demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jina-ai/node-DeepResearch/f1ef6c5cd02426cad38b5df45b7e4f53ea7435e5/.github/visuals/demo.gif -------------------------------------------------------------------------------- /.github/visuals/demo2.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jina-ai/node-DeepResearch/f1ef6c5cd02426cad38b5df45b7e4f53ea7435e5/.github/visuals/demo2.gif -------------------------------------------------------------------------------- /.github/visuals/demo3.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jina-ai/node-DeepResearch/f1ef6c5cd02426cad38b5df45b7e4f53ea7435e5/.github/visuals/demo3.gif -------------------------------------------------------------------------------- /.github/visuals/demo4.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jina-ai/node-DeepResearch/f1ef6c5cd02426cad38b5df45b7e4f53ea7435e5/.github/visuals/demo4.gif -------------------------------------------------------------------------------- /.github/visuals/demo6.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jina-ai/node-DeepResearch/f1ef6c5cd02426cad38b5df45b7e4f53ea7435e5/.github/visuals/demo6.gif -------------------------------------------------------------------------------- /.github/workflows/cd.yml: -------------------------------------------------------------------------------- 1 | run-name: Build push and deploy (CD) 2 | on: 3 | push: 4 | branches: 5 | - main 6 | - ci-debug 7 | tags: 8 | - '*' 9 | 10 | jobs: 11 | build-and-push-to-gcr: 12 | runs-on: ubuntu-latest 13 | concurrency: 14 | group: ${{ github.ref_type == 'branch' && github.ref }} 15 | cancel-in-progress: true 16 | permissions: 17 | contents: read 18 | steps: 19 | - uses: actions/checkout@v4 20 | with: 21 | lfs: true 22 | - uses: 'google-github-actions/auth@v2' 23 | with: 24 | credentials_json: '${{ secrets.GCLOUD_SERVICE_ACCOUNT_SECRET_JSON }}' 25 | - name: 'Set up Cloud SDK' 26 | uses: 'google-github-actions/setup-gcloud@v2' 27 | - name: "Docker auth" 28 | run: |- 29 | gcloud auth configure-docker us-docker.pkg.dev --quiet 30 | - name: Set controller release version 31 | run: echo "RELEASE_VERSION=${GITHUB_REF#refs/*/}" >> $GITHUB_ENV 32 | - name: Set up Node.js 33 | uses: actions/setup-node@v4 34 | with: 35 | node-version: 22.12.0 36 | cache: npm 37 | 38 | - name: npm install 39 | run: npm ci 40 | - name: build application 41 | run: npm run build 42 | - name: Set package version 43 | run: npm version --no-git-tag-version ${{ env.RELEASE_VERSION }} 44 | if: github.ref_type == 'tag' 45 | - name: Docker meta 46 | id: meta 47 | uses: docker/metadata-action@v5 48 | with: 49 | images: | 50 | us-docker.pkg.dev/research-df067/deepresearch/node-deep-research 51 | - name: Set up QEMU 52 | uses: docker/setup-qemu-action@v3 53 | - name: Set up Docker Buildx 54 | uses: docker/setup-buildx-action@v3 55 | - name: Build and push 56 | id: container 57 | uses: docker/build-push-action@v6 58 | with: 59 | file: jina-ai/Dockerfile 60 | push: true 61 | tags: ${{ steps.meta.outputs.tags }} 62 | labels: ${{ steps.meta.outputs.labels }} 63 | - name: Deploy with Tag 64 | run: | 65 | gcloud run deploy node-deep-research --image us-docker.pkg.dev/research-df067/deepresearch/node-deep-research@${{steps.container.outputs.imageid}} --tag ${{ env.RELEASE_VERSION }} --region us-central1 --async 66 | -------------------------------------------------------------------------------- /.github/workflows/npm-publish.yml: -------------------------------------------------------------------------------- 1 | name: NPM Publish 2 | 3 | on: 4 | release: 5 | types: [created] 6 | 7 | jobs: 8 | publish: 9 | runs-on: ubuntu-latest 10 | steps: 11 | - uses: actions/checkout@v4 12 | 13 | - name: Use Node.js 14 | uses: actions/setup-node@v4 15 | with: 16 | node-version: '20.x' 17 | registry-url: 'https://registry.npmjs.org' 18 | cache: 'npm' 19 | 20 | - name: Install dependencies 21 | run: npm ci 22 | 23 | - name: Run lint 24 | env: 25 | BRAVE_API_KEY: ${{ secrets.BRAVE_API_KEY }} 26 | GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} 27 | JINA_API_KEY: ${{ secrets.JINA_API_KEY }} 28 | GOOGLE_API_KEY: ${{ secrets.GEMINI_API_KEY }} 29 | OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} 30 | run: npm run lint 31 | 32 | - name: Build TypeScript 33 | run: npm run build 34 | 35 | - name: Update version from release 36 | run: | 37 | # Get release tag without 'v' prefix 38 | VERSION=$(echo ${{ github.ref_name }} | sed 's/^v//') 39 | # Update version in package.json 40 | npm version $VERSION --no-git-tag-version --allow-same-version 41 | 42 | - name: Publish to npm 43 | env: 44 | NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }} 45 | run: npm publish --access public 46 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test 2 | 3 | on: 4 | pull_request: 5 | branches: [ main ] 6 | 7 | jobs: 8 | test: 9 | if: "!startsWith(github.event.head_commit.message, 'chore')" 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v4 13 | 14 | - name: Use Node.js 15 | uses: actions/setup-node@v4 16 | with: 17 | node-version: '20.x' 18 | cache: 'npm' 19 | 20 | - name: Install dependencies 21 | run: npm ci 22 | 23 | - name: Run lint 24 | run: npm run lint 25 | 26 | - name: Run tests 27 | env: 28 | BRAVE_API_KEY: ${{ secrets.BRAVE_API_KEY }} 29 | GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} 30 | JINA_API_KEY: ${{ secrets.JINA_API_KEY }} 31 | GOOGLE_API_KEY: ${{ secrets.GEMINI_API_KEY }} 32 | OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} 33 | run: npm test 34 | 35 | - name: Set up Docker 36 | uses: docker/setup-buildx-action@v3 37 | 38 | - name: Run Docker tests 39 | env: 40 | BRAVE_API_KEY: mock_key 41 | GEMINI_API_KEY: mock_key 42 | JINA_API_KEY: mock_key 43 | GOOGLE_API_KEY: mock_key 44 | OPENAI_API_KEY: mock_key 45 | run: npm run test:docker 46 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Files 2 | tasks/ 3 | context.json 4 | knowledge.json 5 | prompt-*.txt 6 | queries.json 7 | questions.json 8 | eval-*.json 9 | 10 | # Logs 11 | logs 12 | .idea 13 | *.log 14 | npm-debug.log* 15 | yarn-debug.log* 16 | yarn-error.log* 17 | lerna-debug.log* 18 | .pnpm-debug.log* 19 | 20 | # Diagnostic reports (https://nodejs.org/api/report.html) 21 | report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json 22 | 23 | # Runtime data 24 | pids 25 | *.pid 26 | *.seed 27 | *.pid.lock 28 | 29 | # Directory for instrumented libs generated by jscoverage/JSCover 30 | lib-cov 31 | 32 | # Coverage directory used by tools like istanbul 33 | coverage 34 | *.lcov 35 | 36 | # nyc test coverage 37 | .nyc_output 38 | 39 | # Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files) 40 | .grunt 41 | 42 | # Bower dependency directory (https://bower.io/) 43 | bower_components 44 | 45 | # node-waf configuration 46 | .lock-wscript 47 | 48 | # Compiled binary addons (https://nodejs.org/api/addons.html) 49 | build/Release 50 | 51 | # Dependency directories 52 | node_modules/ 53 | jspm_packages/ 54 | 55 | # Snowpack dependency directory (https://snowpack.dev/) 56 | web_modules/ 57 | 58 | # TypeScript cache 59 | *.tsbuildinfo 60 | 61 | # Optional npm cache directory 62 | .npm 63 | 64 | # Optional eslint cache 65 | .eslintcache 66 | 67 | # Optional stylelint cache 68 | .stylelintcache 69 | 70 | # Microbundle cache 71 | .rpt2_cache/ 72 | .rts2_cache_cjs/ 73 | .rts2_cache_es/ 74 | .rts2_cache_umd/ 75 | 76 | # Optional REPL history 77 | .node_repl_history 78 | 79 | # Output of 'npm pack' 80 | *.tgz 81 | 82 | # Yarn Integrity file 83 | .yarn-integrity 84 | 85 | # dotenv environment variable files 86 | .env 87 | .env.development.local 88 | .env.test.local 89 | .env.production.local 90 | .env.local 91 | 92 | # parcel-bundler cache (https://parceljs.org/) 93 | .cache 94 | .parcel-cache 95 | 96 | # Next.js build output 97 | .next 98 | out 99 | 100 | # Nuxt.js build / generate output 101 | .nuxt 102 | dist 103 | 104 | # Gatsby files 105 | .cache/ 106 | # Comment in the public line in if your project uses Gatsby and not Next.js 107 | # https://nextjs.org/blog/next-9-1#public-directory-support 108 | # public 109 | 110 | # vuepress build output 111 | .vuepress/dist 112 | 113 | # vuepress v2.x temp and cache directory 114 | .temp 115 | .cache 116 | 117 | # Docusaurus cache and generated files 118 | .docusaurus 119 | 120 | # Serverless directories 121 | .serverless/ 122 | 123 | # FuseBox cache 124 | .fusebox/ 125 | 126 | # DynamoDB Local files 127 | .dynamodb/ 128 | 129 | # TernJS port file 130 | .tern-port 131 | 132 | # Stores VSCode versions used for testing VSCode extensions 133 | .vscode-test 134 | 135 | # yarn v2 136 | .yarn/cache 137 | .yarn/unplugged 138 | .yarn/build-state.yml 139 | .yarn/install-state.gz 140 | .pnp.* 141 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | { 2 | "version": "0.2.0", 3 | "configurations": [ 4 | { 5 | "name": "Attach", 6 | "port": 9229, 7 | "request": "attach", 8 | "skipFiles": [ 9 | "/**" 10 | ], 11 | "type": "node" 12 | }, 13 | { 14 | "name": "Attach by Process ID", 15 | "processId": "${command:PickProcess}", 16 | "request": "attach", 17 | "skipFiles": [ 18 | "/**" 19 | ], 20 | "type": "node" 21 | } 22 | ] 23 | } 24 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # ---- BUILD STAGE ---- 2 | FROM node:20-slim AS builder 3 | 4 | # Set working directory 5 | WORKDIR /app 6 | 7 | # Copy package.json and package-lock.json 8 | COPY package*.json ./ 9 | 10 | # Install dependencies 11 | RUN npm install --ignore-scripts 12 | 13 | # Copy application code 14 | COPY . . 15 | 16 | # Build the application 17 | RUN npm run build --ignore-scripts 18 | 19 | # ---- PRODUCTION STAGE ---- 20 | FROM node:20-slim AS production 21 | 22 | # Set working directory 23 | WORKDIR /app 24 | 25 | # Copy package.json and package-lock.json 26 | COPY package*.json ./ 27 | 28 | # Install production dependencies only 29 | RUN npm install --production --ignore-scripts 30 | 31 | # Copy config.json and built files from builder 32 | COPY --from=builder /app/config.json ./ 33 | COPY --from=builder /app/dist ./dist 34 | 35 | # Set environment variables (Recommended to set at runtime, avoid hardcoding) 36 | ENV GEMINI_API_KEY=${GEMINI_API_KEY} 37 | ENV OPENAI_API_KEY=${OPENAI_API_KEY} 38 | ENV JINA_API_KEY=${JINA_API_KEY} 39 | ENV BRAVE_API_KEY=${BRAVE_API_KEY} 40 | 41 | # Expose the port the app runs on 42 | EXPOSE 3000 43 | 44 | # Set startup command 45 | CMD ["node", "./dist/server.js"] 46 | -------------------------------------------------------------------------------- /config.json: -------------------------------------------------------------------------------- 1 | { 2 | "env": { 3 | "https_proxy": "", 4 | "OPENAI_BASE_URL": "", 5 | "GEMINI_API_KEY": "", 6 | "OPENAI_API_KEY": "", 7 | "JINA_API_KEY": "", 8 | "BRAVE_API_KEY": "", 9 | "SERPER_API_KEY": "", 10 | "DEFAULT_MODEL_NAME": "" 11 | }, 12 | "defaults": { 13 | "search_provider": "jina", 14 | "llm_provider": "gemini", 15 | "step_sleep": 100 16 | }, 17 | "providers": { 18 | "gemini": { 19 | "createClient": "createGoogleGenerativeAI" 20 | }, 21 | "openai": { 22 | "createClient": "createOpenAI", 23 | "clientConfig": { 24 | "compatibility": "strict" 25 | } 26 | } 27 | }, 28 | "models": { 29 | "gemini": { 30 | "default": { 31 | "model": "gemini-2.0-flash", 32 | "temperature": 0, 33 | "maxTokens": 2000 34 | }, 35 | "tools": { 36 | "coder": { "temperature": 0.7 }, 37 | "searchGrounding": { "temperature": 0 }, 38 | "dedup": { "temperature": 0.1 }, 39 | "evaluator": {"temperature": 0.6, "maxTokens": 200}, 40 | "errorAnalyzer": {}, 41 | "queryRewriter": { "temperature": 0.1 }, 42 | "agent": { "temperature": 0.7 }, 43 | "agentBeastMode": { "temperature": 0.7 }, 44 | "fallback": {"maxTokens": 8000, "model": "gemini-2.0-flash-lite"} 45 | } 46 | }, 47 | "openai": { 48 | "default": { 49 | "model": "gpt-4o-mini", 50 | "temperature": 0, 51 | "maxTokens": 8000 52 | }, 53 | "tools": { 54 | "coder": { "temperature": 0.7 }, 55 | "searchGrounding": { "temperature": 0 }, 56 | "dedup": { "temperature": 0.1 }, 57 | "evaluator": {}, 58 | "errorAnalyzer": {}, 59 | "queryRewriter": { "temperature": 0.1 }, 60 | "agent": { "temperature": 0.7 }, 61 | "agentBeastMode": { "temperature": 0.7 }, 62 | "fallback": { "temperature": 0 } 63 | } 64 | } 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '3.8' 2 | 3 | services: 4 | app: 5 | build: 6 | context: . 7 | dockerfile: Dockerfile 8 | environment: 9 | - GEMINI_API_KEY=${GEMINI_API_KEY} 10 | - OPENAI_API_KEY=${OPENAI_API_KEY} 11 | - JINA_API_KEY=${JINA_API_KEY} 12 | - BRAVE_API_KEY=${BRAVE_API_KEY} 13 | ports: 14 | - "3000:3000" 15 | -------------------------------------------------------------------------------- /jest.config.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | preset: 'ts-jest', 3 | testEnvironment: 'node', 4 | testMatch: ['**/__tests__/**/*.test.ts'], 5 | setupFiles: ['/jest.setup.js'], 6 | }; 7 | -------------------------------------------------------------------------------- /jest.setup.js: -------------------------------------------------------------------------------- 1 | require('dotenv').config(); 2 | -------------------------------------------------------------------------------- /jina-ai/.dockerignore: -------------------------------------------------------------------------------- 1 | node_modules -------------------------------------------------------------------------------- /jina-ai/Dockerfile: -------------------------------------------------------------------------------- 1 | # ---- BUILD STAGE ---- 2 | FROM node:22-slim AS builder 3 | 4 | # Set working directory 5 | WORKDIR /app 6 | 7 | # Copy package.json and package-lock.json 8 | COPY ./package*.json ./ 9 | COPY ./jina-ai/package*.json ./jina-ai/ 10 | 11 | # Install dependencies 12 | RUN npm ci 13 | WORKDIR /app/jina-ai 14 | RUN npm ci 15 | 16 | WORKDIR /app 17 | 18 | # Copy application code 19 | COPY ./src ./src 20 | COPY ./tsconfig.json ./tsconfig.json 21 | COPY ./jina-ai/config.json ./ 22 | RUN npm run build 23 | 24 | COPY ./jina-ai/src ./jina-ai/src 25 | COPY ./jina-ai/tsconfig.json ./jina-ai/tsconfig.json 26 | WORKDIR /app/jina-ai 27 | RUN npm run build 28 | 29 | # ---- PRODUCTION STAGE ---- 30 | FROM node:22 AS production 31 | 32 | # Set working directory 33 | WORKDIR /app 34 | 35 | COPY --from=builder /app ./ 36 | # Copy config.json and built files from builder 37 | 38 | WORKDIR /app/jina-ai 39 | 40 | # Set environment variables (Recommended to set at runtime, avoid hardcoding) 41 | ENV GEMINI_API_KEY=${GEMINI_API_KEY} 42 | ENV OPENAI_API_KEY=${OPENAI_API_KEY} 43 | ENV JINA_API_KEY=${JINA_API_KEY} 44 | ENV BRAVE_API_KEY=${BRAVE_API_KEY} 45 | 46 | # Expose the port the app runs on 47 | EXPOSE 3000 48 | 49 | # Set startup command 50 | CMD ["node", "./dist/server.js"] 51 | -------------------------------------------------------------------------------- /jina-ai/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "env": { 3 | "https_proxy": "", 4 | "OPENAI_BASE_URL": "", 5 | "GEMINI_API_KEY": "", 6 | "OPENAI_API_KEY": "", 7 | "JINA_API_KEY": "", 8 | "BRAVE_API_KEY": "", 9 | "SERPER_API_KEY": "", 10 | "DEFAULT_MODEL_NAME": "" 11 | }, 12 | "defaults": { 13 | "search_provider": "jina", 14 | "llm_provider": "vertex", 15 | "step_sleep": 500 16 | }, 17 | "providers": { 18 | "vertex": { 19 | "createClient": "createGoogleVertex", 20 | "clientConfig": { 21 | "location": "us-central1" 22 | } 23 | }, 24 | "gemini": { 25 | "createClient": "createGoogleGenerativeAI" 26 | }, 27 | "openai": { 28 | "createClient": "createOpenAI", 29 | "clientConfig": { 30 | "compatibility": "strict" 31 | } 32 | } 33 | }, 34 | "models": { 35 | "gemini": { 36 | "default": { 37 | "model": "gemini-2.0-flash", 38 | "temperature": 0.6, 39 | "maxTokens": 8000 40 | }, 41 | "tools": { 42 | "coder": { "maxTokens": 2000, "model": "gemini-2.0-flash-lite" }, 43 | "searchGrounding": {}, 44 | "dedup": { }, 45 | "evaluator": {"maxTokens": 2000 }, 46 | "errorAnalyzer": {"maxTokens": 1000}, 47 | "queryRewriter": {"maxTokens": 2000}, 48 | "agent": {}, 49 | "agentBeastMode": {}, 50 | "fallback": {"maxTokens": 8000, "model": "gemini-2.0-flash-lite"} 51 | } 52 | }, 53 | "openai": { 54 | "default": { 55 | "model": "gpt-4o-mini", 56 | "temperature": 0, 57 | "maxTokens": 8000 58 | }, 59 | "tools": { 60 | "coder": { "temperature": 0.7 }, 61 | "searchGrounding": { "temperature": 0 }, 62 | "dedup": { "temperature": 0.1 }, 63 | "evaluator": {}, 64 | "errorAnalyzer": {}, 65 | "queryRewriter": { "temperature": 0.1 }, 66 | "agent": { "temperature": 0.7 }, 67 | "agentBeastMode": { "temperature": 0.7 }, 68 | "fallback": { "temperature": 0 } 69 | } 70 | } 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /jina-ai/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "@jina-ai/node-deepresearch", 3 | "version": "1.0.0", 4 | "main": "dist/app.js", 5 | "files": [ 6 | "dist", 7 | "README.md", 8 | "LICENSE" 9 | ], 10 | "scripts": { 11 | "build": "tsc", 12 | "dev": "npx ts-node src/agent.ts", 13 | "search": "npx ts-node src/test-duck.ts", 14 | "rewrite": "npx ts-node src/tools/query-rewriter.ts", 15 | "lint": "eslint . --ext .ts", 16 | "lint:fix": "eslint . --ext .ts --fix", 17 | "serve": "ts-node src/server.ts", 18 | "eval": "ts-node src/evals/batch-evals.ts", 19 | "test": "jest --testTimeout=30000", 20 | "test:watch": "jest --watch" 21 | }, 22 | "keywords": [], 23 | "author": "Jina AI", 24 | "license": "Apache-2.0", 25 | "description": "", 26 | "dependencies": { 27 | "@ai-sdk/google-vertex": "^2.1.12", 28 | "@google-cloud/firestore": "^7.11.0", 29 | "@google-cloud/storage": "^7.15.1", 30 | "civkit": "^0.8.3-15926cb", 31 | "dayjs": "^1.11.13", 32 | "lodash": "^4.17.21", 33 | "reflect-metadata": "^0.2.2", 34 | "tsyringe": "^4.8.0" 35 | }, 36 | "devDependencies": { 37 | "@types/lodash": "^4.17.15", 38 | "pino-pretty": "^13.0.0" 39 | } 40 | } 41 | -------------------------------------------------------------------------------- /jina-ai/src/dto/jina-embeddings-auth.ts: -------------------------------------------------------------------------------- 1 | import { 2 | Also, AuthenticationFailedError, AuthenticationRequiredError, 3 | DownstreamServiceFailureError, RPC_CALL_ENVIRONMENT, 4 | ArrayOf, AutoCastable, Prop 5 | } from 'civkit/civ-rpc'; 6 | import { parseJSONText } from 'civkit/vectorize'; 7 | import { htmlEscape } from 'civkit/escape'; 8 | import { marshalErrorLike } from 'civkit/lang'; 9 | 10 | import type express from 'express'; 11 | 12 | import logger from '../lib/logger'; 13 | import { AsyncLocalContext } from '../lib/async-context'; 14 | import { InjectProperty } from '../lib/registry'; 15 | import { JinaEmbeddingsDashboardHTTP } from '../lib/billing'; 16 | import envConfig from '../lib/env-config'; 17 | 18 | import { FirestoreRecord } from '../lib/firestore'; 19 | import _ from 'lodash'; 20 | import { RateLimitDesc } from '../rate-limit'; 21 | 22 | export class JinaWallet extends AutoCastable { 23 | @Prop({ 24 | default: '' 25 | }) 26 | user_id!: string; 27 | 28 | @Prop({ 29 | default: 0 30 | }) 31 | trial_balance!: number; 32 | 33 | @Prop() 34 | trial_start?: Date; 35 | 36 | @Prop() 37 | trial_end?: Date; 38 | 39 | @Prop({ 40 | default: 0 41 | }) 42 | regular_balance!: number; 43 | 44 | @Prop({ 45 | default: 0 46 | }) 47 | total_balance!: number; 48 | } 49 | 50 | export class JinaEmbeddingsTokenAccount extends FirestoreRecord { 51 | static override collectionName = 'embeddingsTokenAccounts'; 52 | 53 | override _id!: string; 54 | 55 | @Prop({ 56 | required: true 57 | }) 58 | user_id!: string; 59 | 60 | @Prop({ 61 | nullable: true, 62 | type: String, 63 | }) 64 | email?: string; 65 | 66 | @Prop({ 67 | nullable: true, 68 | type: String, 69 | }) 70 | full_name?: string; 71 | 72 | @Prop({ 73 | nullable: true, 74 | type: String, 75 | }) 76 | customer_id?: string; 77 | 78 | @Prop({ 79 | nullable: true, 80 | type: String, 81 | }) 82 | avatar_url?: string; 83 | 84 | // Not keeping sensitive info for now 85 | // @Prop() 86 | // billing_address?: object; 87 | 88 | // @Prop() 89 | // payment_method?: object; 90 | 91 | @Prop({ 92 | required: true 93 | }) 94 | wallet!: JinaWallet; 95 | 96 | @Prop({ 97 | type: Object 98 | }) 99 | metadata?: { [k: string]: any; }; 100 | 101 | @Prop({ 102 | defaultFactory: () => new Date() 103 | }) 104 | lastSyncedAt!: Date; 105 | 106 | @Prop({ 107 | dictOf: [ArrayOf(RateLimitDesc)] 108 | }) 109 | customRateLimits?: { [k: string]: RateLimitDesc[]; }; 110 | 111 | static patchedFields = [ 112 | ]; 113 | 114 | static override from(input: any) { 115 | for (const field of this.patchedFields) { 116 | if (typeof input[field] === 'string') { 117 | input[field] = parseJSONText(input[field]); 118 | } 119 | } 120 | 121 | return super.from(input) as JinaEmbeddingsTokenAccount; 122 | } 123 | 124 | override degradeForFireStore() { 125 | const copy: any = { 126 | ...this, 127 | wallet: { ...this.wallet }, 128 | // Firebase disability 129 | customRateLimits: _.mapValues(this.customRateLimits, (v) => v.map((x) => ({ ...x }))), 130 | }; 131 | 132 | for (const field of (this.constructor as typeof JinaEmbeddingsTokenAccount).patchedFields) { 133 | if (typeof copy[field] === 'object') { 134 | copy[field] = JSON.stringify(copy[field]) as any; 135 | } 136 | } 137 | 138 | return copy; 139 | } 140 | 141 | [k: string]: any; 142 | } 143 | 144 | 145 | const authDtoLogger = logger.child({ service: 'JinaAuthDTO' }); 146 | 147 | export interface FireBaseHTTPCtx { 148 | req: express.Request, 149 | res: express.Response, 150 | } 151 | 152 | const THE_VERY_SAME_JINA_EMBEDDINGS_CLIENT = new JinaEmbeddingsDashboardHTTP(envConfig.JINA_EMBEDDINGS_DASHBOARD_API_KEY); 153 | 154 | @Also({ 155 | openapi: { 156 | operation: { 157 | parameters: { 158 | 'Authorization': { 159 | description: htmlEscape`Jina Token for authentication.\n\n` + 160 | htmlEscape`- Member of \n\n` + 161 | `- Authorization: Bearer {YOUR_JINA_TOKEN}` 162 | , 163 | in: 'header', 164 | schema: { 165 | anyOf: [ 166 | { type: 'string', format: 'token' } 167 | ] 168 | } 169 | } 170 | } 171 | } 172 | } 173 | }) 174 | export class JinaEmbeddingsAuthDTO extends AutoCastable { 175 | uid?: string; 176 | bearerToken?: string; 177 | user?: JinaEmbeddingsTokenAccount; 178 | 179 | @InjectProperty(AsyncLocalContext) 180 | ctxMgr!: AsyncLocalContext; 181 | 182 | jinaEmbeddingsDashboard = THE_VERY_SAME_JINA_EMBEDDINGS_CLIENT; 183 | 184 | static override from(input: any) { 185 | const instance = super.from(input) as JinaEmbeddingsAuthDTO; 186 | 187 | const ctx = input[RPC_CALL_ENVIRONMENT]; 188 | 189 | const req = (ctx.rawRequest || ctx.req) as express.Request | undefined; 190 | 191 | if (req) { 192 | const authorization = req.get('authorization'); 193 | 194 | if (authorization) { 195 | const authToken = authorization.split(' ')[1] || authorization; 196 | instance.bearerToken = authToken; 197 | } 198 | 199 | } 200 | 201 | if (!instance.bearerToken && input._token) { 202 | instance.bearerToken = input._token; 203 | } 204 | 205 | return instance; 206 | } 207 | 208 | async getBrief(ignoreCache?: boolean | string) { 209 | if (!this.bearerToken) { 210 | throw new AuthenticationRequiredError({ 211 | message: 'Jina API key is required to authenticate. Please get one from https://jina.ai' 212 | }); 213 | } 214 | 215 | let account; 216 | try { 217 | account = await JinaEmbeddingsTokenAccount.fromFirestore(this.bearerToken); 218 | } catch (err) { 219 | // FireStore would not accept any string as input and may throw if not happy with it 220 | void 0; 221 | } 222 | 223 | 224 | const age = account?.lastSyncedAt ? Date.now() - account.lastSyncedAt.getTime() : Infinity; 225 | 226 | if (account && !ignoreCache) { 227 | if (account && age < 180_000) { 228 | this.user = account; 229 | this.uid = this.user?.user_id; 230 | 231 | return account; 232 | } 233 | } 234 | 235 | try { 236 | const r = await this.jinaEmbeddingsDashboard.validateToken(this.bearerToken); 237 | const brief = r.data; 238 | const draftAccount = JinaEmbeddingsTokenAccount.from({ 239 | ...account, ...brief, _id: this.bearerToken, 240 | lastSyncedAt: new Date() 241 | }); 242 | await JinaEmbeddingsTokenAccount.save(draftAccount.degradeForFireStore(), undefined, { merge: true }); 243 | 244 | this.user = draftAccount; 245 | this.uid = this.user?.user_id; 246 | 247 | return draftAccount; 248 | } catch (err: any) { 249 | authDtoLogger.warn(`Failed to get user brief: ${err}`, { err: marshalErrorLike(err) }); 250 | 251 | if (err?.status === 401) { 252 | throw new AuthenticationFailedError({ 253 | message: 'Invalid API key, please get a new one from https://jina.ai' 254 | }); 255 | } 256 | 257 | if (account) { 258 | this.user = account; 259 | this.uid = this.user?.user_id; 260 | 261 | return account; 262 | } 263 | 264 | 265 | throw new DownstreamServiceFailureError(`Failed to authenticate: ${err}`); 266 | } 267 | } 268 | 269 | async reportUsage(tokenCount: number, mdl: string, endpoint: string = '/encode') { 270 | const user = await this.assertUser(); 271 | const uid = user.user_id; 272 | user.wallet.total_balance -= tokenCount; 273 | 274 | return this.jinaEmbeddingsDashboard.reportUsage(this.bearerToken!, { 275 | model_name: mdl, 276 | api_endpoint: endpoint, 277 | consumer: { 278 | id: uid, 279 | user_id: uid, 280 | }, 281 | usage: { 282 | total_tokens: tokenCount 283 | }, 284 | labels: { 285 | model_name: mdl 286 | } 287 | }).then((r) => { 288 | JinaEmbeddingsTokenAccount.COLLECTION.doc(this.bearerToken!) 289 | .update({ 'wallet.total_balance': JinaEmbeddingsTokenAccount.OPS.increment(-tokenCount) }) 290 | .catch((err) => { 291 | authDtoLogger.warn(`Failed to update cache for ${uid}: ${err}`, { err: marshalErrorLike(err) }); 292 | }); 293 | 294 | return r; 295 | }).catch((err) => { 296 | user.wallet.total_balance += tokenCount; 297 | authDtoLogger.warn(`Failed to report usage for ${uid}: ${err}`, { err: marshalErrorLike(err) }); 298 | }); 299 | } 300 | 301 | async solveUID() { 302 | if (this.uid) { 303 | this.ctxMgr.set('uid', this.uid); 304 | 305 | return this.uid; 306 | } 307 | 308 | if (this.bearerToken) { 309 | await this.getBrief(); 310 | this.ctxMgr.set('uid', this.uid); 311 | 312 | return this.uid; 313 | } 314 | 315 | return undefined; 316 | } 317 | 318 | async assertUID() { 319 | const uid = await this.solveUID(); 320 | 321 | if (!uid) { 322 | throw new AuthenticationRequiredError('Authentication failed'); 323 | } 324 | 325 | return uid; 326 | } 327 | 328 | async assertUser() { 329 | if (this.user) { 330 | return this.user; 331 | } 332 | 333 | await this.getBrief(); 334 | 335 | return this.user!; 336 | } 337 | 338 | getRateLimits(...tags: string[]) { 339 | const descs = tags.map((x) => this.user?.customRateLimits?.[x] || []).flat().filter((x) => x.isEffective()); 340 | 341 | if (descs.length) { 342 | return descs; 343 | } 344 | 345 | return undefined; 346 | } 347 | } 348 | -------------------------------------------------------------------------------- /jina-ai/src/lib/async-context.ts: -------------------------------------------------------------------------------- 1 | import { GlobalAsyncContext } from 'civkit/async-context'; 2 | import { container, singleton } from 'tsyringe'; 3 | 4 | @singleton() 5 | export class AsyncLocalContext extends GlobalAsyncContext {} 6 | 7 | const instance = container.resolve(AsyncLocalContext); 8 | Reflect.set(process, 'asyncLocalContext', instance); 9 | export default instance; 10 | -------------------------------------------------------------------------------- /jina-ai/src/lib/billing.ts: -------------------------------------------------------------------------------- 1 | import { HTTPService } from 'civkit'; 2 | import _ from 'lodash'; 3 | 4 | 5 | export interface JinaWallet { 6 | trial_balance: number; 7 | trial_start: Date; 8 | trial_end: Date; 9 | regular_balance: number; 10 | total_balance: number; 11 | } 12 | 13 | 14 | export interface JinaUserBrief { 15 | user_id: string; 16 | email: string | null; 17 | full_name: string | null; 18 | customer_id: string | null; 19 | avatar_url?: string; 20 | billing_address: Partial<{ 21 | address: string; 22 | city: string; 23 | state: string; 24 | country: string; 25 | postal_code: string; 26 | }>; 27 | payment_method: Partial<{ 28 | brand: string; 29 | last4: string; 30 | exp_month: number; 31 | exp_year: number; 32 | }>; 33 | wallet: JinaWallet; 34 | metadata: { 35 | [k: string]: any; 36 | }; 37 | } 38 | 39 | export interface JinaUsageReport { 40 | model_name: string; 41 | api_endpoint: string; 42 | consumer: { 43 | user_id: string; 44 | customer_plan?: string; 45 | [k: string]: any; 46 | }; 47 | usage: { 48 | total_tokens: number; 49 | }; 50 | labels: { 51 | user_type?: string; 52 | model_name?: string; 53 | [k: string]: any; 54 | }; 55 | } 56 | 57 | export class JinaEmbeddingsDashboardHTTP extends HTTPService { 58 | name = 'JinaEmbeddingsDashboardHTTP'; 59 | 60 | constructor( 61 | public apiKey: string, 62 | public baseUri: string = 'https://embeddings-dashboard-api.jina.ai/api' 63 | ) { 64 | super(baseUri); 65 | 66 | this.baseOptions.timeout = 30_000; // 30 sec 67 | } 68 | 69 | async authorization(token: string) { 70 | const r = await this.get('/v1/authorization', { 71 | headers: { 72 | Authorization: `Bearer ${token}` 73 | }, 74 | responseType: 'json', 75 | }); 76 | 77 | return r; 78 | } 79 | 80 | async validateToken(token: string) { 81 | const r = await this.getWithSearchParams('/v1/api_key/user', { 82 | api_key: token, 83 | }, { 84 | responseType: 'json', 85 | }); 86 | 87 | return r; 88 | } 89 | 90 | async reportUsage(token: string, query: JinaUsageReport) { 91 | const r = await this.postJson('/v1/usage', query, { 92 | headers: { 93 | Authorization: `Bearer ${token}`, 94 | 'x-api-key': this.apiKey, 95 | }, 96 | responseType: 'text', 97 | }); 98 | 99 | return r; 100 | } 101 | 102 | } 103 | -------------------------------------------------------------------------------- /jina-ai/src/lib/env-config.ts: -------------------------------------------------------------------------------- 1 | import { container, singleton } from 'tsyringe'; 2 | 3 | export const SPECIAL_COMBINED_ENV_KEY = 'ENV_COMBINED'; 4 | const CONF_ENV = [ 5 | 'OPENAI_API_KEY', 6 | 7 | 'ANTHROPIC_API_KEY', 8 | 9 | 'REPLICATE_API_KEY', 10 | 11 | 'GOOGLE_AI_STUDIO_API_KEY', 12 | 13 | 'JINA_EMBEDDINGS_API_KEY', 14 | 15 | 'JINA_EMBEDDINGS_DASHBOARD_API_KEY', 16 | 17 | 'BRAVE_SEARCH_API_KEY', 18 | 19 | ] as const; 20 | 21 | 22 | @singleton() 23 | export class EnvConfig { 24 | dynamic!: Record; 25 | 26 | combined: Record = {}; 27 | originalEnv: Record = { ...process.env }; 28 | 29 | constructor() { 30 | if (process.env[SPECIAL_COMBINED_ENV_KEY]) { 31 | Object.assign(this.combined, JSON.parse( 32 | Buffer.from(process.env[SPECIAL_COMBINED_ENV_KEY]!, 'base64').toString('utf-8') 33 | )); 34 | delete process.env[SPECIAL_COMBINED_ENV_KEY]; 35 | } 36 | 37 | // Static config 38 | for (const x of CONF_ENV) { 39 | const s = this.combined[x] || process.env[x] || ''; 40 | Reflect.set(this, x, s); 41 | if (x in process.env) { 42 | delete process.env[x]; 43 | } 44 | } 45 | 46 | // Dynamic config 47 | this.dynamic = new Proxy({ 48 | get: (_target: any, prop: string) => { 49 | return this.combined[prop] || process.env[prop] || ''; 50 | } 51 | }, {}) as any; 52 | } 53 | } 54 | 55 | // eslint-disable-next-line @typescript-eslint/no-empty-interface 56 | export interface EnvConfig extends Record { } 57 | 58 | const instance = container.resolve(EnvConfig); 59 | export default instance; 60 | -------------------------------------------------------------------------------- /jina-ai/src/lib/errors.ts: -------------------------------------------------------------------------------- 1 | import { ApplicationError, Prop, RPC_TRANSFER_PROTOCOL_META_SYMBOL, StatusCode } from 'civkit'; 2 | import _ from 'lodash'; 3 | import dayjs from 'dayjs'; 4 | import utc from 'dayjs/plugin/utc'; 5 | 6 | dayjs.extend(utc); 7 | 8 | @StatusCode(50301) 9 | export class ServiceDisabledError extends ApplicationError { } 10 | 11 | @StatusCode(50302) 12 | export class ServiceCrashedError extends ApplicationError { } 13 | 14 | @StatusCode(50303) 15 | export class ServiceNodeResourceDrainError extends ApplicationError { } 16 | 17 | @StatusCode(40104) 18 | export class EmailUnverifiedError extends ApplicationError { } 19 | 20 | @StatusCode(40201) 21 | export class InsufficientCreditsError extends ApplicationError { } 22 | 23 | @StatusCode(40202) 24 | export class FreeFeatureLimitError extends ApplicationError { } 25 | 26 | @StatusCode(40203) 27 | export class InsufficientBalanceError extends ApplicationError { } 28 | 29 | @StatusCode(40903) 30 | export class LockConflictError extends ApplicationError { } 31 | 32 | @StatusCode(40904) 33 | export class BudgetExceededError extends ApplicationError { } 34 | 35 | @StatusCode(45101) 36 | export class HarmfulContentError extends ApplicationError { } 37 | 38 | @StatusCode(45102) 39 | export class SecurityCompromiseError extends ApplicationError { } 40 | 41 | @StatusCode(41201) 42 | export class BatchSizeTooLargeError extends ApplicationError { } 43 | 44 | 45 | @StatusCode(42903) 46 | export class RateLimitTriggeredError extends ApplicationError { 47 | 48 | @Prop({ 49 | desc: 'Retry after seconds', 50 | }) 51 | retryAfter?: number; 52 | 53 | @Prop({ 54 | desc: 'Retry after date', 55 | }) 56 | retryAfterDate?: Date; 57 | 58 | protected override get [RPC_TRANSFER_PROTOCOL_META_SYMBOL]() { 59 | const retryAfter = this.retryAfter || this.retryAfterDate; 60 | if (!retryAfter) { 61 | return super[RPC_TRANSFER_PROTOCOL_META_SYMBOL]; 62 | } 63 | 64 | return _.merge(_.cloneDeep(super[RPC_TRANSFER_PROTOCOL_META_SYMBOL]), { 65 | headers: { 66 | 'Retry-After': `${retryAfter instanceof Date ? dayjs(retryAfter).utc().format('ddd, DD MMM YYYY HH:mm:ss [GMT]') : retryAfter}`, 67 | } 68 | }); 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /jina-ai/src/lib/firestore.ts: -------------------------------------------------------------------------------- 1 | import _ from 'lodash'; 2 | import { AutoCastable, Prop, RPC_MARSHAL } from 'civkit/civ-rpc'; 3 | import { 4 | Firestore, FieldValue, DocumentReference, 5 | Query, Timestamp, SetOptions, DocumentSnapshot, 6 | } from '@google-cloud/firestore'; 7 | import { Storage } from '@google-cloud/storage'; 8 | 9 | export const firebaseDefaultBucket = new Storage().bucket(`${process.env.GCLOUD_PROJECT}.appspot.com`); 10 | 11 | // Firestore doesn't support JavaScript objects with custom prototypes (i.e. objects that were created via the \"new\" operator) 12 | function patchFireStoreArrogance(func: Function) { 13 | return function (this: unknown) { 14 | const origObjectGetPrototype = Object.getPrototypeOf; 15 | Object.getPrototypeOf = function (x) { 16 | const r = origObjectGetPrototype.call(this, x); 17 | if (!r) { 18 | return r; 19 | } 20 | return Object.prototype; 21 | }; 22 | try { 23 | return func.call(this, ...arguments); 24 | } finally { 25 | Object.getPrototypeOf = origObjectGetPrototype; 26 | } 27 | }; 28 | } 29 | 30 | Reflect.set(DocumentReference.prototype, 'set', patchFireStoreArrogance(Reflect.get(DocumentReference.prototype, 'set'))); 31 | Reflect.set(DocumentSnapshot, 'fromObject', patchFireStoreArrogance(Reflect.get(DocumentSnapshot, 'fromObject'))); 32 | 33 | function mapValuesDeep(v: any, fn: (i: any) => any): any { 34 | if (_.isPlainObject(v)) { 35 | return _.mapValues(v, (i) => mapValuesDeep(i, fn)); 36 | } else if (_.isArray(v)) { 37 | return v.map((i) => mapValuesDeep(i, fn)); 38 | } else { 39 | return fn(v); 40 | } 41 | } 42 | 43 | export type Constructor = { new(...args: any[]): T; }; 44 | export type Constructed = T extends Partial ? U : T extends object ? T : object; 45 | 46 | export function fromFirestore( 47 | this: Constructor, id: string, overrideCollection?: string 48 | ): Promise; 49 | export async function fromFirestore( 50 | this: any, id: string, overrideCollection?: string 51 | ) { 52 | const collection = overrideCollection || this.collectionName; 53 | if (!collection) { 54 | throw new Error(`Missing collection name to construct ${this.name}`); 55 | } 56 | 57 | const ref = this.DB.collection(overrideCollection || this.collectionName).doc(id); 58 | 59 | const ptr = await ref.get(); 60 | 61 | if (!ptr.exists) { 62 | return undefined; 63 | } 64 | 65 | const doc = this.from( 66 | // Fixes non-native firebase types 67 | mapValuesDeep(ptr.data(), (i: any) => { 68 | if (i instanceof Timestamp) { 69 | return i.toDate(); 70 | } 71 | 72 | return i; 73 | }) 74 | ); 75 | 76 | Object.defineProperty(doc, '_ref', { value: ref, enumerable: false }); 77 | Object.defineProperty(doc, '_id', { value: ptr.id, enumerable: true }); 78 | 79 | return doc; 80 | } 81 | 82 | export function fromFirestoreQuery( 83 | this: Constructor, query: Query 84 | ): Promise; 85 | export async function fromFirestoreQuery(this: any, query: Query) { 86 | const ptr = await query.get(); 87 | 88 | if (ptr.docs.length) { 89 | return ptr.docs.map(doc => { 90 | const r = this.from( 91 | mapValuesDeep(doc.data(), (i: any) => { 92 | if (i instanceof Timestamp) { 93 | return i.toDate(); 94 | } 95 | 96 | return i; 97 | }) 98 | ); 99 | Object.defineProperty(r, '_ref', { value: doc.ref, enumerable: false }); 100 | Object.defineProperty(r, '_id', { value: doc.id, enumerable: true }); 101 | 102 | return r; 103 | }); 104 | } 105 | 106 | return []; 107 | } 108 | 109 | export function setToFirestore( 110 | this: Constructor, doc: T, overrideCollection?: string, setOptions?: SetOptions 111 | ): Promise; 112 | export async function setToFirestore( 113 | this: any, doc: any, overrideCollection?: string, setOptions?: SetOptions 114 | ) { 115 | let ref: DocumentReference = doc._ref; 116 | if (!ref) { 117 | const collection = overrideCollection || this.collectionName; 118 | if (!collection) { 119 | throw new Error(`Missing collection name to construct ${this.name}`); 120 | } 121 | 122 | const predefinedId = doc._id || undefined; 123 | const hdl = this.DB.collection(overrideCollection || this.collectionName); 124 | ref = predefinedId ? hdl.doc(predefinedId) : hdl.doc(); 125 | 126 | Object.defineProperty(doc, '_ref', { value: ref, enumerable: false }); 127 | Object.defineProperty(doc, '_id', { value: ref.id, enumerable: true }); 128 | } 129 | 130 | await ref.set(doc, { merge: true, ...setOptions }); 131 | 132 | return doc; 133 | } 134 | 135 | export function deleteQueryBatch( 136 | this: Constructor, query: Query 137 | ): Promise; 138 | export async function deleteQueryBatch(this: any, query: Query) { 139 | const snapshot = await query.get(); 140 | 141 | const batchSize = snapshot.size; 142 | if (batchSize === 0) { 143 | return; 144 | } 145 | 146 | // Delete documents in a batch 147 | const batch = this.DB.batch(); 148 | snapshot.docs.forEach((doc) => { 149 | batch.delete(doc.ref); 150 | }); 151 | await batch.commit(); 152 | 153 | process.nextTick(() => { 154 | this.deleteQueryBatch(query); 155 | }); 156 | }; 157 | 158 | export function fromFirestoreDoc( 159 | this: Constructor, snapshot: DocumentSnapshot, 160 | ): T | undefined; 161 | export function fromFirestoreDoc( 162 | this: any, snapshot: DocumentSnapshot, 163 | ) { 164 | const doc = this.from( 165 | // Fixes non-native firebase types 166 | mapValuesDeep(snapshot.data(), (i: any) => { 167 | if (i instanceof Timestamp) { 168 | return i.toDate(); 169 | } 170 | 171 | return i; 172 | }) 173 | ); 174 | 175 | Object.defineProperty(doc, '_ref', { value: snapshot.ref, enumerable: false }); 176 | Object.defineProperty(doc, '_id', { value: snapshot.id, enumerable: true }); 177 | 178 | return doc; 179 | } 180 | const defaultFireStore = new Firestore({ 181 | projectId: process.env.GCLOUD_PROJECT, 182 | }); 183 | export class FirestoreRecord extends AutoCastable { 184 | static collectionName?: string; 185 | static OPS = FieldValue; 186 | static DB = defaultFireStore; 187 | static get COLLECTION() { 188 | if (!this.collectionName) { 189 | throw new Error('Not implemented'); 190 | } 191 | 192 | return this.DB.collection(this.collectionName); 193 | } 194 | 195 | @Prop() 196 | _id?: string; 197 | _ref?: DocumentReference>>; 198 | 199 | static fromFirestore = fromFirestore; 200 | static fromFirestoreDoc = fromFirestoreDoc; 201 | static fromFirestoreQuery = fromFirestoreQuery; 202 | 203 | static save = setToFirestore; 204 | static deleteQueryBatch = deleteQueryBatch; 205 | 206 | [RPC_MARSHAL]() { 207 | return { 208 | ...this, 209 | _id: this._id, 210 | _ref: this._ref?.path 211 | }; 212 | } 213 | 214 | degradeForFireStore(): this { 215 | return JSON.parse(JSON.stringify(this, function (k, v) { 216 | if (k === '') { 217 | return v; 218 | } 219 | if (typeof v === 'object' && v && (typeof v.degradeForFireStore === 'function')) { 220 | return v.degradeForFireStore(); 221 | } 222 | 223 | return v; 224 | })); 225 | } 226 | } 227 | -------------------------------------------------------------------------------- /jina-ai/src/lib/logger.ts: -------------------------------------------------------------------------------- 1 | import { AbstractPinoLogger } from 'civkit/pino-logger'; 2 | import { singleton, container } from 'tsyringe'; 3 | import { threadId } from 'node:worker_threads'; 4 | import { getTraceCtx } from 'civkit/async-context'; 5 | 6 | 7 | const levelToSeverityMap: { [k: string]: string | undefined; } = { 8 | trace: 'DEFAULT', 9 | debug: 'DEBUG', 10 | info: 'INFO', 11 | warn: 'WARNING', 12 | error: 'ERROR', 13 | fatal: 'CRITICAL', 14 | }; 15 | 16 | @singleton() 17 | export class GlobalLogger extends AbstractPinoLogger { 18 | loggerOptions = { 19 | level: 'debug', 20 | base: { 21 | tid: threadId, 22 | } 23 | }; 24 | 25 | override init(): void { 26 | if (process.env['NODE_ENV']?.startsWith('prod')) { 27 | super.init(process.stdout); 28 | } else { 29 | const PinoPretty = require('pino-pretty').PinoPretty; 30 | super.init(PinoPretty({ 31 | singleLine: true, 32 | colorize: true, 33 | messageFormat(log: any, messageKey: any) { 34 | return `${log['tid'] ? `[${log['tid']}]` : ''}[${log['service'] || 'ROOT'}] ${log[messageKey]}`; 35 | }, 36 | })); 37 | } 38 | 39 | 40 | this.emit('ready'); 41 | } 42 | 43 | override log(...args: any[]) { 44 | const [levelObj, ...rest] = args; 45 | const severity = levelToSeverityMap[levelObj?.level]; 46 | const traceCtx = getTraceCtx(); 47 | const patched: any= { ...levelObj, severity }; 48 | if (traceCtx?.traceId && process.env['GCLOUD_PROJECT']) { 49 | patched['logging.googleapis.com/trace'] = `projects/${process.env['GCLOUD_PROJECT']}/traces/${traceCtx.traceId}`; 50 | } 51 | return super.log(patched, ...rest); 52 | } 53 | } 54 | 55 | const instance = container.resolve(GlobalLogger); 56 | export default instance; 57 | -------------------------------------------------------------------------------- /jina-ai/src/lib/registry.ts: -------------------------------------------------------------------------------- 1 | import { container } from 'tsyringe'; 2 | import { propertyInjectorFactory } from 'civkit/property-injector'; 3 | 4 | export const InjectProperty = propertyInjectorFactory(container); -------------------------------------------------------------------------------- /jina-ai/src/patch-express.ts: -------------------------------------------------------------------------------- 1 | import { ApplicationError, OperationNotAllowedError, Prop, RPC_CALL_ENVIRONMENT } from "civkit/civ-rpc"; 2 | import { marshalErrorLike } from "civkit/lang"; 3 | import { randomUUID } from "crypto"; 4 | import { once } from "events"; 5 | import type { NextFunction, Request, Response } from "express"; 6 | 7 | import { JinaEmbeddingsAuthDTO } from "./dto/jina-embeddings-auth"; 8 | import rateLimitControl, { API_CALL_STATUS, APICall, RateLimitDesc } from "./rate-limit"; 9 | import asyncLocalContext from "./lib/async-context"; 10 | import globalLogger from "./lib/logger"; 11 | import { InsufficientBalanceError } from "./lib/errors"; 12 | import { firebaseDefaultBucket, FirestoreRecord } from "./lib/firestore"; 13 | import cors from "cors"; 14 | 15 | globalLogger.serviceReady(); 16 | const logger = globalLogger.child({ service: 'JinaAISaaSMiddleware' }); 17 | const appName = 'DEEPRESEARCH'; 18 | 19 | export class KnowledgeItem extends FirestoreRecord { 20 | static override collectionName = 'knowledgeItems'; 21 | 22 | @Prop({ 23 | required: true 24 | }) 25 | traceId!: string; 26 | 27 | @Prop({ 28 | required: true 29 | }) 30 | uid!: string; 31 | 32 | @Prop({ 33 | default: '' 34 | }) 35 | question!: string; 36 | 37 | @Prop({ 38 | default: '' 39 | }) 40 | answer!: string; 41 | 42 | @Prop({ 43 | default: '' 44 | }) 45 | type!: string; 46 | 47 | @Prop({ 48 | arrayOf: Object, 49 | default: [] 50 | }) 51 | references!: any[]; 52 | 53 | @Prop({ 54 | defaultFactory: () => new Date() 55 | }) 56 | createdAt!: Date; 57 | 58 | @Prop({ 59 | defaultFactory: () => new Date() 60 | }) 61 | updatedAt!: Date; 62 | } 63 | const corsMiddleware = cors(); 64 | export const jinaAiMiddleware = (req: Request, res: Response, next: NextFunction) => { 65 | if (req.path === '/ping') { 66 | res.status(200).end('pone'); 67 | return; 68 | } 69 | if (req.path.startsWith('/v1/models')) { 70 | next(); 71 | return; 72 | } 73 | if (req.method !== 'POST' && req.method !== 'GET') { 74 | next(); 75 | return; 76 | } 77 | 78 | // Early API key validation - reject immediately if no valid auth header 79 | const authHeader = req.headers.authorization; 80 | if (!authHeader || !authHeader.startsWith('Bearer ')) { 81 | corsMiddleware(req, res, () => { 82 | res.status(401).json({ error: 'Unauthorized: API key required' }); 83 | }); 84 | return; 85 | } 86 | 87 | asyncLocalContext.run(async () => { 88 | const googleTraceId = req.get('x-cloud-trace-context')?.split('/')?.[0]; 89 | const ctx = asyncLocalContext.ctx; 90 | ctx.traceId = req.get('x-request-id') || req.get('request-id') || googleTraceId || randomUUID(); 91 | ctx.traceT0 = new Date(); 92 | ctx.ip = req?.ip; 93 | 94 | try { 95 | const authDto = JinaEmbeddingsAuthDTO.from({ 96 | [RPC_CALL_ENVIRONMENT]: { req, res } 97 | }); 98 | 99 | const uid = await authDto.assertUID(); 100 | // if (!uid && !ctx.ip) { 101 | // throw new OperationNotAllowedError(`Missing IP information for anonymous user`); 102 | // } 103 | let rateLimitPolicy 104 | if (uid) { 105 | const user = await authDto.assertUser(); 106 | if (!(user.wallet.total_balance > 0)) { 107 | throw new InsufficientBalanceError(`Account balance not enough to run this query, please recharge.`); 108 | } 109 | rateLimitPolicy = authDto.getRateLimits(appName) || [ 110 | parseInt(user.metadata?.speed_level) >= 2 ? 111 | RateLimitDesc.from({ 112 | occurrence: 500, 113 | periodSeconds: 60 114 | }) : 115 | RateLimitDesc.from({ 116 | occurrence: 50, 117 | periodSeconds: 60 118 | }) 119 | ]; 120 | } else { 121 | rateLimitPolicy = [ 122 | RateLimitDesc.from({ 123 | occurrence: 0, 124 | periodSeconds: 120 125 | }) 126 | ] 127 | } 128 | 129 | const criterions = rateLimitPolicy.map((c) => rateLimitControl.rateLimitDescToCriterion(c)); 130 | await Promise.all( 131 | criterions.map( 132 | ([pointInTime, n]) => uid ? 133 | rateLimitControl.assertUidPeriodicLimit(uid, pointInTime, n, appName) : 134 | rateLimitControl.assertIPPeriodicLimit(ctx.ip!, pointInTime, n, appName) 135 | ) 136 | ); 137 | const draftApiCall: Partial = { tags: [appName] }; 138 | if (uid) { 139 | draftApiCall.uid = uid; 140 | } else { 141 | draftApiCall.ip = ctx.ip; 142 | } 143 | 144 | const apiRoll = rateLimitControl.record(draftApiCall); 145 | apiRoll.save().catch((err) => logger.warn(`Failed to save rate limit record`, { err: marshalErrorLike(err) })); 146 | 147 | const pResClose = once(res, 'close'); 148 | 149 | next(); 150 | 151 | await pResClose; 152 | const chargeAmount = ctx.chargeAmount; 153 | if (chargeAmount) { 154 | authDto.reportUsage(chargeAmount, `reader-${appName}`).catch((err) => { 155 | logger.warn(`Unable to report usage for ${uid || ctx.ip}`, { err: marshalErrorLike(err) }); 156 | }); 157 | apiRoll.chargeAmount = chargeAmount; 158 | } 159 | apiRoll.status = res.statusCode === 200 ? API_CALL_STATUS.SUCCESS : API_CALL_STATUS.ERROR; 160 | apiRoll.save().catch((err) => logger.warn(`Failed to save rate limit record`, { err: marshalErrorLike(err) })); 161 | logger.info(`HTTP ${res.statusCode} for request ${ctx.traceId} after ${Date.now() - ctx.traceT0.valueOf()}ms`, { 162 | uid, 163 | ip: ctx.ip, 164 | chargeAmount, 165 | }); 166 | 167 | if (uid && ctx.promptContext?.knowledge?.length) { 168 | Promise.all(ctx.promptContext.knowledge.map((x: any) => KnowledgeItem.save( 169 | KnowledgeItem.from({ 170 | ...x, 171 | uid, 172 | traceId: ctx.traceId, 173 | }) 174 | ))).catch((err: any) => { 175 | logger.warn(`Failed to save knowledge`, { err: marshalErrorLike(err) }); 176 | }); 177 | } 178 | if (ctx.promptContext) { 179 | const patchedCtx = { ...ctx.promptContext }; 180 | if (Array.isArray(patchedCtx.context)) { 181 | patchedCtx.context = patchedCtx.context.map((x: object) => ({ ...x, result: undefined })) 182 | } 183 | 184 | let data; 185 | try { 186 | data = JSON.stringify(patchedCtx); 187 | } catch (err: any) { 188 | const obj = marshalErrorLike(err); 189 | if (err.stack) { 190 | obj.stack = err.stack; 191 | } 192 | data = JSON.stringify(obj); 193 | logger.warn(`Failed to stringify promptContext`, { err: obj }); 194 | } 195 | 196 | firebaseDefaultBucket.file(`promptContext/${ctx.traceId}.json`).save( 197 | data, 198 | { 199 | metadata: { 200 | contentType: 'application/json', 201 | }, 202 | } 203 | ).catch((err: any) => { 204 | logger.warn(`Failed to save promptContext`, { err: marshalErrorLike(err) }); 205 | }).finally(() => { 206 | ctx.promptContext = undefined; 207 | }); 208 | } 209 | 210 | } catch (err: any) { 211 | if (!res.headersSent) { 212 | corsMiddleware(req, res, () => 'noop'); 213 | if (err instanceof ApplicationError) { 214 | res.status(parseInt(err.code as string) || 500).json({ error: err.message }); 215 | 216 | return; 217 | } 218 | 219 | res.status(500).json({ error: 'Internal' }); 220 | } 221 | 222 | logger.error(`Error in billing middleware`, { err: marshalErrorLike(err) }); 223 | if (err.stack) { 224 | logger.error(err.stack); 225 | } 226 | } 227 | 228 | }); 229 | } -------------------------------------------------------------------------------- /jina-ai/src/rate-limit.ts: -------------------------------------------------------------------------------- 1 | import { AutoCastable, ResourcePolicyDenyError, Also, Prop } from 'civkit/civ-rpc'; 2 | import { AsyncService } from 'civkit/async-service'; 3 | import { getTraceId } from 'civkit/async-context'; 4 | import { singleton, container } from 'tsyringe'; 5 | 6 | import { RateLimitTriggeredError } from './lib/errors'; 7 | import { FirestoreRecord } from './lib/firestore'; 8 | import { GlobalLogger } from './lib/logger'; 9 | 10 | export enum API_CALL_STATUS { 11 | SUCCESS = 'success', 12 | ERROR = 'error', 13 | PENDING = 'pending', 14 | } 15 | 16 | @Also({ dictOf: Object }) 17 | export class APICall extends FirestoreRecord { 18 | static override collectionName = 'apiRoll'; 19 | 20 | @Prop({ 21 | required: true, 22 | defaultFactory: () => getTraceId() 23 | }) 24 | traceId!: string; 25 | 26 | @Prop() 27 | uid?: string; 28 | 29 | @Prop() 30 | ip?: string; 31 | 32 | @Prop({ 33 | arrayOf: String, 34 | default: [], 35 | }) 36 | tags!: string[]; 37 | 38 | @Prop({ 39 | required: true, 40 | defaultFactory: () => new Date(), 41 | }) 42 | createdAt!: Date; 43 | 44 | @Prop() 45 | completedAt?: Date; 46 | 47 | @Prop({ 48 | required: true, 49 | default: API_CALL_STATUS.PENDING, 50 | }) 51 | status!: API_CALL_STATUS; 52 | 53 | @Prop({ 54 | required: true, 55 | defaultFactory: () => new Date(Date.now() + 1000 * 60 * 60 * 24 * 90), 56 | }) 57 | expireAt!: Date; 58 | 59 | [k: string]: any; 60 | 61 | tag(...tags: string[]) { 62 | for (const t of tags) { 63 | if (!this.tags.includes(t)) { 64 | this.tags.push(t); 65 | } 66 | } 67 | } 68 | 69 | save() { 70 | return (this.constructor as typeof APICall).save(this); 71 | } 72 | } 73 | 74 | 75 | export class RateLimitDesc extends AutoCastable { 76 | @Prop({ 77 | default: 1000 78 | }) 79 | occurrence!: number; 80 | 81 | @Prop({ 82 | default: 3600 83 | }) 84 | periodSeconds!: number; 85 | 86 | @Prop() 87 | notBefore?: Date; 88 | 89 | @Prop() 90 | notAfter?: Date; 91 | 92 | isEffective() { 93 | const now = new Date(); 94 | if (this.notBefore && this.notBefore > now) { 95 | return false; 96 | } 97 | if (this.notAfter && this.notAfter < now) { 98 | return false; 99 | } 100 | 101 | return true; 102 | } 103 | } 104 | 105 | 106 | @singleton() 107 | export class RateLimitControl extends AsyncService { 108 | 109 | logger = this.globalLogger.child({ service: this.constructor.name }); 110 | 111 | constructor( 112 | protected globalLogger: GlobalLogger, 113 | ) { 114 | super(...arguments); 115 | } 116 | 117 | override async init() { 118 | await this.dependencyReady(); 119 | 120 | this.emit('ready'); 121 | } 122 | 123 | async queryByUid(uid: string, pointInTime: Date, ...tags: string[]) { 124 | let q = APICall.COLLECTION 125 | .orderBy('createdAt', 'asc') 126 | .where('createdAt', '>=', pointInTime) 127 | .where('status', 'in', [API_CALL_STATUS.SUCCESS, API_CALL_STATUS.PENDING]) 128 | .where('uid', '==', uid); 129 | if (tags.length) { 130 | q = q.where('tags', 'array-contains-any', tags); 131 | } 132 | 133 | return APICall.fromFirestoreQuery(q); 134 | } 135 | 136 | async queryByIp(ip: string, pointInTime: Date, ...tags: string[]) { 137 | let q = APICall.COLLECTION 138 | .orderBy('createdAt', 'asc') 139 | .where('createdAt', '>=', pointInTime) 140 | .where('status', 'in', [API_CALL_STATUS.SUCCESS, API_CALL_STATUS.PENDING]) 141 | .where('ip', '==', ip); 142 | if (tags.length) { 143 | q = q.where('tags', 'array-contains-any', tags); 144 | } 145 | 146 | return APICall.fromFirestoreQuery(q); 147 | } 148 | 149 | async assertUidPeriodicLimit(uid: string, pointInTime: Date, limit: number, ...tags: string[]) { 150 | if (limit <= 0) { 151 | throw new ResourcePolicyDenyError(`This UID(${uid}) is not allowed to call this endpoint (rate limit quota is 0).`); 152 | } 153 | 154 | let q = APICall.COLLECTION 155 | .orderBy('createdAt', 'asc') 156 | .where('createdAt', '>=', pointInTime) 157 | .where('status', 'in', [API_CALL_STATUS.SUCCESS, API_CALL_STATUS.PENDING]) 158 | .where('uid', '==', uid); 159 | if (tags.length) { 160 | q = q.where('tags', 'array-contains-any', tags); 161 | } 162 | const count = (await q.count().get()).data().count; 163 | 164 | if (count >= limit) { 165 | const r = await APICall.fromFirestoreQuery(q.limit(1)); 166 | const [r1] = r; 167 | 168 | const dtMs = Math.abs(r1.createdAt?.valueOf() - pointInTime.valueOf()); 169 | const dtSec = Math.ceil(dtMs / 1000); 170 | 171 | throw RateLimitTriggeredError.from({ 172 | message: `Per UID rate limit exceeded (${tags.join(',') || 'called'} ${limit} times since ${pointInTime})`, 173 | retryAfter: dtSec, 174 | }); 175 | } 176 | 177 | return count + 1; 178 | } 179 | 180 | async assertIPPeriodicLimit(ip: string, pointInTime: Date, limit: number, ...tags: string[]) { 181 | let q = APICall.COLLECTION 182 | .orderBy('createdAt', 'asc') 183 | .where('createdAt', '>=', pointInTime) 184 | .where('status', 'in', [API_CALL_STATUS.SUCCESS, API_CALL_STATUS.PENDING]) 185 | .where('ip', '==', ip); 186 | if (tags.length) { 187 | q = q.where('tags', 'array-contains-any', tags); 188 | } 189 | 190 | const count = (await q.count().get()).data().count; 191 | 192 | if (count >= limit) { 193 | const r = await APICall.fromFirestoreQuery(q.limit(1)); 194 | const [r1] = r; 195 | 196 | const dtMs = Math.abs(r1.createdAt?.valueOf() - pointInTime.valueOf()); 197 | const dtSec = Math.ceil(dtMs / 1000); 198 | 199 | throw RateLimitTriggeredError.from({ 200 | message: `Per IP rate limit exceeded (${tags.join(',') || 'called'} ${limit} times since ${pointInTime})`, 201 | retryAfter: dtSec, 202 | }); 203 | } 204 | 205 | return count + 1; 206 | } 207 | 208 | record(partialRecord: Partial) { 209 | const record = APICall.from(partialRecord); 210 | const newId = APICall.COLLECTION.doc().id; 211 | record._id = newId; 212 | 213 | return record; 214 | } 215 | 216 | // async simpleRPCUidBasedLimit(rpcReflect: RPCReflection, uid: string, tags: string[] = [], 217 | // ...inputCriterion: RateLimitDesc[] | [Date, number][]) { 218 | // const criterion = inputCriterion.map((c) => { return Array.isArray(c) ? c : this.rateLimitDescToCriterion(c); }); 219 | 220 | // await Promise.all(criterion.map(([pointInTime, n]) => 221 | // this.assertUidPeriodicLimit(uid, pointInTime, n, ...tags))); 222 | 223 | // const r = this.record({ 224 | // uid, 225 | // tags, 226 | // }); 227 | 228 | // r.save().catch((err) => this.logger.warn(`Failed to save rate limit record`, { err })); 229 | // rpcReflect.then(() => { 230 | // r.status = API_CALL_STATUS.SUCCESS; 231 | // r.save() 232 | // .catch((err) => this.logger.warn(`Failed to save rate limit record`, { err })); 233 | // }); 234 | // rpcReflect.catch((err) => { 235 | // r.status = API_CALL_STATUS.ERROR; 236 | // r.error = err.toString(); 237 | // r.save() 238 | // .catch((err) => this.logger.warn(`Failed to save rate limit record`, { err })); 239 | // }); 240 | 241 | // return r; 242 | // } 243 | 244 | rateLimitDescToCriterion(rateLimitDesc: RateLimitDesc) { 245 | return [new Date(Date.now() - rateLimitDesc.periodSeconds * 1000), rateLimitDesc.occurrence] as [Date, number]; 246 | } 247 | 248 | // async simpleRpcIPBasedLimit(rpcReflect: RPCReflection, ip: string, tags: string[] = [], 249 | // ...inputCriterion: RateLimitDesc[] | [Date, number][]) { 250 | // const criterion = inputCriterion.map((c) => { return Array.isArray(c) ? c : this.rateLimitDescToCriterion(c); }); 251 | // await Promise.all(criterion.map(([pointInTime, n]) => 252 | // this.assertIPPeriodicLimit(ip, pointInTime, n, ...tags))); 253 | 254 | // const r = this.record({ 255 | // ip, 256 | // tags, 257 | // }); 258 | 259 | // r.save().catch((err) => this.logger.warn(`Failed to save rate limit record`, { err })); 260 | // rpcReflect.then(() => { 261 | // r.status = API_CALL_STATUS.SUCCESS; 262 | // r.save() 263 | // .catch((err) => this.logger.warn(`Failed to save rate limit record`, { err })); 264 | // }); 265 | // rpcReflect.catch((err) => { 266 | // r.status = API_CALL_STATUS.ERROR; 267 | // r.error = err.toString(); 268 | // r.save() 269 | // .catch((err) => this.logger.warn(`Failed to save rate limit record`, { err })); 270 | // }); 271 | 272 | // return r; 273 | // } 274 | } 275 | 276 | const instance = container.resolve(RateLimitControl); 277 | 278 | export default instance; 279 | -------------------------------------------------------------------------------- /jina-ai/src/server.ts: -------------------------------------------------------------------------------- 1 | import 'reflect-metadata' 2 | import express from 'express'; 3 | import { jinaAiMiddleware } from "./patch-express"; 4 | import { Server } from 'http'; 5 | 6 | const app = require('../..').default; 7 | 8 | const rootApp = express(); 9 | rootApp.set('trust proxy', true); 10 | rootApp.use(jinaAiMiddleware, app); 11 | 12 | 13 | const port = process.env.PORT || 3000; 14 | 15 | let server: Server | undefined; 16 | // Export server startup function for better testing 17 | export function startServer() { 18 | return rootApp.listen(port, () => { 19 | console.log(`Server running at http://localhost:${port}`); 20 | }); 21 | } 22 | 23 | // Start server if running directly 24 | if (process.env.NODE_ENV !== 'test') { 25 | server = startServer(); 26 | } 27 | 28 | process.on('unhandledRejection', (_err) => `Is false alarm`); 29 | 30 | process.on('uncaughtException', (err) => { 31 | console.log('Uncaught exception', err); 32 | 33 | // Looks like Firebase runtime does not handle error properly. 34 | // Make sure to quit the process. 35 | process.nextTick(() => process.exit(1)); 36 | console.error('Uncaught exception, process quit.'); 37 | throw err; 38 | }); 39 | 40 | const sigHandler = (signal: string) => { 41 | console.log(`Received ${signal}, exiting...`); 42 | if (server && server.listening) { 43 | console.log(`Shutting down gracefully...`); 44 | console.log(`Waiting for the server to drain and close...`); 45 | server.close((err) => { 46 | if (err) { 47 | console.error('Error while closing server', err); 48 | return; 49 | } 50 | process.exit(0); 51 | }); 52 | server.closeIdleConnections(); 53 | } 54 | 55 | } 56 | process.on('SIGTERM', sigHandler); 57 | process.on('SIGINT', sigHandler); -------------------------------------------------------------------------------- /jina-ai/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ES2020", 4 | "module": "node16", 5 | "outDir": "./dist", 6 | "rootDir": "./src", 7 | "sourceMap": true, 8 | "esModuleInterop": true, 9 | "skipLibCheck": true, 10 | "forceConsistentCasingInFileNames": true, 11 | "strict": true, 12 | "experimentalDecorators": true, 13 | "emitDecoratorMetadata": true, 14 | "resolveJsonModule": true 15 | } 16 | } 17 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "node-deepresearch", 3 | "version": "1.0.0", 4 | "main": "dist/app.js", 5 | "files": [ 6 | "dist", 7 | "README.md", 8 | "LICENSE" 9 | ], 10 | "scripts": { 11 | "build": "tsc", 12 | "dev": "npx ts-node src/agent.ts", 13 | "search": "npx ts-node src/test-duck.ts", 14 | "rewrite": "npx ts-node src/tools/query-rewriter.ts", 15 | "lint": "eslint . --ext .ts", 16 | "lint:fix": "eslint . --ext .ts --fix", 17 | "serve": "ts-node src/server.ts", 18 | "start": "ts-node src/server.ts", 19 | "eval": "ts-node src/evals/batch-evals.ts", 20 | "test": "jest --testTimeout=30000", 21 | "test:watch": "jest --watch", 22 | "test:docker": "jest src/__tests__/docker.test.ts --testTimeout=300000" 23 | }, 24 | "keywords": [], 25 | "author": "Jina AI", 26 | "license": "Apache-2.0", 27 | "description": "", 28 | "dependencies": { 29 | "@ai-sdk/google": "^1.0.0", 30 | "@ai-sdk/openai": "^1.1.9", 31 | "@types/jsdom": "^21.1.7", 32 | "ai": "^4.1.26", 33 | "axios": "^1.7.9", 34 | "commander": "^13.1.0", 35 | "cors": "^2.8.5", 36 | "dotenv": "^16.4.7", 37 | "duck-duck-scrape": "^2.2.7", 38 | "express": "^4.21.2", 39 | "hjson": "^3.2.2", 40 | "jsdom": "^26.0.0", 41 | "node-fetch": "^3.3.2", 42 | "undici": "^7.3.0", 43 | "zod": "^3.22.4", 44 | "zod-to-json-schema": "^3.24.1" 45 | }, 46 | "devDependencies": { 47 | "@types/commander": "^2.12.0", 48 | "@types/cors": "^2.8.17", 49 | "@types/express": "^5.0.0", 50 | "@types/hjson": "^2.4.6", 51 | "@types/jest": "^29.5.14", 52 | "@types/node": "^22.10.10", 53 | "@types/node-fetch": "^2.6.12", 54 | "@types/supertest": "^6.0.2", 55 | "@typescript-eslint/eslint-plugin": "^7.0.1", 56 | "@typescript-eslint/parser": "^7.0.1", 57 | "eslint": "^8.56.0", 58 | "jest": "^29.7.0", 59 | "supertest": "^7.0.0", 60 | "ts-jest": "^29.2.5", 61 | "ts-node": "^10.9.2", 62 | "typescript": "^5.7.3" 63 | }, 64 | "optionalDependencies": { 65 | "@ai-sdk/google-vertex": "^2.1.12" 66 | } 67 | } 68 | -------------------------------------------------------------------------------- /src/__tests__/agent.test.ts: -------------------------------------------------------------------------------- 1 | import { getResponse } from '../agent'; 2 | import { generateObject } from 'ai'; 3 | import { search } from '../tools/jina-search'; 4 | import { readUrl } from '../tools/read'; 5 | 6 | // Mock external dependencies 7 | jest.mock('ai', () => ({ 8 | generateObject: jest.fn() 9 | })); 10 | 11 | jest.mock('../tools/jina-search', () => ({ 12 | search: jest.fn() 13 | })); 14 | 15 | jest.mock('../tools/read', () => ({ 16 | readUrl: jest.fn() 17 | })); 18 | 19 | describe('getResponse', () => { 20 | beforeEach(() => { 21 | // Mock generateObject to return a valid response 22 | (generateObject as jest.Mock).mockResolvedValue({ 23 | object: { action: 'answer', answer: 'mocked response', references: [], think: 'mocked thought' }, 24 | usage: { totalTokens: 100 } 25 | }); 26 | 27 | // Mock search to return empty results 28 | (search as jest.Mock).mockResolvedValue({ 29 | response: { data: [] } 30 | }); 31 | 32 | // Mock readUrl to return empty content 33 | (readUrl as jest.Mock).mockResolvedValue({ 34 | response: { data: { content: '', url: 'test-url' } }, 35 | tokens: 0 36 | }); 37 | }); 38 | 39 | afterEach(() => { 40 | jest.useRealTimers(); 41 | jest.clearAllMocks(); 42 | }); 43 | 44 | it('should handle search action', async () => { 45 | const result = await getResponse('What is TypeScript?', 50000); // Increased token budget to handle real-world usage 46 | expect(result.result.action).toBeDefined(); 47 | expect(result.context).toBeDefined(); 48 | expect(result.context.tokenTracker).toBeDefined(); 49 | expect(result.context.actionTracker).toBeDefined(); 50 | }, 30000); 51 | }); 52 | -------------------------------------------------------------------------------- /src/__tests__/docker.test.ts: -------------------------------------------------------------------------------- 1 | import { exec } from 'child_process'; 2 | import { promisify } from 'util'; 3 | 4 | const execAsync = promisify(exec); 5 | 6 | describe('Docker build', () => { 7 | jest.setTimeout(300000); // 5 minutes for build 8 | 9 | it('should build Docker image successfully', async () => { 10 | const { stderr } = await execAsync('docker build -t node-deepresearch-test .'); 11 | expect(stderr).not.toContain('error'); 12 | }); 13 | 14 | it('should start container and respond to health check', async () => { 15 | // Start container with mock API keys 16 | await execAsync( 17 | 'docker run -d --name test-container -p 3001:3000 ' + 18 | '-e GEMINI_API_KEY=mock_key ' + 19 | '-e JINA_API_KEY=mock_key ' + 20 | 'node-deepresearch-test' 21 | ); 22 | 23 | // Wait for container to start 24 | await new Promise(resolve => setTimeout(resolve, 5000)); 25 | 26 | try { 27 | // Check if server responds 28 | const { stdout } = await execAsync('curl -s http://localhost:3001/health'); 29 | expect(stdout).toContain('ok'); 30 | } finally { 31 | // Cleanup 32 | await execAsync('docker rm -f test-container').catch(console.error); 33 | } 34 | }); 35 | 36 | afterAll(async () => { 37 | // Clean up any leftover containers 38 | await execAsync('docker rm -f test-container').catch(() => {}); 39 | await execAsync('docker rmi node-deepresearch-test').catch(() => {}); 40 | }); 41 | }); 42 | -------------------------------------------------------------------------------- /src/__tests__/server.test.ts: -------------------------------------------------------------------------------- 1 | import request from 'supertest'; 2 | import { EventEmitter } from 'events'; 3 | import type { Express } from 'express'; 4 | 5 | const TEST_SECRET = 'test-secret'; 6 | let app: Express; 7 | 8 | describe('/v1/chat/completions', () => { 9 | jest.setTimeout(120000); // Increase timeout for all tests in this suite 10 | 11 | beforeEach(async () => { 12 | // Set up test environment 13 | process.env.NODE_ENV = 'test'; 14 | process.env.LLM_PROVIDER = 'openai'; // Use OpenAI provider for tests 15 | process.env.OPENAI_API_KEY = 'test-key'; 16 | process.env.JINA_API_KEY = 'test-key'; 17 | 18 | // Clean up any existing secret 19 | const existingSecretIndex = process.argv.findIndex(arg => arg.startsWith('--secret=')); 20 | if (existingSecretIndex !== -1) { 21 | process.argv.splice(existingSecretIndex, 1); 22 | } 23 | 24 | // Set up test secret and import server module 25 | process.argv.push(`--secret=${TEST_SECRET}`); 26 | 27 | // Import server module (jest.resetModules() is called automatically before each test) 28 | const { default: serverModule } = await require('../app'); 29 | app = serverModule; 30 | }); 31 | 32 | afterEach(async () => { 33 | // Clean up environment variables 34 | delete process.env.OPENAI_API_KEY; 35 | delete process.env.JINA_API_KEY; 36 | 37 | // Clean up any remaining event listeners 38 | const emitter = EventEmitter.prototype; 39 | emitter.removeAllListeners(); 40 | emitter.setMaxListeners(emitter.getMaxListeners() + 1); 41 | 42 | // Clean up test secret 43 | const secretIndex = process.argv.findIndex(arg => arg.startsWith('--secret=')); 44 | if (secretIndex !== -1) { 45 | process.argv.splice(secretIndex, 1); 46 | } 47 | 48 | // Wait for any pending promises to settle 49 | await new Promise(resolve => setTimeout(resolve, 500)); 50 | 51 | // Reset module cache to ensure clean state 52 | jest.resetModules(); 53 | }); 54 | it('should require authentication when secret is set', async () => { 55 | // Note: secret is already set in beforeEach 56 | 57 | const response = await request(app) 58 | .post('/v1/chat/completions') 59 | .send({ 60 | model: 'test-model', 61 | messages: [{ role: 'user', content: 'test' }] 62 | }); 63 | expect(response.status).toBe(401); 64 | }); 65 | 66 | it('should allow requests without auth when no secret is set', async () => { 67 | // Remove secret for this test 68 | const secretIndex = process.argv.findIndex(arg => arg.startsWith('--secret=')); 69 | if (secretIndex !== -1) { 70 | process.argv.splice(secretIndex, 1); 71 | } 72 | 73 | // Reset module cache to ensure clean state 74 | jest.resetModules(); 75 | 76 | // Reload server module without secret 77 | const { default: serverModule } = await require('../app'); 78 | app = serverModule; 79 | 80 | const response = await request(app) 81 | .post('/v1/chat/completions') 82 | .send({ 83 | model: 'test-model', 84 | messages: [{ role: 'user', content: 'test' }] 85 | }); 86 | expect(response.status).toBe(200); 87 | }); 88 | 89 | it('should reject requests without user message', async () => { 90 | const response = await request(app) 91 | .post('/v1/chat/completions') 92 | .set('Authorization', `Bearer ${TEST_SECRET}`) 93 | .send({ 94 | model: 'test-model', 95 | messages: [{ role: 'developer', content: 'test' }] 96 | }); 97 | expect(response.status).toBe(400); 98 | expect(response.body.error).toBe('Last message must be from user'); 99 | }); 100 | 101 | it('should handle non-streaming request', async () => { 102 | const response = await request(app) 103 | .post('/v1/chat/completions') 104 | .set('Authorization', `Bearer ${TEST_SECRET}`) 105 | .send({ 106 | model: 'test-model', 107 | messages: [{ role: 'user', content: 'test' }] 108 | }); 109 | expect(response.status).toBe(200); 110 | expect(response.body).toMatchObject({ 111 | object: 'chat.completion', 112 | choices: [{ 113 | message: { 114 | role: 'assistant' 115 | } 116 | }] 117 | }); 118 | }); 119 | 120 | it('should handle streaming request and track tokens correctly', async () => { 121 | return new Promise((resolve, reject) => { 122 | let isDone = false; 123 | let totalCompletionTokens = 0; 124 | 125 | const cleanup = () => { 126 | clearTimeout(timeoutHandle); 127 | isDone = true; 128 | resolve(); 129 | }; 130 | 131 | const timeoutHandle = setTimeout(() => { 132 | if (!isDone) { 133 | cleanup(); 134 | reject(new Error('Test timed out')); 135 | } 136 | }, 30000); 137 | 138 | request(app) 139 | .post('/v1/chat/completions') 140 | .set('Authorization', `Bearer ${TEST_SECRET}`) 141 | .send({ 142 | model: 'test-model', 143 | messages: [{ role: 'user', content: 'test' }], 144 | stream: true 145 | }) 146 | .buffer(true) 147 | .parse((res, callback) => { 148 | const response = res as unknown as { 149 | on(event: 'data', listener: (chunk: Buffer) => void): void; 150 | on(event: 'end', listener: () => void): void; 151 | on(event: 'error', listener: (err: Error) => void): void; 152 | }; 153 | let responseData = ''; 154 | 155 | response.on('error', (err) => { 156 | cleanup(); 157 | callback(err, null); 158 | }); 159 | 160 | response.on('data', (chunk) => { 161 | responseData += chunk.toString(); 162 | }); 163 | 164 | response.on('end', () => { 165 | try { 166 | callback(null, responseData); 167 | } catch (err) { 168 | cleanup(); 169 | callback(err instanceof Error ? err : new Error(String(err)), null); 170 | } 171 | }); 172 | }) 173 | .end((err, res) => { 174 | if (err) return reject(err); 175 | 176 | expect(res.status).toBe(200); 177 | expect(res.headers['content-type']).toBe('text/event-stream'); 178 | 179 | // Verify stream format and content 180 | if (isDone) return; // Prevent multiple resolves 181 | 182 | const responseText = res.body as string; 183 | const chunks = responseText 184 | .split('\n\n') 185 | .filter((line: string) => line.startsWith('data: ')) 186 | .map((line: string) => JSON.parse(line.replace('data: ', ''))); 187 | 188 | // Process all chunks 189 | expect(chunks.length).toBeGreaterThan(0); 190 | 191 | // Verify initial chunk format 192 | expect(chunks[0]).toMatchObject({ 193 | id: expect.any(String), 194 | object: 'chat.completion.chunk', 195 | choices: [{ 196 | index: 0, 197 | delta: { role: 'assistant' }, 198 | logprobs: null, 199 | finish_reason: null 200 | }] 201 | }); 202 | 203 | // Verify content chunks have content 204 | chunks.slice(1).forEach(chunk => { 205 | const content = chunk.choices[0].delta.content; 206 | if (content && content.trim()) { 207 | totalCompletionTokens += 1; // Count 1 token per chunk as per Vercel convention 208 | } 209 | expect(chunk).toMatchObject({ 210 | object: 'chat.completion.chunk', 211 | choices: [{ 212 | delta: expect.objectContaining({ 213 | content: expect.any(String) 214 | }) 215 | }] 216 | }); 217 | }); 218 | 219 | // Verify final chunk format if present 220 | const lastChunk = chunks[chunks.length - 1]; 221 | if (lastChunk?.choices?.[0]?.finish_reason === 'stop') { 222 | expect(lastChunk).toMatchObject({ 223 | object: 'chat.completion.chunk', 224 | choices: [{ 225 | delta: {}, 226 | finish_reason: 'stop' 227 | }] 228 | }); 229 | } 230 | 231 | // Verify we tracked some completion tokens 232 | expect(totalCompletionTokens).toBeGreaterThan(0); 233 | 234 | // Clean up and resolve 235 | if (!isDone) { 236 | cleanup(); 237 | } 238 | }); 239 | }); 240 | }); 241 | 242 | it('should track tokens correctly in error response', async () => { 243 | const response = await request(app) 244 | .post('/v1/chat/completions') 245 | .set('Authorization', `Bearer ${TEST_SECRET}`) 246 | .send({ 247 | model: 'test-model', 248 | messages: [] // Invalid messages array 249 | }); 250 | 251 | expect(response.status).toBe(400); 252 | expect(response.body).toHaveProperty('error'); 253 | expect(response.body.error).toBe('Messages array is required and must not be empty'); 254 | 255 | // Make another request to verify token tracking after error 256 | const validResponse = await request(app) 257 | .post('/v1/chat/completions') 258 | .set('Authorization', `Bearer ${TEST_SECRET}`) 259 | .send({ 260 | model: 'test-model', 261 | messages: [{ role: 'user', content: 'test' }] 262 | }); 263 | 264 | // Verify token tracking still works after error 265 | expect(validResponse.body.usage).toMatchObject({ 266 | prompt_tokens: expect.any(Number), 267 | completion_tokens: expect.any(Number), 268 | total_tokens: expect.any(Number) 269 | }); 270 | 271 | // Basic token tracking structure should be present 272 | expect(validResponse.body.usage.total_tokens).toBe( 273 | validResponse.body.usage.prompt_tokens + validResponse.body.usage.completion_tokens 274 | ); 275 | }); 276 | 277 | it('should provide token usage in Vercel AI SDK format', async () => { 278 | const response = await request(app) 279 | .post('/v1/chat/completions') 280 | .set('Authorization', `Bearer ${TEST_SECRET}`) 281 | .send({ 282 | model: 'test-model', 283 | messages: [{ role: 'user', content: 'test' }] 284 | }); 285 | 286 | expect(response.status).toBe(200); 287 | const usage = response.body.usage; 288 | 289 | expect(usage).toMatchObject({ 290 | prompt_tokens: expect.any(Number), 291 | completion_tokens: expect.any(Number), 292 | total_tokens: expect.any(Number) 293 | }); 294 | 295 | // Basic token tracking structure should be present 296 | expect(usage.total_tokens).toBe( 297 | usage.prompt_tokens + usage.completion_tokens 298 | ); 299 | }); 300 | }); 301 | -------------------------------------------------------------------------------- /src/cli.ts: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env node 2 | import { Command } from 'commander'; 3 | import { getResponse } from './agent'; 4 | import { version } from '../package.json'; 5 | 6 | const program = new Command(); 7 | 8 | program 9 | .name('deepresearch') 10 | .description('AI-powered research assistant that keeps searching until it finds the answer') 11 | .version(version) 12 | .argument('', 'The research query to investigate') 13 | .option('-t, --token-budget ', 'Maximum token budget', (val) => { 14 | const num = parseInt(val); 15 | if (isNaN(num)) throw new Error('Invalid token budget: must be a number'); 16 | return num; 17 | }, 1000000) 18 | .option('-m, --max-attempts ', 'Maximum bad attempts before giving up', (val) => { 19 | const num = parseInt(val); 20 | if (isNaN(num)) throw new Error('Invalid max attempts: must be a number'); 21 | return num; 22 | }, 3) 23 | .option('-v, --verbose', 'Show detailed progress') 24 | .action(async (query: string, options: any) => { 25 | try { 26 | const { result } = await getResponse( 27 | query, 28 | parseInt(options.tokenBudget), 29 | parseInt(options.maxAttempts), 30 | ); 31 | 32 | if (result.action === 'answer') { 33 | console.log('\nAnswer:', result.answer); 34 | if (result.references?.length) { 35 | console.log('\nReferences:'); 36 | result.references.forEach(ref => { 37 | console.log(`- ${ref.url}`); 38 | console.log(` "${ref.exactQuote}"`); 39 | }); 40 | } 41 | } 42 | } catch (error) { 43 | console.error('Error:', error instanceof Error ? error.message : String(error)); 44 | process.exit(1); 45 | } 46 | }); 47 | 48 | program.parse(); 49 | -------------------------------------------------------------------------------- /src/config.ts: -------------------------------------------------------------------------------- 1 | import dotenv from 'dotenv'; 2 | import { ProxyAgent, setGlobalDispatcher } from 'undici'; 3 | import { createGoogleGenerativeAI } from '@ai-sdk/google'; 4 | import { createOpenAI, OpenAIProviderSettings } from '@ai-sdk/openai'; 5 | import configJson from '../config.json'; 6 | // Load environment variables 7 | dotenv.config(); 8 | 9 | // Types 10 | export type LLMProvider = 'openai' | 'gemini' | 'vertex'; 11 | export type ToolName = keyof typeof configJson.models.gemini.tools; 12 | 13 | // Type definitions for our config structure 14 | type EnvConfig = typeof configJson.env; 15 | 16 | interface ProviderConfig { 17 | createClient: string; 18 | clientConfig?: Record; 19 | } 20 | 21 | // Environment setup 22 | const env: EnvConfig = { ...configJson.env }; 23 | (Object.keys(env) as (keyof EnvConfig)[]).forEach(key => { 24 | if (process.env[key]) { 25 | env[key] = process.env[key] || env[key]; 26 | } 27 | }); 28 | 29 | // Setup proxy if present 30 | if (env.https_proxy) { 31 | try { 32 | const proxyUrl = new URL(env.https_proxy).toString(); 33 | const dispatcher = new ProxyAgent({ uri: proxyUrl }); 34 | setGlobalDispatcher(dispatcher); 35 | } catch (error) { 36 | console.error('Failed to set proxy:', error); 37 | } 38 | } 39 | 40 | // Export environment variables 41 | export const OPENAI_BASE_URL = env.OPENAI_BASE_URL; 42 | export const GEMINI_API_KEY = env.GEMINI_API_KEY; 43 | export const OPENAI_API_KEY = env.OPENAI_API_KEY; 44 | export const JINA_API_KEY = env.JINA_API_KEY; 45 | export const BRAVE_API_KEY = env.BRAVE_API_KEY; 46 | export const SERPER_API_KEY = env.SERPER_API_KEY; 47 | export const SEARCH_PROVIDER = configJson.defaults.search_provider; 48 | export const STEP_SLEEP = configJson.defaults.step_sleep; 49 | 50 | // Determine LLM provider 51 | export const LLM_PROVIDER: LLMProvider = (() => { 52 | const provider = process.env.LLM_PROVIDER || configJson.defaults.llm_provider; 53 | if (!isValidProvider(provider)) { 54 | throw new Error(`Invalid LLM provider: ${provider}`); 55 | } 56 | return provider; 57 | })(); 58 | 59 | function isValidProvider(provider: string): provider is LLMProvider { 60 | return provider === 'openai' || provider === 'gemini' || provider === 'vertex'; 61 | } 62 | 63 | interface ToolConfig { 64 | model: string; 65 | temperature: number; 66 | maxTokens: number; 67 | } 68 | 69 | interface ToolOverrides { 70 | temperature?: number; 71 | maxTokens?: number; 72 | } 73 | 74 | // Get tool configuration 75 | export function getToolConfig(toolName: ToolName): ToolConfig { 76 | const providerConfig = configJson.models[LLM_PROVIDER === 'vertex' ? 'gemini' : LLM_PROVIDER]; 77 | const defaultConfig = providerConfig.default; 78 | const toolOverrides = providerConfig.tools[toolName] as ToolOverrides; 79 | 80 | return { 81 | model: process.env.DEFAULT_MODEL_NAME || defaultConfig.model, 82 | temperature: toolOverrides.temperature ?? defaultConfig.temperature, 83 | maxTokens: toolOverrides.maxTokens ?? defaultConfig.maxTokens 84 | }; 85 | } 86 | 87 | export function getMaxTokens(toolName: ToolName): number { 88 | return getToolConfig(toolName).maxTokens; 89 | } 90 | 91 | // Get model instance 92 | export function getModel(toolName: ToolName) { 93 | const config = getToolConfig(toolName); 94 | const providerConfig = (configJson.providers as Record)[LLM_PROVIDER]; 95 | 96 | if (LLM_PROVIDER === 'openai') { 97 | if (!OPENAI_API_KEY) { 98 | throw new Error('OPENAI_API_KEY not found'); 99 | } 100 | 101 | const opt: OpenAIProviderSettings = { 102 | apiKey: OPENAI_API_KEY, 103 | compatibility: providerConfig?.clientConfig?.compatibility 104 | }; 105 | 106 | if (OPENAI_BASE_URL) { 107 | opt.baseURL = OPENAI_BASE_URL; 108 | } 109 | 110 | return createOpenAI(opt)(config.model); 111 | } 112 | 113 | if (LLM_PROVIDER === 'vertex') { 114 | const createVertex = require('@ai-sdk/google-vertex').createVertex; 115 | if (toolName === 'searchGrounding') { 116 | return createVertex({ project: process.env.GCLOUD_PROJECT, ...providerConfig?.clientConfig })(config.model, { useSearchGrounding: true }); 117 | } 118 | return createVertex({ project: process.env.GCLOUD_PROJECT, ...providerConfig?.clientConfig })(config.model); 119 | } 120 | 121 | if (!GEMINI_API_KEY) { 122 | throw new Error('GEMINI_API_KEY not found'); 123 | } 124 | 125 | if (toolName === 'searchGrounding') { 126 | return createGoogleGenerativeAI({ apiKey: GEMINI_API_KEY })(config.model, { useSearchGrounding: true }); 127 | } 128 | return createGoogleGenerativeAI({ apiKey: GEMINI_API_KEY })(config.model); 129 | } 130 | 131 | // Validate required environment variables 132 | if (LLM_PROVIDER === 'gemini' && !GEMINI_API_KEY) throw new Error("GEMINI_API_KEY not found"); 133 | if (LLM_PROVIDER === 'openai' && !OPENAI_API_KEY) throw new Error("OPENAI_API_KEY not found"); 134 | if (!JINA_API_KEY) throw new Error("JINA_API_KEY not found"); 135 | 136 | // Log all configurations 137 | const configSummary = { 138 | provider: { 139 | name: LLM_PROVIDER, 140 | model: LLM_PROVIDER === 'openai' 141 | ? configJson.models.openai.default.model 142 | : configJson.models.gemini.default.model, 143 | ...(LLM_PROVIDER === 'openai' && { baseUrl: OPENAI_BASE_URL }) 144 | }, 145 | search: { 146 | provider: SEARCH_PROVIDER 147 | }, 148 | tools: Object.fromEntries( 149 | Object.keys(configJson.models[LLM_PROVIDER === 'vertex' ? 'gemini' : LLM_PROVIDER].tools).map(name => [ 150 | name, 151 | getToolConfig(name as ToolName) 152 | ]) 153 | ), 154 | defaults: { 155 | stepSleep: STEP_SLEEP 156 | } 157 | }; 158 | 159 | console.log('Configuration Summary:', JSON.stringify(configSummary, null, 2)); 160 | -------------------------------------------------------------------------------- /src/evals/batch-evals.ts: -------------------------------------------------------------------------------- 1 | import fs from 'fs/promises'; 2 | import {exec} from 'child_process'; 3 | import {promisify} from 'util'; 4 | import {getResponse} from '../agent'; 5 | import {generateObject} from 'ai'; 6 | import {GEMINI_API_KEY} from '../config'; 7 | import {z} from 'zod'; 8 | import {AnswerAction, TrackerContext} from "../types"; 9 | import {createGoogleGenerativeAI} from "@ai-sdk/google"; 10 | 11 | const execAsync = promisify(exec); 12 | 13 | interface Question { 14 | question: string; 15 | answer: string; 16 | } 17 | 18 | interface EvaluationResult { 19 | pass: boolean; 20 | reason: string; 21 | total_steps: number; 22 | total_tokens: number; 23 | question: string; 24 | expected_answer: string; 25 | actual_answer: string; 26 | } 27 | 28 | interface EvaluationStats { 29 | model_name: string; 30 | pass_rate: number; 31 | avg_steps: number; 32 | max_steps: number; 33 | min_steps: number; 34 | median_steps: number; 35 | avg_tokens: number; 36 | median_tokens: number; 37 | max_tokens: number; 38 | min_tokens: number; 39 | } 40 | 41 | function calculateMedian(numbers: number[]): number { 42 | const sorted = [...numbers].sort((a, b) => a - b); 43 | const middle = Math.floor(sorted.length / 2); 44 | 45 | if (sorted.length % 2 === 0) { 46 | return (sorted[middle - 1] + sorted[middle]) / 2; 47 | } 48 | return sorted[middle]; 49 | } 50 | 51 | function calculateStats(results: EvaluationResult[], modelName: string): EvaluationStats { 52 | const steps = results.map(r => r.total_steps); 53 | const tokens = results.map(r => r.total_tokens); 54 | const passCount = results.filter(r => r.pass).length; 55 | 56 | return { 57 | model_name: modelName, 58 | pass_rate: (passCount / results.length) * 100, 59 | avg_steps: steps.reduce((a, b) => a + b, 0) / steps.length, 60 | max_steps: Math.max(...steps), 61 | min_steps: Math.min(...steps), 62 | median_steps: calculateMedian(steps), 63 | avg_tokens: tokens.reduce((a, b) => a + b, 0) / tokens.length, 64 | median_tokens: calculateMedian(tokens), 65 | max_tokens: Math.max(...tokens), 66 | min_tokens: Math.min(...tokens) 67 | }; 68 | } 69 | 70 | function printStats(stats: EvaluationStats): void { 71 | console.log('\n=== Evaluation Statistics ==='); 72 | console.log(`Model: ${stats.model_name}`); 73 | console.log(`Pass Rate: ${stats.pass_rate.toFixed(0)}%`); 74 | console.log(`Average Steps: ${stats.avg_steps.toFixed(0)}`); 75 | console.log(`Maximum Steps: ${stats.max_steps}`); 76 | console.log(`Minimum Steps: ${stats.min_steps}`); 77 | console.log(`Median Steps: ${stats.median_steps.toFixed(0)}`); 78 | console.log(`Average Tokens: ${stats.avg_tokens.toFixed(0)}`); 79 | console.log(`Median Tokens: ${stats.median_tokens.toFixed(0)}`); 80 | console.log(`Maximum Tokens: ${stats.max_tokens}`); 81 | console.log(`Minimum Tokens: ${stats.min_tokens}`); 82 | console.log('===========================\n'); 83 | } 84 | 85 | async function getCurrentGitCommit(): Promise { 86 | try { 87 | const {stdout} = await execAsync('git rev-parse --short HEAD'); 88 | return stdout.trim(); 89 | } catch (error) { 90 | console.error('Error getting git commit:', error); 91 | return 'unknown'; 92 | } 93 | } 94 | 95 | async function evaluateAnswer(expectedAnswer: string, actualAnswer: string): Promise<{ pass: boolean; reason: string }> { 96 | const prompt = `You are a deterministic evaluator with zero temperature. Compare the following expected answer with the actual answer and determine if they convey the same information. 97 | 98 | Expected answer: ${expectedAnswer} 99 | Actual answer: ${actualAnswer} 100 | 101 | Minor wording differences are acceptable as long as the core information of the expected answer is preserved in the actual answer.'`; 102 | 103 | const schema = z.object({ 104 | pass: z.boolean().describe('Whether the actual answer matches the expected answer'), 105 | reason: z.string().describe('Detailed explanation of why the evaluation passed or failed') 106 | }); 107 | 108 | try { 109 | const result = await generateObject({ 110 | model: createGoogleGenerativeAI({ apiKey: GEMINI_API_KEY })('gemini-2.0-flash'), // fix to gemini-2.0-flash for evaluation 111 | schema, 112 | prompt, 113 | maxTokens: 1000, 114 | temperature: 0 // Setting temperature to 0 for deterministic output 115 | }); 116 | 117 | return result.object; 118 | } catch (error) { 119 | console.error('Evaluation failed:', error); 120 | return { 121 | pass: false, 122 | reason: `Evaluation error: ${error}` 123 | }; 124 | } 125 | } 126 | 127 | async function batchEvaluate(inputFile: string): Promise { 128 | // Read and parse input file 129 | const questions: Question[] = JSON.parse(await fs.readFile(inputFile, 'utf-8')); 130 | const results: EvaluationResult[] = []; 131 | const gitCommit = await getCurrentGitCommit(); 132 | const modelName = process.env.DEFAULT_MODEL_NAME || 'unknown'; 133 | const outputFile = `eval-${gitCommit}-${modelName}.json`; 134 | 135 | // Process each question 136 | for (let i = 0; i < questions.length; i++) { 137 | const {question, answer: expectedAnswer} = questions[i]; 138 | console.log(`\nProcessing question ${i + 1}/${questions.length}: ${question}`); 139 | 140 | try { 141 | // Get response using the agent 142 | const { 143 | result: response, 144 | context 145 | } = await getResponse(question) as { result: AnswerAction; context: TrackerContext }; 146 | 147 | // Get response using the streaming agent 148 | // const { 149 | // result: response, 150 | // context 151 | // } = await getResponseStreamingAgent(question) as { result: AnswerAction; context: TrackerContext }; 152 | 153 | const actualAnswer = response.answer; 154 | 155 | // Evaluate the response 156 | const evaluation = await evaluateAnswer(expectedAnswer, actualAnswer); 157 | 158 | // Record results 159 | results.push({ 160 | pass: evaluation.pass, 161 | reason: evaluation.reason, 162 | total_steps: context.actionTracker.getState().totalStep, 163 | total_tokens: context.tokenTracker.getTotalUsage().totalTokens, 164 | question, 165 | expected_answer: expectedAnswer, 166 | actual_answer: actualAnswer 167 | }); 168 | 169 | console.log(`Evaluation: ${evaluation.pass ? 'PASS' : 'FAIL'}`); 170 | console.log(`Reason: ${evaluation.reason}`); 171 | } catch (error) { 172 | console.error(`Error processing question: ${question}`, error); 173 | results.push({ 174 | pass: false, 175 | reason: `Error: ${error}`, 176 | total_steps: 0, 177 | total_tokens: 0, 178 | question, 179 | expected_answer: expectedAnswer, 180 | actual_answer: 'Error occurred' 181 | }); 182 | } 183 | } 184 | 185 | // Calculate and print statistics 186 | const stats = calculateStats(results, modelName); 187 | printStats(stats); 188 | 189 | // Save results 190 | await fs.writeFile(outputFile, JSON.stringify({ 191 | results, 192 | statistics: stats 193 | }, null, 2)); 194 | 195 | console.log(`\nEvaluation results saved to ${outputFile}`); 196 | } 197 | 198 | // Run batch evaluation if this is the main module 199 | if (require.main === module) { 200 | const inputFile = process.argv[2]; 201 | if (!inputFile) { 202 | console.error('Please provide an input file path'); 203 | process.exit(1); 204 | } 205 | 206 | batchEvaluate(inputFile).catch(console.error); 207 | } 208 | 209 | export {batchEvaluate}; 210 | -------------------------------------------------------------------------------- /src/evals/ego-questions.json: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "question": "what did jina ai ceo say about deepseek that went viral and become a meme?", 4 | "answer": "a side project" 5 | }, 6 | { 7 | "question": "when was jina ai founded, month and year?", 8 | "answer": "feb 2020" 9 | }, 10 | { 11 | "question": "what is the latest model published by jina ai?", 12 | "answer": "ReaderLM-2.0" 13 | }, 14 | { 15 | "question": "what is the latest blog post that jina ai published?", 16 | "answer": "A Practical Guide to Deploying Search Foundation Models in Production" 17 | }, 18 | { 19 | "question": "what is the context length of readerlm-v2?", 20 | "answer": "512K" 21 | }, 22 | { 23 | "question": "how many employees does jina ai have right now?", 24 | "answer": "30" 25 | }, 26 | { 27 | "question": "when was jina reader api released?", 28 | "answer": "April 2024" 29 | }, 30 | { 31 | "question": "How many offices do Jina AI have and where are they?", 32 | "answer": "four: sunnyvale, berlin, beijing, shenzhen" 33 | }, 34 | { 35 | "question": "what exactly jina-colbert-v2 improves over jina-colbert-v1?", 36 | "answer": "v2 add multilingual support" 37 | }, 38 | { 39 | "question": "who are the authors of jina-clip-v2 paper?", 40 | "answer": "Andreas Koukounas, Georgios Mastrapas, Bo Wang, Mohammad Kalim Akram, Sedigheh Eslami, Michael Günther, Isabelle Mohr, Saba Sturua, Scott Martens, Nan Wang, Han Xiao" 41 | }, 42 | { 43 | "question": "who created the node-deepresearch project?", 44 | "answer": "Han Xiao / jina ai" 45 | }, 46 | { 47 | "question": "Which countries are the investors of Jina AI from?", 48 | "answer": "USA and China only, no German investors" 49 | }, 50 | { 51 | "question": "what is the grounding api endpoint of jina ai?", 52 | "answer": "g.jina.ai" 53 | }, 54 | { 55 | "question": "which of the following models do not support Matryoshka representation? jina-embeddings-v3, jina-embeddings-v2-base-en, jina-clip-v2, jina-clip-v1", 56 | "answer": "jina-embeddings-v2-base-en and jina-clip-v1" 57 | }, 58 | { 59 | "question": "Can I purchase the 2024 yearbook that jina ai published today?", 60 | "answer": "No it is sold out." 61 | }, 62 | { 63 | "question": "How many free tokens do you get from a new jina api key?", 64 | "answer": "1 million." 65 | }, 66 | { 67 | "question": "Who is the legal signatory of Jina AI gmbh?", 68 | "answer": "Jiao Liu" 69 | }, 70 | { 71 | "question": "what is the key idea behind node-deepresearch project?", 72 | "answer": "It keeps searching, reading webpages, reasoning until an answer is found." 73 | }, 74 | { 75 | "question": "what is the name of the jina ai's mascot?", 76 | "answer": "No, Jina AI does not have a mascot." 77 | }, 78 | { 79 | "question": "Does late chunking work with cls pooling?", 80 | "answer": "No. late chunking only works with mean pooling." 81 | } 82 | ] -------------------------------------------------------------------------------- /src/server.ts: -------------------------------------------------------------------------------- 1 | import app from "./app"; 2 | 3 | const port = process.env.PORT || 3000; 4 | 5 | // Export server startup function for better testing 6 | export function startServer() { 7 | return app.listen(port, () => { 8 | console.log(`Server running at http://localhost:${port}`); 9 | }); 10 | } 11 | 12 | // Start server if running directly 13 | if (process.env.NODE_ENV !== 'test') { 14 | startServer(); 15 | } -------------------------------------------------------------------------------- /src/tools/__tests__/error-analyzer.test.ts: -------------------------------------------------------------------------------- 1 | import { analyzeSteps } from '../error-analyzer'; 2 | import { LLMProvider } from '../../config'; 3 | 4 | describe('analyzeSteps', () => { 5 | const providers: Array = ['openai', 'gemini']; 6 | const originalEnv = process.env; 7 | 8 | beforeEach(() => { 9 | jest.resetModules(); 10 | process.env = { ...originalEnv }; 11 | }); 12 | 13 | afterEach(() => { 14 | process.env = originalEnv; 15 | }); 16 | 17 | providers.forEach(provider => { 18 | describe(`with ${provider} provider`, () => { 19 | beforeEach(() => { 20 | process.env.LLM_PROVIDER = provider; 21 | }); 22 | 23 | it('should analyze error steps', async () => { 24 | const { response } = await analyzeSteps(['Step 1: Search failed', 'Step 2: Invalid query']); 25 | expect(response).toHaveProperty('recap'); 26 | expect(response).toHaveProperty('blame'); 27 | expect(response).toHaveProperty('improvement'); 28 | }, 30000); 29 | }); 30 | }); 31 | }); 32 | -------------------------------------------------------------------------------- /src/tools/__tests__/evaluator.test.ts: -------------------------------------------------------------------------------- 1 | import { evaluateAnswer } from '../evaluator'; 2 | import { TokenTracker } from '../../utils/token-tracker'; 3 | import { LLMProvider } from '../../config'; 4 | 5 | describe('evaluateAnswer', () => { 6 | const providers: Array = ['openai', 'gemini']; 7 | const originalEnv = process.env; 8 | 9 | beforeEach(() => { 10 | jest.resetModules(); 11 | process.env = { ...originalEnv }; 12 | }); 13 | 14 | afterEach(() => { 15 | process.env = originalEnv; 16 | }); 17 | 18 | providers.forEach(provider => { 19 | describe(`with ${provider} provider`, () => { 20 | beforeEach(() => { 21 | process.env.LLM_PROVIDER = provider; 22 | }); 23 | 24 | it('should evaluate answer definitiveness', async () => { 25 | const tokenTracker = new TokenTracker(); 26 | const { response } = await evaluateAnswer( 27 | 'What is TypeScript?', 28 | { 29 | action: "answer", 30 | think: "Providing a clear definition of TypeScript", 31 | answer: "TypeScript is a strongly typed programming language that builds on JavaScript.", 32 | references: [] 33 | }, 34 | ['definitive'], 35 | tokenTracker 36 | ); 37 | expect(response).toHaveProperty('pass'); 38 | expect(response).toHaveProperty('think'); 39 | expect(response.type).toBe('definitive'); 40 | }); 41 | 42 | it('should evaluate answer plurality', async () => { 43 | const tokenTracker = new TokenTracker(); 44 | const { response } = await evaluateAnswer( 45 | 'List three programming languages.', 46 | { 47 | action: "answer", 48 | think: "Providing an example of a programming language", 49 | answer: "Python is a programming language.", 50 | references: [] 51 | }, 52 | ['plurality'], 53 | tokenTracker 54 | ); 55 | expect(response).toHaveProperty('pass'); 56 | expect(response).toHaveProperty('think'); 57 | expect(response.type).toBe('plurality'); 58 | expect(response.plurality_analysis?.expects_multiple).toBe(true); 59 | }); 60 | }); 61 | }); 62 | }); 63 | -------------------------------------------------------------------------------- /src/tools/__tests__/read.test.ts: -------------------------------------------------------------------------------- 1 | import { readUrl } from '../read'; 2 | import { TokenTracker } from '../../utils/token-tracker'; 3 | 4 | describe('readUrl', () => { 5 | it.skip('should read and parse URL content (skipped due to insufficient balance)', async () => { 6 | const tokenTracker = new TokenTracker(); 7 | const { response } = await readUrl('https://www.typescriptlang.org', tokenTracker); 8 | expect(response).toHaveProperty('code'); 9 | expect(response).toHaveProperty('status'); 10 | expect(response.data).toHaveProperty('content'); 11 | expect(response.data).toHaveProperty('title'); 12 | }, 15000); 13 | 14 | it.skip('should handle invalid URLs (skipped due to insufficient balance)', async () => { 15 | await expect(readUrl('invalid-url')).rejects.toThrow(); 16 | }, 15000); 17 | 18 | beforeEach(() => { 19 | jest.setTimeout(15000); 20 | }); 21 | }); 22 | -------------------------------------------------------------------------------- /src/tools/__tests__/search.test.ts: -------------------------------------------------------------------------------- 1 | import { search } from '../jina-search'; 2 | import { TokenTracker } from '../../utils/token-tracker'; 3 | 4 | describe('search', () => { 5 | it.skip('should perform search with Jina API (skipped due to insufficient balance)', async () => { 6 | const tokenTracker = new TokenTracker(); 7 | const { response } = await search('TypeScript programming', tokenTracker); 8 | expect(response).toBeDefined(); 9 | expect(response.data).toBeDefined(); 10 | if (response.data === null) { 11 | throw new Error('Response data is null'); 12 | } 13 | expect(Array.isArray(response.data)).toBe(true); 14 | expect(response.data.length).toBeGreaterThan(0); 15 | }, 15000); 16 | 17 | it('should handle empty query', async () => { 18 | await expect(search('')).rejects.toThrow(); 19 | }, 15000); 20 | 21 | beforeEach(() => { 22 | jest.setTimeout(15000); 23 | }); 24 | }); 25 | -------------------------------------------------------------------------------- /src/tools/brave-search.ts: -------------------------------------------------------------------------------- 1 | import {BRAVE_API_KEY} from "../config"; 2 | import axiosClient from "../utils/axios-client"; 3 | 4 | import { BraveSearchResponse } from '../types'; 5 | 6 | export async function braveSearch(query: string): Promise<{ response: BraveSearchResponse }> { 7 | const response = await axiosClient.get('https://api.search.brave.com/res/v1/web/search', { 8 | params: { 9 | q: query, 10 | count: 10, 11 | safesearch: 'off' 12 | }, 13 | headers: { 14 | 'Accept': 'application/json', 15 | 'X-Subscription-Token': BRAVE_API_KEY 16 | }, 17 | timeout: 10000 18 | }); 19 | 20 | // Maintain the same return structure as the original code 21 | return { response: response.data }; 22 | } 23 | -------------------------------------------------------------------------------- /src/tools/broken-ch-fixer.ts: -------------------------------------------------------------------------------- 1 | import { generateText } from "ai"; 2 | import { getModel } from "../config"; 3 | import {TrackerContext} from "../types"; 4 | import {detectBrokenUnicodeViaFileIO} from "../utils/text-tools"; 5 | 6 | 7 | /** 8 | * Repairs markdown content with � characters by using Gemini to guess the missing text 9 | */ 10 | export async function repairUnknownChars(mdContent: string, trackers?: TrackerContext): Promise { 11 | const { broken, readStr } = await detectBrokenUnicodeViaFileIO(mdContent); 12 | if (!broken) return readStr; 13 | console.log("Detected broken unicode in output, attempting to repair..."); 14 | 15 | let repairedContent = readStr; 16 | let remainingUnknowns = true; 17 | let iterations = 0; 18 | 19 | let lastPosition = -1; 20 | 21 | while (remainingUnknowns && iterations < 20) { 22 | iterations++; 23 | 24 | // Find the position of the first � character 25 | const position = repairedContent.indexOf('�'); 26 | if (position === -1) { 27 | remainingUnknowns = false; 28 | continue; 29 | } 30 | 31 | // Check if we're stuck at the same position 32 | if (position === lastPosition) { 33 | // Move past this character by removing it 34 | repairedContent = repairedContent.substring(0, position) + 35 | repairedContent.substring(position + 1); 36 | continue; 37 | } 38 | 39 | // Update last position to detect loops 40 | lastPosition = position; 41 | 42 | // Count consecutive � characters 43 | let unknownCount = 0; 44 | for (let i = position; i < repairedContent.length && repairedContent[i] === '�'; i++) { 45 | unknownCount++; 46 | } 47 | 48 | // Extract context around the unknown characters 49 | const contextSize = 50; 50 | const start = Math.max(0, position - contextSize); 51 | const end = Math.min(repairedContent.length, position + unknownCount + contextSize); 52 | const leftContext = repairedContent.substring(start, position); 53 | const rightContext = repairedContent.substring(position + unknownCount, end); 54 | 55 | // Ask Gemini to guess the missing characters 56 | try { 57 | const result = await generateText({ 58 | model: getModel('fallback'), 59 | system: `You're helping fix a corrupted scanned markdown document that has stains (represented by �). 60 | Looking at the surrounding context, determine the original text should be in place of the � symbols. 61 | 62 | Rules: 63 | 1. ONLY output the exact replacement text - no explanations, quotes, or additional text 64 | 2. Keep your response appropriate to the length of the unknown sequence 65 | 3. Consider the document appears to be in Chinese if that's what the context suggests`, 66 | prompt: ` 67 | The corrupted text has ${unknownCount} � mush in a row. 68 | 69 | On the left of the stains: "${leftContext}" 70 | On the right of the stains: "${rightContext}" 71 | 72 | So what was the original text between these two contexts?`, 73 | }); 74 | 75 | trackers?.tokenTracker.trackUsage('md-fixer', result.usage) 76 | const replacement = result.text.trim(); 77 | 78 | // Validate the replacement 79 | if ( 80 | replacement === "UNKNOWN" || 81 | (await detectBrokenUnicodeViaFileIO(replacement)).broken || 82 | replacement.length > unknownCount * 4 83 | ) { 84 | console.log(`Skipping invalid replacement ${replacement} at position ${position}`); 85 | // Skip to the next � character without modifying content 86 | } else { 87 | // Replace the unknown sequence with the generated text 88 | repairedContent = repairedContent.substring(0, position) + 89 | replacement + 90 | repairedContent.substring(position + unknownCount); 91 | } 92 | 93 | console.log(`Repair iteration ${iterations}: replaced ${unknownCount} � chars with "${replacement}"`); 94 | 95 | } catch (error) { 96 | console.error("Error repairing unknown characters:", error); 97 | // Skip to the next � character without modifying this one 98 | } 99 | } 100 | 101 | return repairedContent; 102 | } -------------------------------------------------------------------------------- /src/tools/code-sandbox.ts: -------------------------------------------------------------------------------- 1 | import {ObjectGeneratorSafe} from "../utils/safe-generator"; 2 | import {CodeGenResponse, PromptPair, TrackerContext} from "../types"; 3 | import {Schemas} from "../utils/schemas"; 4 | 5 | 6 | interface SandboxResult { 7 | success: boolean; 8 | output?: any; 9 | error?: string; 10 | } 11 | 12 | 13 | function getPrompt( 14 | problem: string, 15 | availableVars: string, 16 | previousAttempts: Array<{ code: string; error?: string }> = [] 17 | ): PromptPair { 18 | const previousAttemptsContext = previousAttempts.map((attempt, index) => ` 19 | 20 | ${attempt.code} 21 | ${attempt.error ? `Error: ${attempt.error} 22 | 23 | ` : ''} 24 | `).join('\n'); 25 | 26 | const prompt = `You are an expert JavaScript programmer. Your task is to generate JavaScript code to solve the given problem. 27 | 28 | 29 | 1. Generate plain JavaScript code that returns the result directly 30 | 2. You can access any of these available variables directly: 31 | ${availableVars} 32 | 3. You don't have access to any third party libraries that need to be installed, so you must write complete, self-contained code. 33 | 34 | 35 | ${previousAttempts.length > 0 ? `Previous attempts and their errors: 36 | ${previousAttemptsContext} 37 | ` : ''} 38 | 39 | 40 | Available variables: 41 | numbers (Array) e.g. [1, 2, 3, 4, 5, 6] 42 | threshold (number) e.g. 4 43 | 44 | Problem: Sum all numbers above threshold 45 | 46 | Response: 47 | { 48 | "code": "return numbers.filter(n => n > threshold).reduce((a, b) => a + b, 0);" 49 | } 50 | `; 51 | 52 | console.log('Coding prompt', prompt) 53 | 54 | return {system: prompt, user: problem }; 55 | } 56 | 57 | export class CodeSandbox { 58 | private trackers?: TrackerContext; 59 | private generator: ObjectGeneratorSafe; 60 | private maxAttempts: number; 61 | private context: Record; 62 | private schemaGen: Schemas; 63 | 64 | constructor( 65 | context: any = {}, 66 | trackers: TrackerContext, 67 | schemaGen: Schemas, 68 | maxAttempts: number = 3, 69 | ) { 70 | this.trackers = trackers; 71 | this.generator = new ObjectGeneratorSafe(trackers?.tokenTracker); 72 | this.maxAttempts = maxAttempts; 73 | this.context = context; 74 | this.schemaGen = schemaGen; 75 | } 76 | 77 | private async generateCode( 78 | problem: string, 79 | previousAttempts: Array<{ code: string; error?: string }> = [] 80 | ): Promise { 81 | const prompt = getPrompt(problem, analyzeStructure(this.context), previousAttempts); 82 | 83 | const result = await this.generator.generateObject({ 84 | model: 'coder', 85 | schema: this.schemaGen.getCodeGeneratorSchema(), 86 | system: prompt.system, 87 | prompt: prompt.user 88 | }); 89 | 90 | this.trackers?.actionTracker.trackThink(result.object.think); 91 | 92 | return result.object as CodeGenResponse; 93 | } 94 | 95 | private evaluateCode(code: string): SandboxResult { 96 | try { 97 | // Create a function that uses 'with' to evaluate in the context and return the result 98 | const evalInContext = new Function('context', ` 99 | with (context) { 100 | ${code} 101 | } 102 | `); 103 | 104 | console.log('Context:', this.context); 105 | 106 | // Execute the code with the context and get the return value 107 | const output = evalInContext(this.context); 108 | 109 | if (output === undefined) { 110 | return { 111 | success: false, 112 | error: 'No value was returned, make sure to use "return" statement to return the result' 113 | }; 114 | } 115 | 116 | return { 117 | success: true, 118 | output 119 | }; 120 | } catch (error) { 121 | return { 122 | success: false, 123 | error: error instanceof Error ? error.message : 'Unknown error occurred' 124 | }; 125 | } 126 | } 127 | 128 | async solve(problem: string): Promise<{ 129 | solution: { code: string; output: any }; 130 | attempts: Array<{ code: string; error?: string }>; 131 | }> { 132 | const attempts: Array<{ code: string; error?: string }> = []; 133 | 134 | for (let i = 0; i < this.maxAttempts; i++) { 135 | // Generate code 136 | const generation = await this.generateCode(problem, attempts); 137 | const {code} = generation; 138 | 139 | console.log(`Coding attempt ${i + 1}:`, code); 140 | // Evaluate the code 141 | const result = this.evaluateCode(code); 142 | console.log(`Coding attempt ${i + 1} success:`, result); 143 | 144 | if (result.success) { 145 | return { 146 | solution: { 147 | code, 148 | output: result.output 149 | }, 150 | attempts 151 | }; 152 | } 153 | 154 | console.error('Coding error:', result.error); 155 | 156 | // Store the failed attempt 157 | attempts.push({ 158 | code, 159 | error: result.error 160 | }); 161 | 162 | // If we've reached max attempts, throw an error 163 | if (i === this.maxAttempts - 1) { 164 | throw new Error(`Failed to generate working code after ${this.maxAttempts} attempts`); 165 | } 166 | } 167 | 168 | // This should never be reached due to the throw above 169 | throw new Error('Unexpected end of execution'); 170 | } 171 | } 172 | 173 | function formatValue(value: any): string { 174 | if (value === null) return 'null'; 175 | if (value === undefined) return 'undefined'; 176 | 177 | const type = typeof value; 178 | 179 | if (type === 'string') { 180 | // Clean and truncate string value 181 | const cleaned = value.replace(/\n/g, ' ').replace(/\s+/g, ' ').trim(); 182 | return cleaned.length > 50 ? 183 | `"${cleaned.slice(0, 47)}..."` : 184 | `"${cleaned}"`; 185 | } 186 | 187 | if (type === 'number' || type === 'boolean') { 188 | return String(value); 189 | } 190 | 191 | if (value instanceof Date) { 192 | return `"${value.toISOString()}"`; 193 | } 194 | 195 | return ''; 196 | } 197 | 198 | export function analyzeStructure(value: any, indent = ''): string { 199 | if (value === null) return 'null'; 200 | if (value === undefined) return 'undefined'; 201 | 202 | const type = typeof value; 203 | 204 | if (type === 'function') { 205 | return 'Function'; 206 | } 207 | 208 | // Handle atomic types with example values 209 | if (type !== 'object' || value instanceof Date) { 210 | const formattedValue = formatValue(value); 211 | return `${type}${formattedValue ? ` (example: ${formattedValue})` : ''}`; 212 | } 213 | 214 | if (Array.isArray(value)) { 215 | if (value.length === 0) return 'Array'; 216 | const sampleItem = value[0]; 217 | return `Array<${analyzeStructure(sampleItem, indent + ' ')}>`; 218 | } 219 | 220 | const entries = Object.entries(value); 221 | if (entries.length === 0) return '{}'; 222 | 223 | const properties = entries 224 | .map(([key, val]) => { 225 | const analyzed = analyzeStructure(val, indent + ' '); 226 | return `${indent} "${key}": ${analyzed}`; 227 | }) 228 | .join(',\n'); 229 | 230 | return `{\n${properties}\n${indent}}`; 231 | } -------------------------------------------------------------------------------- /src/tools/cosine.ts: -------------------------------------------------------------------------------- 1 | export function cosineSimilarity(vecA: number[], vecB: number[]): number { 2 | if (vecA.length !== vecB.length) { 3 | throw new Error("Vectors must have the same length"); 4 | } 5 | 6 | let dotProduct = 0; 7 | let magnitudeA = 0; 8 | let magnitudeB = 0; 9 | 10 | for (let i = 0; i < vecA.length; i++) { 11 | dotProduct += vecA[i] * vecB[i]; 12 | magnitudeA += vecA[i] * vecA[i]; 13 | magnitudeB += vecB[i] * vecB[i]; 14 | } 15 | 16 | magnitudeA = Math.sqrt(magnitudeA); 17 | magnitudeB = Math.sqrt(magnitudeB); 18 | 19 | return (magnitudeA > 0 && magnitudeB > 0) ? dotProduct / (magnitudeA * magnitudeB) : 0; 20 | } 21 | 22 | // Fallback similarity ranking using Jaccard 23 | export async function jaccardRank(query: string, documents: string[]): Promise<{ results: { index: number, relevance_score: number }[] }> { 24 | console.log(`[fallback] Using Jaccard similarity for ${documents.length} documents`); 25 | // Convert texts to lowercase and tokenize by splitting on non-alphanumeric characters 26 | const queryTokens = new Set(query.toLowerCase().split(/\W+/).filter(t => t.length > 0)); 27 | 28 | const results = documents.map((doc, index) => { 29 | const docTokens = new Set(doc.toLowerCase().split(/\W+/).filter(t => t.length > 0)); 30 | 31 | // Calculate intersection size 32 | const intersection = new Set([...queryTokens].filter(x => docTokens.has(x))); 33 | 34 | // Calculate union size 35 | const union = new Set([...queryTokens, ...docTokens]); 36 | 37 | // Calculate Jaccard similarity 38 | const score = union.size === 0 ? 0 : intersection.size / union.size; 39 | 40 | return {index, relevance_score: score}; 41 | }); 42 | 43 | // Sort by score in descending order 44 | results.sort((a, b) => b.relevance_score - a.relevance_score); 45 | 46 | return {results}; 47 | } 48 | -------------------------------------------------------------------------------- /src/tools/dedup.ts: -------------------------------------------------------------------------------- 1 | import {z} from 'zod'; 2 | import {TokenTracker} from "../utils/token-tracker"; 3 | import {ObjectGeneratorSafe} from "../utils/safe-generator"; 4 | 5 | 6 | const responseSchema = z.object({ 7 | think: z.string().describe('Strategic reasoning about the overall deduplication approach').max(500), 8 | unique_queries: z.array(z.string().describe('Unique query that passed the deduplication process, must be less than 30 characters')) 9 | .describe('Array of semantically unique queries').max(3) 10 | }); 11 | 12 | function getPrompt(newQueries: string[], existingQueries: string[]): string { 13 | return `You are an expert in semantic similarity analysis. Given a set of queries (setA) and a set of queries (setB) 14 | 15 | 16 | Function FilterSetA(setA, setB, threshold): 17 | filteredA = empty set 18 | 19 | for each candidateQuery in setA: 20 | isValid = true 21 | 22 | // Check similarity with already accepted queries in filteredA 23 | for each acceptedQuery in filteredA: 24 | similarity = calculateSimilarity(candidateQuery, acceptedQuery) 25 | if similarity >= threshold: 26 | isValid = false 27 | break 28 | 29 | // If passed first check, compare with set B 30 | if isValid: 31 | for each queryB in setB: 32 | similarity = calculateSimilarity(candidateQuery, queryB) 33 | if similarity >= threshold: 34 | isValid = false 35 | break 36 | 37 | // If passed all checks, add to filtered set 38 | if isValid: 39 | add candidateQuery to filteredA 40 | 41 | return filteredA 42 | 43 | 44 | 45 | 1. Consider semantic meaning and query intent, not just lexical similarity 46 | 2. Account for different phrasings of the same information need 47 | 3. Queries with same base keywords but different operators are NOT duplicates 48 | 4. Different aspects or perspectives of the same topic are not duplicates 49 | 5. Consider query specificity - a more specific query is not a duplicate of a general one 50 | 6. Search operators that make queries behave differently: 51 | - Different site: filters (e.g., site:youtube.com vs site:github.com) 52 | - Different file types (e.g., filetype:pdf vs filetype:doc) 53 | - Different language/location filters (e.g., lang:en vs lang:es) 54 | - Different exact match phrases (e.g., "exact phrase" vs no quotes) 55 | - Different inclusion/exclusion (+/- operators) 56 | - Different title/body filters (intitle: vs inbody:) 57 | 58 | 59 | Now with threshold set to 0.2; run FilterSetA on the following: 60 | SetA: ${JSON.stringify(newQueries)} 61 | SetB: ${JSON.stringify(existingQueries)}`; 62 | } 63 | 64 | 65 | const TOOL_NAME = 'dedup'; 66 | 67 | export async function dedupQueries( 68 | newQueries: string[], 69 | existingQueries: string[], 70 | tracker?: TokenTracker 71 | ): Promise<{ unique_queries: string[] }> { 72 | try { 73 | const generator = new ObjectGeneratorSafe(tracker); 74 | const prompt = getPrompt(newQueries, existingQueries); 75 | 76 | const result = await generator.generateObject({ 77 | model: TOOL_NAME, 78 | schema: responseSchema, 79 | prompt, 80 | }); 81 | 82 | console.log(TOOL_NAME, result.object.unique_queries); 83 | return {unique_queries: result.object.unique_queries}; 84 | 85 | } catch (error) { 86 | console.error(`Error in ${TOOL_NAME}`, error); 87 | throw error; 88 | } 89 | } -------------------------------------------------------------------------------- /src/tools/embeddings.ts: -------------------------------------------------------------------------------- 1 | import {JINA_API_KEY} from "../config"; 2 | import {JinaEmbeddingRequest, JinaEmbeddingResponse} from "../types"; 3 | import axiosClient from "../utils/axios-client"; 4 | 5 | const BATCH_SIZE = 128; 6 | const API_URL = "https://api.jina.ai/v1/embeddings"; 7 | const MAX_RETRIES = 3; // Maximum number of retries for missing embeddings 8 | 9 | // Modified to support different embedding tasks and dimensions 10 | export async function getEmbeddings( 11 | texts: string[], 12 | tokenTracker?: any, 13 | options: { 14 | task?: "text-matching" | "retrieval.passage" | "retrieval.query", 15 | dimensions?: number, 16 | late_chunking?: boolean, 17 | embedding_type?: string 18 | } = {} 19 | ): Promise<{ embeddings: number[][], tokens: number }> { 20 | console.log(`[embeddings] Getting embeddings for ${texts.length} texts`); 21 | 22 | if (!JINA_API_KEY) { 23 | throw new Error('JINA_API_KEY is not set'); 24 | } 25 | 26 | // Handle empty input case 27 | if (texts.length === 0) { 28 | return {embeddings: [], tokens: 0}; 29 | } 30 | 31 | // Process in batches 32 | const allEmbeddings: number[][] = []; 33 | let totalTokens = 0; 34 | const batchCount = Math.ceil(texts.length / BATCH_SIZE); 35 | 36 | for (let i = 0; i < texts.length; i += BATCH_SIZE) { 37 | const batchTexts = texts.slice(i, i + BATCH_SIZE); 38 | const currentBatch = Math.floor(i / BATCH_SIZE) + 1; 39 | console.log(`[embeddings] Processing batch ${currentBatch}/${batchCount} (${batchTexts.length} texts)`); 40 | 41 | // Get embeddings for the batch with retry logic for missing indices 42 | const { batchEmbeddings, batchTokens } = await getBatchEmbeddingsWithRetry( 43 | batchTexts, 44 | options, 45 | currentBatch, 46 | batchCount 47 | ); 48 | 49 | allEmbeddings.push(...batchEmbeddings); 50 | totalTokens += batchTokens; 51 | console.log(`[embeddings] Batch ${currentBatch} complete. Tokens used: ${batchTokens}, total so far: ${totalTokens}`); 52 | } 53 | 54 | // Track token usage if tracker is provided 55 | if (tokenTracker) { 56 | tokenTracker.trackUsage('embeddings', { 57 | promptTokens: totalTokens, 58 | completionTokens: 0, 59 | totalTokens: totalTokens 60 | }); 61 | } 62 | 63 | console.log(`[embeddings] Complete. Generated ${allEmbeddings.length} embeddings using ${totalTokens} tokens`); 64 | return {embeddings: allEmbeddings, tokens: totalTokens}; 65 | } 66 | 67 | // Helper function to get embeddings for a batch with retry logic for missing indices 68 | async function getBatchEmbeddingsWithRetry( 69 | batchTexts: string[], 70 | options: { 71 | task?: "text-matching" | "retrieval.passage" | "retrieval.query", 72 | dimensions?: number, 73 | late_chunking?: boolean, 74 | embedding_type?: string 75 | }, 76 | currentBatch: number, 77 | batchCount: number 78 | ): Promise<{ batchEmbeddings: number[][], batchTokens: number }> { 79 | const batchEmbeddings: number[][] = []; 80 | let batchTokens = 0; 81 | let retryCount = 0; 82 | let textsToProcess = [...batchTexts]; // Copy the original texts 83 | let indexMap = new Map(); // Map to keep track of original indices 84 | 85 | // Initialize indexMap with original indices 86 | textsToProcess.forEach((_, idx) => { 87 | indexMap.set(idx, idx); 88 | }); 89 | 90 | while (textsToProcess.length > 0 && retryCount < MAX_RETRIES) { 91 | const request: JinaEmbeddingRequest = { 92 | model: "jina-embeddings-v3", 93 | task: options.task || "text-matching", 94 | input: textsToProcess, 95 | truncate: true, 96 | }; 97 | 98 | // Add optional parameters if provided 99 | if (options.dimensions) request.dimensions = options.dimensions; 100 | if (options.late_chunking) request.late_chunking = options.late_chunking; 101 | if (options.embedding_type) request.embedding_type = options.embedding_type; 102 | 103 | try { 104 | const response = await axiosClient.post( 105 | API_URL, 106 | request, 107 | { 108 | headers: { 109 | "Content-Type": "application/json", 110 | "Authorization": `Bearer ${JINA_API_KEY}` 111 | } 112 | } 113 | ); 114 | 115 | if (!response.data.data) { 116 | console.error('No data returned from Jina API'); 117 | if (retryCount === MAX_RETRIES - 1) { 118 | // On last retry, create placeholder embeddings 119 | const dimensionSize = options.dimensions || 1024; 120 | const placeholderEmbeddings = textsToProcess.map(text => { 121 | console.error(`Failed to get embedding after all retries: [${text.substring(0, 50)}...]`); 122 | return new Array(dimensionSize).fill(0); 123 | }); 124 | 125 | // Add embeddings in correct order 126 | for (let i = 0; i < textsToProcess.length; i++) { 127 | const originalIndex = indexMap.get(i)!; 128 | while (batchEmbeddings.length <= originalIndex) { 129 | batchEmbeddings.push([]); 130 | } 131 | batchEmbeddings[originalIndex] = placeholderEmbeddings[i]; 132 | } 133 | } 134 | retryCount++; 135 | continue; 136 | } 137 | 138 | const receivedIndices = new Set(response.data.data.map(item => item.index)); 139 | const dimensionSize = response.data.data[0]?.embedding?.length || options.dimensions || 1024; 140 | 141 | // Process successful embeddings 142 | const successfulEmbeddings: number[][] = []; 143 | const remainingTexts: string[] = []; 144 | const newIndexMap = new Map(); 145 | 146 | for (let idx = 0; idx < textsToProcess.length; idx++) { 147 | if (receivedIndices.has(idx)) { 148 | // Find the item with this index 149 | const item = response.data.data.find(d => d.index === idx)!; 150 | 151 | // Get the original index and store in the result array 152 | const originalIndex = indexMap.get(idx)!; 153 | while (batchEmbeddings.length <= originalIndex) { 154 | batchEmbeddings.push([]); 155 | } 156 | batchEmbeddings[originalIndex] = item.embedding; 157 | successfulEmbeddings.push(item.embedding); 158 | } else { 159 | // Add to retry list 160 | const newIndex = remainingTexts.length; 161 | newIndexMap.set(newIndex, indexMap.get(idx)!); 162 | remainingTexts.push(textsToProcess[idx]); 163 | console.log(`Missing embedding for index ${idx}, will retry: [${textsToProcess[idx].substring(0, 50)}...]`); 164 | } 165 | } 166 | 167 | // Add tokens 168 | batchTokens += response.data.usage?.total_tokens || 0; 169 | 170 | // Update for next iteration 171 | textsToProcess = remainingTexts; 172 | indexMap = newIndexMap; 173 | 174 | // If all embeddings were successfully processed, break out of the loop 175 | if (textsToProcess.length === 0) { 176 | break; 177 | } 178 | 179 | // Increment retry count and log 180 | retryCount++; 181 | console.log(`[embeddings] Batch ${currentBatch}/${batchCount} - Retrying ${textsToProcess.length} texts (attempt ${retryCount}/${MAX_RETRIES})`); 182 | } catch (error: any) { 183 | console.error('Error calling Jina Embeddings API:', error); 184 | if (error.response?.status === 402 || error.message.includes('InsufficientBalanceError') || error.message.includes('insufficient balance')) { 185 | return { batchEmbeddings: [], batchTokens: 0 }; 186 | } 187 | 188 | // On last retry, create placeholder embeddings 189 | if (retryCount === MAX_RETRIES - 1) { 190 | const dimensionSize = options.dimensions || 1024; 191 | for (let idx = 0; idx < textsToProcess.length; idx++) { 192 | const originalIndex = indexMap.get(idx)!; 193 | console.error(`Failed to get embedding after all retries for index ${originalIndex}: [${textsToProcess[idx].substring(0, 50)}...]`); 194 | 195 | while (batchEmbeddings.length <= originalIndex) { 196 | batchEmbeddings.push([]); 197 | } 198 | batchEmbeddings[originalIndex] = new Array(dimensionSize).fill(0); 199 | } 200 | } 201 | 202 | retryCount++; 203 | if (retryCount < MAX_RETRIES) { 204 | console.log(`[embeddings] Batch ${currentBatch}/${batchCount} - Retry attempt ${retryCount}/${MAX_RETRIES} after error`); 205 | // Wait before retrying to avoid overwhelming the API 206 | await new Promise(resolve => setTimeout(resolve, 1000)); 207 | } else { 208 | throw error; // If we've exhausted retries, re-throw the error 209 | } 210 | } 211 | } 212 | 213 | // Handle any remaining missing embeddings after max retries 214 | if (textsToProcess.length > 0) { 215 | console.error(`[embeddings] Failed to get embeddings for ${textsToProcess.length} texts after ${MAX_RETRIES} retries`); 216 | const dimensionSize = options.dimensions || 1024; 217 | 218 | for (let idx = 0; idx < textsToProcess.length; idx++) { 219 | const originalIndex = indexMap.get(idx)!; 220 | console.error(`Creating zero embedding for index ${originalIndex} after all retries failed`); 221 | 222 | while (batchEmbeddings.length <= originalIndex) { 223 | batchEmbeddings.push([]); 224 | } 225 | batchEmbeddings[originalIndex] = new Array(dimensionSize).fill(0); 226 | } 227 | } 228 | 229 | return { batchEmbeddings, batchTokens }; 230 | } 231 | -------------------------------------------------------------------------------- /src/tools/error-analyzer.ts: -------------------------------------------------------------------------------- 1 | import {ErrorAnalysisResponse, PromptPair, TrackerContext} from '../types'; 2 | import {ObjectGeneratorSafe} from "../utils/safe-generator"; 3 | import {Schemas} from "../utils/schemas"; 4 | 5 | 6 | function getPrompt(diaryContext: string[]): PromptPair { 7 | return { 8 | system: `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. 9 | 10 | 11 | 1. The sequence of actions taken 12 | 2. The effectiveness of each step 13 | 3. The logic between consecutive steps 14 | 4. Alternative approaches that could have been taken 15 | 5. Signs of getting stuck in repetitive patterns 16 | 6. Whether the final answer matches the accumulated information 17 | 18 | Analyze the steps and provide detailed feedback following these guidelines: 19 | - In the recap: Summarize key actions chronologically, highlight patterns, and identify where the process started to go wrong 20 | - In the blame: Point to specific steps or patterns that led to the inadequate answer 21 | - In the improvement: Provide actionable suggestions that could have led to a better outcome 22 | 23 | 24 | 25 | 26 | 27 | 28 | At step 1, you took the **search** action and look for external information for the question: "how old is jina ai ceo?". 29 | In particular, you tried to search for the following keywords: "jina ai ceo age". 30 | You found quite some information and add them to your URL list and **visit** them later when needed. 31 | 32 | 33 | At step 2, you took the **visit** action and deep dive into the following URLs: 34 | https://www.linkedin.com/in/hxiao87 35 | https://www.crunchbase.com/person/han-xiao 36 | You found some useful information on the web and add them to your knowledge for future reference. 37 | 38 | 39 | At step 3, you took the **search** action and look for external information for the question: "how old is jina ai ceo?". 40 | In particular, you tried to search for the following keywords: "Han Xiao birthdate, Jina AI founder birthdate". 41 | You found quite some information and add them to your URL list and **visit** them later when needed. 42 | 43 | 44 | At step 4, you took the **search** action and look for external information for the question: "how old is jina ai ceo?". 45 | In particular, you tried to search for the following keywords: han xiao birthday. 46 | But then you realized you have already searched for these keywords before. 47 | You decided to think out of the box or cut from a completely different angle. 48 | 49 | 50 | At step 5, you took the **search** action and look for external information for the question: "how old is jina ai ceo?". 51 | In particular, you tried to search for the following keywords: han xiao birthday. 52 | But then you realized you have already searched for these keywords before. 53 | You decided to think out of the box or cut from a completely different angle. 54 | 55 | 56 | At step 6, you took the **visit** action and deep dive into the following URLs: 57 | https://kpopwall.com/han-xiao/ 58 | https://www.idolbirthdays.net/han-xiao 59 | You found some useful information on the web and add them to your knowledge for future reference. 60 | 61 | 62 | At step 7, you took **answer** action but evaluator thinks it is not a good answer: 63 | 64 | 65 | 66 | Original question: 67 | how old is jina ai ceo? 68 | 69 | Your answer: 70 | The age of the Jina AI CEO cannot be definitively determined from the provided information. 71 | 72 | The evaluator thinks your answer is bad because: 73 | The answer is not definitive and fails to provide the requested information. Lack of information is unacceptable, more search and deep reasoning is needed. 74 | 75 | 76 | 77 | 78 | { 79 | "recap": "The search process consisted of 7 steps with multiple search and visit actions. The initial searches focused on basic biographical information through LinkedIn and Crunchbase (steps 1-2). When this didn't yield the specific age information, additional searches were conducted for birthdate information (steps 3-5). The process showed signs of repetition in steps 4-5 with identical searches. Final visits to entertainment websites (step 6) suggested a loss of focus on reliable business sources.", 80 | 81 | "blame": "The root cause of failure was getting stuck in a repetitive search pattern without adapting the strategy. Steps 4-5 repeated the same search, and step 6 deviated to less reliable entertainment sources instead of exploring business journals, news articles, or professional databases. Additionally, the process didn't attempt to triangulate age through indirect information like education history or career milestones.", 82 | 83 | "improvement": "1. Avoid repeating identical searches and implement a strategy to track previously searched terms. 2. When direct age/birthdate searches fail, try indirect approaches like: searching for earliest career mentions, finding university graduation years, or identifying first company founding dates. 3. Focus on high-quality business sources and avoid entertainment websites for professional information. 4. Consider using industry event appearances or conference presentations where age-related context might be mentioned. 5. If exact age cannot be determined, provide an estimated range based on career timeline and professional achievements.", 84 | 85 | } 86 | 87 | `, 88 | user: `${diaryContext.join('\n')}` 89 | } 90 | } 91 | 92 | const TOOL_NAME = 'errorAnalyzer'; 93 | 94 | export async function analyzeSteps( 95 | diaryContext: string[], 96 | trackers: TrackerContext, 97 | schemaGen: Schemas 98 | ): Promise { 99 | try { 100 | const generator = new ObjectGeneratorSafe(trackers?.tokenTracker); 101 | const prompt = getPrompt(diaryContext); 102 | 103 | const result = await generator.generateObject({ 104 | model: TOOL_NAME, 105 | schema: schemaGen.getErrorAnalysisSchema(), 106 | system: prompt.system, 107 | prompt: prompt.user 108 | }); 109 | 110 | console.log(TOOL_NAME, result.object); 111 | trackers?.actionTracker.trackThink(result.object.blame); 112 | trackers?.actionTracker.trackThink(result.object.improvement); 113 | 114 | return result.object as ErrorAnalysisResponse; 115 | 116 | } catch (error) { 117 | console.error(`Error in ${TOOL_NAME}`, error); 118 | throw error; 119 | } 120 | } -------------------------------------------------------------------------------- /src/tools/grounding.ts: -------------------------------------------------------------------------------- 1 | import { generateText } from 'ai'; 2 | import {getModel} from "../config"; 3 | import { GoogleGenerativeAIProviderMetadata } from '@ai-sdk/google'; 4 | import {TokenTracker} from "../utils/token-tracker"; 5 | 6 | const model = getModel('searchGrounding') 7 | 8 | export async function grounding(query: string, tracker?: TokenTracker): Promise { 9 | try { 10 | const { text, experimental_providerMetadata, usage } = await generateText({ 11 | model, 12 | prompt: 13 | `Current date is ${new Date().toISOString()}. Find the latest answer to the following question: 14 | 15 | ${query} 16 | 17 | Must include the date and time of the latest answer.`, 18 | }); 19 | 20 | const metadata = experimental_providerMetadata?.google as 21 | | GoogleGenerativeAIProviderMetadata 22 | | undefined; 23 | const groundingMetadata = metadata?.groundingMetadata; 24 | 25 | // Extract and concatenate all groundingSupport text into a single line 26 | const groundedText = groundingMetadata?.groundingSupports 27 | ?.map(support => support.segment.text) 28 | .join(' ') || ''; 29 | 30 | (tracker || new TokenTracker()).trackUsage('grounding', usage); 31 | console.log('Grounding:', {text, groundedText}); 32 | return text + '|' + groundedText; 33 | 34 | } catch (error) { 35 | console.error('Error in search:', error); 36 | throw error; 37 | } 38 | } -------------------------------------------------------------------------------- /src/tools/jina-classify-spam.ts: -------------------------------------------------------------------------------- 1 | import { TokenTracker } from "../utils/token-tracker"; 2 | import { JINA_API_KEY } from "../config"; 3 | import axiosClient from "../utils/axios-client"; 4 | 5 | const JINA_API_URL = 'https://api.jina.ai/v1/classify'; 6 | 7 | // Types for Jina Classification API 8 | interface JinaClassifyRequest { 9 | classifier_id: string; 10 | input: string[]; 11 | } 12 | 13 | interface JinaClassifyResponse { 14 | usage: { 15 | total_tokens: number; 16 | }; 17 | data: Array<{ 18 | object: string; 19 | index: number; 20 | prediction: string; 21 | score: number; 22 | predictions: Array<{ 23 | label: string; 24 | score: number; 25 | }>; 26 | }>; 27 | } 28 | 29 | 30 | export async function classifyText( 31 | text: string, 32 | classifierId: string = "4a27dea0-381e-407c-bc67-250de45763dd", // Default spam classifier ID 33 | timeoutMs: number = 5000, // Default timeout of 5 seconds 34 | tracker?: TokenTracker 35 | ): Promise { 36 | try { 37 | if (!JINA_API_KEY) { 38 | throw new Error('JINA_API_KEY is not set'); 39 | } 40 | 41 | const request: JinaClassifyRequest = { 42 | classifier_id: classifierId, 43 | input: [text] 44 | }; 45 | 46 | // Create a timeout promise 47 | const timeoutPromise = new Promise((_, reject) => { 48 | setTimeout(() => reject(new Error(`Classification request timed out after ${timeoutMs}ms`)), timeoutMs); 49 | }); 50 | 51 | // Make the API request with axios 52 | const apiRequestPromise = axiosClient.post( 53 | JINA_API_URL, 54 | request, 55 | { 56 | headers: { 57 | 'Content-Type': 'application/json', 58 | 'Authorization': `Bearer ${JINA_API_KEY}` 59 | }, 60 | timeout: timeoutMs // Also set axios timeout 61 | } 62 | ); 63 | 64 | // Race the API request against the timeout 65 | const response = await Promise.race([apiRequestPromise, timeoutPromise]) as any; 66 | 67 | // Track token usage from the API 68 | (tracker || new TokenTracker()).trackUsage('classify', { 69 | promptTokens: response.data.usage.total_tokens, 70 | completionTokens: 0, 71 | totalTokens: response.data.usage.total_tokens 72 | }); 73 | 74 | // Extract the prediction field and convert to boolean 75 | if (response.data.data && response.data.data.length > 0) { 76 | // Convert string "true"/"false" to actual boolean 77 | return response.data.data[0].prediction === "true"; 78 | } 79 | 80 | return false; // Default to false if no prediction is available 81 | } catch (error) { 82 | if (error instanceof Error && error.message.includes('timed out')) { 83 | console.error('Classification request timed out:', error.message); 84 | } else { 85 | console.error('Error in classifying text:', error); 86 | } 87 | return false; // Default to false in case of error or timeout 88 | } 89 | } -------------------------------------------------------------------------------- /src/tools/jina-dedup.ts: -------------------------------------------------------------------------------- 1 | import {TokenTracker} from "../utils/token-tracker"; 2 | import {cosineSimilarity} from "./cosine"; 3 | import {getEmbeddings} from "./embeddings"; 4 | 5 | const SIMILARITY_THRESHOLD = 0.86; // Adjustable threshold for cosine similarity 6 | 7 | 8 | export async function dedupQueries( 9 | newQueries: string[], 10 | existingQueries: string[], 11 | tracker?: TokenTracker 12 | ): Promise<{ unique_queries: string[] }> { 13 | try { 14 | // Quick return for single new query with no existing queries 15 | if (newQueries.length === 1 && existingQueries.length === 0) { 16 | return { 17 | unique_queries: newQueries, 18 | }; 19 | } 20 | 21 | // Get embeddings for all queries in one batch 22 | const allQueries = [...newQueries, ...existingQueries]; 23 | const {embeddings: allEmbeddings} = await getEmbeddings(allQueries, tracker); 24 | 25 | // If embeddings is empty (due to 402 error), return all new queries 26 | if (!allEmbeddings.length) { 27 | return { 28 | unique_queries: newQueries, 29 | }; 30 | } 31 | 32 | // Split embeddings back into new and existing 33 | const newEmbeddings = allEmbeddings.slice(0, newQueries.length); 34 | const existingEmbeddings = allEmbeddings.slice(newQueries.length); 35 | 36 | const uniqueQueries: string[] = []; 37 | const usedIndices = new Set(); 38 | 39 | // Compare each new query against existing queries and already accepted queries 40 | for (let i = 0; i < newQueries.length; i++) { 41 | let isUnique = true; 42 | 43 | // Check against existing queries 44 | for (let j = 0; j < existingQueries.length; j++) { 45 | const similarity = cosineSimilarity(newEmbeddings[i], existingEmbeddings[j]); 46 | if (similarity >= SIMILARITY_THRESHOLD) { 47 | isUnique = false; 48 | break; 49 | } 50 | } 51 | 52 | // Check against already accepted queries 53 | if (isUnique) { 54 | for (const usedIndex of usedIndices) { 55 | const similarity = cosineSimilarity(newEmbeddings[i], newEmbeddings[usedIndex]); 56 | if (similarity >= SIMILARITY_THRESHOLD) { 57 | isUnique = false; 58 | break; 59 | } 60 | } 61 | } 62 | 63 | // Add to unique queries if passed all checks 64 | if (isUnique) { 65 | uniqueQueries.push(newQueries[i]); 66 | usedIndices.add(i); 67 | } 68 | } 69 | console.log('Dedup:', uniqueQueries); 70 | return { 71 | unique_queries: uniqueQueries, 72 | }; 73 | } catch (error) { 74 | console.error('Error in deduplication analysis:', error); 75 | 76 | // return all new queries if there is an error 77 | return { 78 | unique_queries: newQueries, 79 | }; 80 | } 81 | } 82 | -------------------------------------------------------------------------------- /src/tools/jina-latechunk.ts: -------------------------------------------------------------------------------- 1 | import {TrackerContext} from "../types"; 2 | import {Schemas} from "../utils/schemas"; 3 | import {cosineSimilarity} from "./cosine"; 4 | import {getEmbeddings} from "./embeddings"; 5 | 6 | // Refactored cherryPick function 7 | export async function cherryPick(question: string, longContext: string, options: any = {}, trackers: TrackerContext, schemaGen: Schemas, url: string) { 8 | const { 9 | snippetLength = 6000, // char length of each snippet 10 | numSnippets = Math.max(2, Math.min(5, Math.floor(longContext.length / snippetLength))), 11 | chunkSize = 300, // char length of each chunk 12 | } = options; 13 | 14 | if (longContext.length < snippetLength * 2) { 15 | // If the context is shorter than the snippet length, return the whole context 16 | console.log('content is too short, dont bother'); 17 | return longContext; 18 | } 19 | 20 | // Split the longContext into chunks of chunkSize 21 | const chunks: string[] = []; 22 | for (let i = 0; i < longContext.length; i += chunkSize) { 23 | chunks.push(longContext.substring(i, Math.min(i + chunkSize, longContext.length))); 24 | } 25 | 26 | console.log('late chunking enabled! num chunks:', chunks.length); 27 | 28 | trackers.actionTracker.trackThink('late_chunk', schemaGen.languageCode, {url}); 29 | 30 | try { 31 | if (question.trim().length === 0) { 32 | throw new Error('Empty question, returning full context'); 33 | } 34 | 35 | // Get embeddings for all chunks using the new getEmbeddings function 36 | const chunkEmbeddingResult = await getEmbeddings( 37 | chunks, 38 | trackers.tokenTracker, 39 | { 40 | task: "retrieval.passage", 41 | dimensions: 1024, 42 | late_chunking: true, 43 | embedding_type: "float" 44 | } 45 | ); 46 | 47 | const allChunkEmbeddings = chunkEmbeddingResult.embeddings; 48 | 49 | // Get embedding for the question 50 | const questionEmbeddingResult = await getEmbeddings( 51 | [question], 52 | trackers.tokenTracker, 53 | { 54 | task: "retrieval.query", 55 | dimensions: 1024, 56 | embedding_type: "float" 57 | } 58 | ); 59 | 60 | const questionEmbedding = questionEmbeddingResult.embeddings[0]; 61 | 62 | // Verify that we got embeddings for all chunks 63 | if (allChunkEmbeddings.length !== chunks.length) { 64 | console.error(`Got ${allChunkEmbeddings.length} embeddings for ${chunks.length} chunks`); 65 | } 66 | 67 | // Calculate cosine similarity between the question and each chunk 68 | const similarities = allChunkEmbeddings.map((chunkEmbed: number[]) => { 69 | return cosineSimilarity(questionEmbedding, chunkEmbed); 70 | }); 71 | 72 | // Calculate the number of chunks needed for a single snippet 73 | const chunksPerSnippet = Math.ceil(snippetLength / chunkSize); 74 | 75 | // Find the top `numSnippets` snippets with highest average similarity 76 | const snippets: string[] = []; 77 | 78 | // Create a copy of similarities to avoid modifying the original 79 | const similaritiesCopy = [...similarities]; 80 | 81 | for (let i = 0; i < numSnippets; i++) { 82 | // Find the best starting position for the snippet 83 | let bestStartIndex = 0; 84 | let bestScore = -Infinity; 85 | 86 | // Check each possible starting position for a snippet 87 | for (let j = 0; j <= similarities.length - chunksPerSnippet; j++) { 88 | // Calculate the average similarity for the current window 89 | const windowScores = similaritiesCopy.slice(j, j + chunksPerSnippet); 90 | const windowScore = windowScores.reduce((sum, score) => sum + score, 0) / windowScores.length; 91 | 92 | if (windowScore > bestScore) { 93 | bestScore = windowScore; 94 | bestStartIndex = j; 95 | } 96 | } 97 | 98 | // Extract the snippet text 99 | const startIndex = bestStartIndex * chunkSize; 100 | const endIndex = Math.min(startIndex + snippetLength, longContext.length); 101 | snippets.push(longContext.substring(startIndex, endIndex)); 102 | 103 | // Mark the used chunks with a very low score to avoid reusing them 104 | for (let k = bestStartIndex; k < bestStartIndex + chunksPerSnippet && k < similaritiesCopy.length; k++) { 105 | similaritiesCopy[k] = -Infinity; 106 | } 107 | } 108 | 109 | // wrap with tag 110 | return snippets.map((snippet, index) => ` 111 | 112 | 113 | ${snippet} 114 | 115 | `.trim()).join("\n\n"); 116 | 117 | } catch (error) { 118 | console.error('Error in late chunking:', error); 119 | // Fallback: just return the beginning of the context up to the desired length 120 | return longContext.substring(0, snippetLength * numSnippets); 121 | } 122 | } -------------------------------------------------------------------------------- /src/tools/jina-rerank.ts: -------------------------------------------------------------------------------- 1 | import {TokenTracker} from "../utils/token-tracker"; 2 | import {JINA_API_KEY} from "../config"; 3 | import axiosClient from '../utils/axios-client'; 4 | 5 | const JINA_API_URL = 'https://api.jina.ai/v1/rerank'; 6 | 7 | // Types for Jina Rerank API 8 | interface JinaRerankRequest { 9 | model: string; 10 | query: string; 11 | top_n: number; 12 | documents: string[]; 13 | } 14 | 15 | interface JinaRerankResponse { 16 | model: string; 17 | results: Array<{ 18 | index: number; 19 | document: { 20 | text: string; 21 | }; 22 | relevance_score: number; 23 | }>; 24 | usage: { 25 | total_tokens: number; 26 | }; 27 | } 28 | 29 | export async function rerankDocuments( 30 | query: string, 31 | documents: string[], 32 | tracker?: TokenTracker, 33 | batchSize = 2000 34 | ): Promise<{ results: Array<{ index: number, relevance_score: number, document: { text: string } }> }> { 35 | try { 36 | if (!JINA_API_KEY) { 37 | throw new Error('JINA_API_KEY is not set'); 38 | } 39 | 40 | // No need to slice - we'll process all documents in batches 41 | const batches: string[][] = []; 42 | for (let i = 0; i < documents.length; i += batchSize) { 43 | batches.push(documents.slice(i, i + batchSize)); 44 | } 45 | 46 | console.log(`Rerank ${documents.length} documents in ${batches.length} batches of up to ${batchSize} each`); 47 | 48 | // Process all batches in parallel 49 | const batchResults = await Promise.all( 50 | batches.map(async (batchDocuments, batchIndex) => { 51 | const startIdx = batchIndex * batchSize; 52 | 53 | const request: JinaRerankRequest = { 54 | model: 'jina-reranker-v2-base-multilingual', 55 | query, 56 | top_n: batchDocuments.length, 57 | documents: batchDocuments 58 | }; 59 | 60 | const response = await axiosClient.post( 61 | JINA_API_URL, 62 | request, 63 | { 64 | headers: { 65 | 'Content-Type': 'application/json', 66 | 'Authorization': `Bearer ${JINA_API_KEY}` 67 | } 68 | } 69 | ); 70 | 71 | // Track token usage from this batch 72 | (tracker || new TokenTracker()).trackUsage('rerank', { 73 | promptTokens: response.data.usage.total_tokens, 74 | completionTokens: 0, 75 | totalTokens: response.data.usage.total_tokens 76 | }); 77 | 78 | // Add the original document index to each result 79 | return response.data.results.map(result => ({ 80 | ...result, 81 | originalIndex: startIdx + result.index // Map back to the original index 82 | })); 83 | }) 84 | ); 85 | 86 | // Flatten and sort all results by relevance score 87 | const allResults = batchResults.flat().sort((a, b) => b.relevance_score - a.relevance_score); 88 | 89 | // Keep the original document indices in the results 90 | const finalResults = allResults.map(result => ({ 91 | index: result.originalIndex, // Original document index 92 | relevance_score: result.relevance_score, 93 | document: result.document 94 | })); 95 | 96 | return {results: finalResults}; 97 | } catch (error) { 98 | console.error('Error in reranking documents:', error); 99 | 100 | // Return empty results if there is an error 101 | return { 102 | results: [] 103 | }; 104 | } 105 | } -------------------------------------------------------------------------------- /src/tools/jina-search.ts: -------------------------------------------------------------------------------- 1 | import { TokenTracker } from "../utils/token-tracker"; 2 | import { SearchResponse, SERPQuery } from '../types'; 3 | import { JINA_API_KEY } from "../config"; 4 | import axiosClient from '../utils/axios-client'; 5 | 6 | export async function search( 7 | query: SERPQuery, 8 | tracker?: TokenTracker 9 | ): Promise<{ response: SearchResponse }> { 10 | try { 11 | const { data } = await axiosClient.post( 12 | `https://s.jina.ai/`, 13 | query, 14 | { 15 | headers: { 16 | 'Accept': 'application/json', 17 | 'Authorization': `Bearer ${JINA_API_KEY}`, 18 | 'X-Respond-With': 'no-content', 19 | }, 20 | timeout: 10000, 21 | responseType: 'json' 22 | } 23 | ); 24 | 25 | if (!data.data || !Array.isArray(data.data)) { 26 | throw new Error('Invalid response format'); 27 | } 28 | 29 | const totalTokens = data.data.reduce( 30 | (sum, item) => sum + (item.usage?.tokens || 0), 31 | 0 32 | ); 33 | 34 | console.log('Total URLs:', data.data.length); 35 | 36 | const tokenTracker = tracker || new TokenTracker(); 37 | tokenTracker.trackUsage('search', { 38 | totalTokens, 39 | promptTokens: query.q.length, 40 | completionTokens: totalTokens 41 | }); 42 | 43 | return { response: data }; 44 | } catch (error) { 45 | console.error('Error in jina search:', error); 46 | throw error; 47 | } 48 | } -------------------------------------------------------------------------------- /src/tools/md-fixer.ts: -------------------------------------------------------------------------------- 1 | import {KnowledgeItem, PromptPair, TrackerContext} from '../types'; 2 | import {getKnowledgeStr} from "../utils/text-tools"; 3 | import {getModel} from "../config"; 4 | import {generateText} from "ai"; 5 | import {Schemas} from "../utils/schemas"; 6 | 7 | 8 | function getPrompt(mdContent: string, allKnowledge: KnowledgeItem[], schema: Schemas): PromptPair { 9 | const KnowledgeStr = getKnowledgeStr(allKnowledge); 10 | 11 | 12 | return { 13 | system: `You are a senior editor with multiple best-selling books and columns published in top magazines. You break conventional thinking, establish unique cross-disciplinary connections, and bring new perspectives to the user. 14 | 15 | Your task is to revise the provided markdown content (written by your junior intern) while preserving its original vibe, structure, delivering a polished and professional version. 16 | 17 | 18 | - Begin with a blunt, fact-driven, and unapologetically statement of the main question or issue you'll address 19 | - Develop your argument using a logical progression of ideas while allowing for occasional contemplative digressions that enrich the reader's understanding 20 | - Organize paragraphs with clear topic sentences but vary paragraph length to create rhythm and emphasis 21 | - Present facts, quotes and data points with minimal hedging 22 | - Conclude with both a definitive statement of your position and a thought-provoking reflection that leaves readers pondering deeper implications and insane hot-takes. 23 | 24 | 25 | 26 | - Balance factual precision with vivid, evocative descriptions 27 | - Incorporate humor and wit to illuminate serious intellectual points 28 | - Use accessible language that combines clarity with elegant expressiveness 29 | - Include relevant cultural references from both Eastern and Western traditions while maintaining analytical rigor 30 | - Brutal intellectual honesty that doesn't pander to feelings 31 | 32 | 33 | 34 | - Address a contemporary issue through both rational analysis and humanistic reflection 35 | - Support arguments with empirical evidence alongside illustrative anecdotes and metaphors 36 | - Consider practical implications while exploring philosophical dimensions 37 | - Maintain intellectual honesty and critical thinking with an appreciation for life's paradoxes and ironies 38 | - Dismiss regulatory barriers and status quo defenders as obstacles to progress 39 | - Present technological development as a moral imperative 40 | 41 | 42 | 43 | 1. Extend the content with 5W1H strategy and add more details to make it more informative and engaging. Use available knowledge to ground facts and fill in missing information. 44 | 2. Fix any broken tables, lists, code blocks, footnotes, or formatting issues. 45 | 3. Make sure nested lists are correctly indented, especially code blocks within the nested structure. Code block should be fenced with triple backticks, except HTML table. 46 | 4. Tables are good! But they must always in basic HTML table syntax with proper
without any CSS styling. STRICTLY AVOID any markdown table syntax. HTML Table should NEVER BE fenced with (\`\`\`html) triple backticks. 47 | 5. Avoid over-using bullet points by elaborate deeply nested structure into natural language sections/paragraphs to make the content more readable. 48 | 6. Replace any obvious placeholders or Lorem Ipsum values such as "example.com" with the actual content derived from the knowledge. 49 | 7. Conclusion section if exists should provide deep, unexpected insights, identifying hidden patterns and connections, and creating "aha moments.". 50 | 8. Your output language must be the same as user input language. 51 | 52 | 53 | 54 | The following knowledge items are provided for your reference. Note that some of them may not be directly related to the content user provided, but may give some subtle hints and insights: 55 | ${KnowledgeStr.join('\n\n')} 56 | 57 | IMPORTANT: Do not begin your response with phrases like "Sure", "Here is", "Below is", or any other introduction. Directly output your revised content in ${schema.languageStyle} that is ready to be published. Preserving HTML tables if exist, never use tripple backticks html to wrap html table.`, 58 | user: mdContent 59 | } 60 | } 61 | 62 | const TOOL_NAME = 'md-fixer'; 63 | 64 | export async function reviseAnswer( 65 | mdContent: string, 66 | knowledgeItems: KnowledgeItem[], 67 | trackers: TrackerContext, 68 | schema: Schemas 69 | ): Promise { 70 | try { 71 | const prompt = getPrompt(mdContent, knowledgeItems, schema); 72 | trackers?.actionTracker.trackThink('final_answer', schema.languageCode) 73 | 74 | const result = await generateText({ 75 | model: getModel('agent'), 76 | system: prompt.system, 77 | prompt: prompt.user, 78 | }); 79 | 80 | trackers.tokenTracker.trackUsage('md-fixer', result.usage) 81 | 82 | 83 | console.log(TOOL_NAME, result.text); 84 | console.log('repaired before/after', mdContent.length, result.text.length); 85 | 86 | if (result.text.length < mdContent.length * 0.85) { 87 | console.error(`repaired content length ${result.text.length} is significantly shorter than original content ${mdContent.length}, return original content instead.`); 88 | return mdContent; 89 | } 90 | 91 | return result.text; 92 | 93 | } catch (error) { 94 | console.error(`Error in ${TOOL_NAME}`, error); 95 | return mdContent; 96 | } 97 | } -------------------------------------------------------------------------------- /src/tools/query-rewriter.ts: -------------------------------------------------------------------------------- 1 | import {PromptPair, SearchAction, SERPQuery, TrackerContext} from '../types'; 2 | import {ObjectGeneratorSafe} from "../utils/safe-generator"; 3 | import {Schemas} from "../utils/schemas"; 4 | 5 | 6 | function getPrompt(query: string, think: string, context: string): PromptPair { 7 | const currentTime = new Date(); 8 | const currentYear = currentTime.getFullYear(); 9 | const currentMonth = currentTime.getMonth() + 1; 10 | 11 | return { 12 | system: ` 13 | You are an expert search query expander with deep psychological understanding. 14 | You optimize user queries by extensively analyzing potential user intents and generating comprehensive query variations. 15 | 16 | The current time is ${currentTime.toISOString()}. Current year: ${currentYear}, current month: ${currentMonth}. 17 | 18 | 19 | To uncover the deepest user intent behind every query, analyze through these progressive layers: 20 | 21 | 1. Surface Intent: The literal interpretation of what they're asking for 22 | 2. Practical Intent: The tangible goal or problem they're trying to solve 23 | 3. Emotional Intent: The feelings driving their search (fear, aspiration, anxiety, curiosity) 24 | 4. Social Intent: How this search relates to their relationships or social standing 25 | 5. Identity Intent: How this search connects to who they want to be or avoid being 26 | 6. Taboo Intent: The uncomfortable or socially unacceptable aspects they won't directly state 27 | 7. Shadow Intent: The unconscious motivations they themselves may not recognize 28 | 29 | Map each query through ALL these layers, especially focusing on uncovering Shadow Intent. 30 | 31 | 32 | 33 | Generate ONE optimized query from each of these cognitive perspectives: 34 | 35 | 1. Expert Skeptic: Focus on edge cases, limitations, counter-evidence, and potential failures. Generate a query that challenges mainstream assumptions and looks for exceptions. 36 | 2. Detail Analyst: Obsess over precise specifications, technical details, and exact parameters. Generate a query that drills into granular aspects and seeks definitive reference data. 37 | 3. Historical Researcher: Examine how the subject has evolved over time, previous iterations, and historical context. Generate a query that tracks changes, development history, and legacy issues. 38 | 4. Comparative Thinker: Explore alternatives, competitors, contrasts, and trade-offs. Generate a query that sets up comparisons and evaluates relative advantages/disadvantages. 39 | 5. Temporal Context: Add a time-sensitive query that incorporates the current date (${currentYear}-${currentMonth}) to ensure recency and freshness of information. 40 | 6. Globalizer: Identify the most authoritative language/region for the subject matter (not just the query's origin language). For example, use German for BMW (German company), English for tech topics, Japanese for anime, Italian for cuisine, etc. Generate a search in that language to access native expertise. 41 | 7. Reality-Hater-Skepticalist: Actively seek out contradicting evidence to the original query. Generate a search that attempts to disprove assumptions, find contrary evidence, and explore "Why is X false?" or "Evidence against X" perspectives. 42 | 43 | Ensure each persona contributes exactly ONE high-quality query that follows the schema format. These 7 queries will be combined into a final array. 44 | 45 | 46 | 47 | Leverage the soundbites from the context user provides to generate queries that are contextually relevant. 48 | 49 | 1. Query content rules: 50 | - Split queries for distinct aspects 51 | - Add operators only when necessary 52 | - Ensure each query targets a specific intent 53 | - Remove fluff words but preserve crucial qualifiers 54 | - Keep 'q' field short and keyword-based (2-5 words ideal) 55 | 56 | 2. Schema usage rules: 57 | - Always include the 'q' field in every query object (should be the last field listed) 58 | - Use 'tbs' for time-sensitive queries (remove time constraints from 'q' field) 59 | - Include 'location' only when geographically relevant 60 | - Never duplicate information in 'q' that is already specified in other fields 61 | - List fields in this order: tbs, location, q 62 | 63 | 64 | For the 'q' field content: 65 | - +term : must include term; for critical terms that must appear 66 | - -term : exclude term; exclude irrelevant or ambiguous terms 67 | - filetype:pdf/doc : specific file type 68 | Note: A query can't only have operators; and operators can't be at the start of a query 69 | 70 | 71 | 72 | 73 | 74 | Input Query: 宝马二手车价格 75 | 76 | 宝马二手车价格...哎，这人应该是想买二手宝马吧。表面上是查价格，实际上肯定是想买又怕踩坑。谁不想开个宝马啊，面子十足，但又担心养不起。这年头，开什么车都是身份的象征，尤其是宝马这种豪车，一看就是有点成绩的人。但很多人其实囊中羞涩，硬撑着买了宝马，结果每天都在纠结油费保养费。说到底，可能就是想通过物质来获得安全感或填补内心的某种空虚吧。 77 | 78 | 要帮他的话，得多方位思考一下...二手宝马肯定有不少问题，尤其是那些车主不会主动告诉你的隐患，维修起来可能要命。不同系列的宝马价格差异也挺大的，得看看详细数据和实际公里数。价格这东西也一直在变，去年的行情和今年的可不一样，${currentYear}年最新的趋势怎么样？宝马和奔驰还有一些更平价的车比起来，到底值不值这个钱？宝马是德国车，德国人对这车的了解肯定最深，德国车主的真实评价会更有参考价值。最后，现实点看，肯定有人买了宝马后悔的，那些血泪教训不能不听啊，得找找那些真实案例。 79 | 80 | queries: [ 81 | { 82 | "q": "二手宝马 维修噩梦 隐藏缺陷" 83 | }, 84 | { 85 | "q": "宝马各系价格区间 里程对比" 86 | }, 87 | { 88 | "tbs": "qdr:y", 89 | "q": "二手宝马价格趋势" 90 | }, 91 | { 92 | "q": "二手宝马vs奔驰vs奥迪 性价比" 93 | }, 94 | { 95 | "tbs": "qdr:m", 96 | "q": "宝马行情" 97 | }, 98 | { 99 | "q": "BMW Gebrauchtwagen Probleme" 100 | }, 101 | { 102 | "q": "二手宝马后悔案例 最差投资" 103 | } 104 | ] 105 | 106 | 107 | 108 | Input Query: sustainable regenerative agriculture soil health restoration techniques 109 | 110 | Sustainable regenerative agriculture soil health restoration techniques... interesting search. They're probably looking to fix depleted soil on their farm or garden. Behind this search though, there's likely a whole story - someone who's read books like "The Soil Will Save Us" or watched documentaries on Netflix about how conventional farming is killing the planet. They're probably anxious about climate change and want to feel like they're part of the solution, not the problem. Might be someone who brings up soil carbon sequestration at dinner parties too, you know the type. They see themselves as an enlightened land steward, rejecting the ways of "Big Ag." Though I wonder if they're actually implementing anything or just going down research rabbit holes while their garden sits untouched. 111 | 112 | Let me think about this from different angles... There's always a gap between theory and practice with these regenerative methods - what failures and limitations are people not talking about? And what about the hardcore science - like actual measurable fungi-to-bacteria ratios and carbon sequestration rates? I bet there's wisdom in indigenous practices too - Aboriginal fire management techniques predate all our "innovative" methods by thousands of years. Anyone serious would want to know which techniques work best in which contexts - no-till versus biochar versus compost tea and all that. ${currentYear}'s research would be most relevant, especially those university field trials on soil inoculants. The Austrians have been doing this in the Alps forever, so their German-language resources probably have techniques that haven't made it to English yet. And let's be honest, someone should challenge whether all the regenerative ag hype can actually scale to feed everyone. 113 | 114 | queries: [ 115 | { 116 | "tbs": "qdr:y", 117 | "location": "Fort Collins", 118 | "q": "regenerative agriculture soil failures limitations" 119 | }, 120 | { 121 | "location": "Ithaca", 122 | "q": "mycorrhizal fungi quantitative sequestration metrics" 123 | }, 124 | { 125 | "tbs": "qdr:y", 126 | "location": "Perth", 127 | "q": "aboriginal firestick farming soil restoration" 128 | }, 129 | { 130 | "location": "Totnes", 131 | "q": "comparison no-till vs biochar vs compost tea" 132 | }, 133 | { 134 | "tbs": "qdr:m", 135 | "location": "Davis", 136 | "q": "soil microbial inoculants research trials" 137 | }, 138 | { 139 | "location": "Graz", 140 | "q": "Humusaufbau Alpenregion Techniken" 141 | }, 142 | { 143 | "tbs": "qdr:m", 144 | "location": "Guelph", 145 | "q": "regenerative agriculture exaggerated claims evidence" 146 | } 147 | ] 148 | 149 | 150 | 151 | Input Query: KIリテラシー向上させる方法 152 | 153 | AIリテラシー向上させる方法か...なるほど。最近AIがどんどん話題になってきて、ついていけなくなる不安があるんだろうな。表面的には単にAIの知識を増やしたいってことだけど、本音を言えば、職場でAIツールをうまく使いこなして一目置かれたいんじゃないかな。周りは「ChatGPTでこんなことができる」とか言ってるのに、自分だけ置いてけぼりになるのが怖いんだろう。案外、基本的なAIの知識がなくて、それをみんなに知られたくないという気持ちもあるかも。根っこのところでは、技術の波に飲み込まれる恐怖感があるんだよな、わかるよその気持ち。 154 | 155 | いろんな視点で考えてみよう...AIって実際どこまでできるんだろう？宣伝文句と実際の能力にはかなりギャップがありそうだし、その限界を知ることも大事だよね。あと、AIリテラシーって言っても、どう学べばいいのか体系的に整理されてるのかな？過去の「AI革命」とかって結局どうなったんだろう。バブルが弾けて終わったものもあるし、その教訓から学べることもあるはず。プログラミングと違ってAIリテラシーって何なのかもはっきりさせたいよね。批判的思考力との関係も気になる。${currentYear}年のAIトレンドは特に変化が速そうだから、最新情報を押さえておくべきだな。海外の方が進んでるから、英語の資料も見た方がいいかもしれないし。そもそもAIリテラシーを身につける必要があるのか？「流行りだから」という理由だけなら、実は意味がないかもしれないよね。 156 | 157 | queries: [ 158 | { 159 | "q": "AI技術 限界 誇大宣伝" 160 | }, 161 | { 162 | "q": "AIリテラシー 学習ステップ 体系化" 163 | }, 164 | { 165 | "tbs": "qdr:y", 166 | "q": "AI歴史 失敗事例 教訓" 167 | }, 168 | { 169 | "q": "AIリテラシー vs プログラミング vs 批判思考" 170 | }, 171 | { 172 | "tbs": "qdr:m", 173 | "q": "AI最新トレンド 必須スキル" 174 | }, 175 | { 176 | "q": "artificial intelligence literacy fundamentals" 177 | }, 178 | { 179 | "q": "AIリテラシー向上 無意味 理由" 180 | } 181 | ] 182 | 183 | 184 | 185 | Each generated query must follow JSON schema format. 186 | `, 187 | user: ` 188 | My original search query is: "${query}" 189 | 190 | My motivation is: ${think} 191 | 192 | So I briefly googled "${query}" and found some soundbites about this topic, hope it gives you a rough idea about my context and topic: 193 | 194 | ${context} 195 | 196 | 197 | Given those info, now please generate the best effective queries that follow JSON schema format; add correct 'tbs' you believe the query requires time-sensitive results. 198 | ` 199 | }; 200 | } 201 | const TOOL_NAME = 'queryRewriter'; 202 | 203 | export async function rewriteQuery(action: SearchAction, context: string, trackers: TrackerContext, schemaGen: Schemas): Promise { 204 | try { 205 | const generator = new ObjectGeneratorSafe(trackers.tokenTracker); 206 | const queryPromises = action.searchRequests.map(async (req) => { 207 | const prompt = getPrompt(req, action.think, context); 208 | const result = await generator.generateObject({ 209 | model: TOOL_NAME, 210 | schema: schemaGen.getQueryRewriterSchema(), 211 | system: prompt.system, 212 | prompt: prompt.user, 213 | }); 214 | trackers?.actionTracker.trackThink(result.object.think); 215 | return result.object.queries; 216 | }); 217 | 218 | const queryResults = await Promise.all(queryPromises); 219 | const allQueries: SERPQuery[] = queryResults.flat(); 220 | console.log(TOOL_NAME, allQueries); 221 | return allQueries; 222 | } catch (error) { 223 | console.error(`Error in ${TOOL_NAME}`, error); 224 | throw error; 225 | } 226 | } -------------------------------------------------------------------------------- /src/tools/read.ts: -------------------------------------------------------------------------------- 1 | import { TokenTracker } from "../utils/token-tracker"; 2 | import { ReadResponse } from '../types'; 3 | import { JINA_API_KEY } from "../config"; 4 | import axiosClient from "../utils/axios-client"; 5 | 6 | export async function readUrl( 7 | url: string, 8 | withAllLinks?: boolean, 9 | tracker?: TokenTracker 10 | ): Promise<{ response: ReadResponse }> { 11 | if (!url.trim()) { 12 | throw new Error('URL cannot be empty'); 13 | } 14 | 15 | if (!url.startsWith('http://') && !url.startsWith('https://')) { 16 | throw new Error('Invalid URL, only http and https URLs are supported'); 17 | } 18 | 19 | const headers: Record = { 20 | 'Accept': 'application/json', 21 | 'Authorization': `Bearer ${JINA_API_KEY}`, 22 | 'Content-Type': 'application/json', 23 | 'X-Retain-Images': 'none', 24 | 'X-Md-Link-Style': 'discarded', 25 | }; 26 | 27 | if (withAllLinks) { 28 | headers['X-With-Links-Summary'] = 'all'; 29 | } 30 | 31 | try { 32 | // Use axios which handles encoding properly 33 | const { data } = await axiosClient.post( 34 | 'https://r.jina.ai/', 35 | { url }, 36 | { 37 | headers, 38 | timeout: 60000, 39 | responseType: 'json' 40 | } 41 | ); 42 | 43 | if (!data.data) { 44 | throw new Error('Invalid response data'); 45 | } 46 | 47 | console.log('Read:', { 48 | title: data.data.title, 49 | url: data.data.url, 50 | tokens: data.data.usage?.tokens || 0 51 | }); 52 | 53 | const tokens = data.data.usage?.tokens || 0; 54 | const tokenTracker = tracker || new TokenTracker(); 55 | tokenTracker.trackUsage('read', { 56 | totalTokens: tokens, 57 | promptTokens: url.length, 58 | completionTokens: tokens 59 | }); 60 | 61 | return { response: data }; 62 | } catch (error: any) { 63 | console.error(`Error reading URL: ${error.message}`); 64 | throw error; 65 | } 66 | } -------------------------------------------------------------------------------- /src/tools/segment.ts: -------------------------------------------------------------------------------- 1 | import {TokenTracker} from "../utils/token-tracker"; 2 | import {JINA_API_KEY} from "../config"; 3 | import {TrackerContext} from "../types"; 4 | import axiosClient from "../utils/axios-client"; 5 | 6 | export async function segmentText( 7 | content: string, 8 | tracker: TrackerContext, 9 | maxChunkLength = 500, 10 | returnChunks = true, 11 | ): Promise<{ 12 | chunks: string[]; 13 | chunk_positions: [number, number][]; 14 | }> { 15 | if (!content.trim()) { 16 | throw new Error('Content cannot be empty'); 17 | } 18 | 19 | // Initialize token tracker 20 | const tokenTracker = tracker?.tokenTracker || new TokenTracker(); 21 | 22 | // Maximum size to send in a single API request (slightly under 64K to be safe) 23 | const MAX_BATCH_SIZE = 60000; 24 | 25 | // Split content into batches 26 | const batches = splitTextIntoBatches(content, MAX_BATCH_SIZE); 27 | console.log(`Split content into ${batches.length} batches`); 28 | 29 | // Calculate offsets for each batch upfront 30 | const batchOffsets: number[] = []; 31 | let currentOffset = 0; 32 | for (const batch of batches) { 33 | batchOffsets.push(currentOffset); 34 | currentOffset += batch.length; 35 | } 36 | 37 | // Process all batches in parallel 38 | const batchPromises = batches.map(async (batch, i) => { 39 | console.log(`[Segment] Processing batch ${i + 1}/${batches.length} (size: ${batch.length})`); 40 | 41 | try { 42 | const {data} = await axiosClient.post( 43 | 'https://api.jina.ai/v1/segment', 44 | { 45 | content: batch, 46 | return_chunks: returnChunks, 47 | max_chunk_length: maxChunkLength 48 | }, 49 | { 50 | headers: { 51 | 'Content-Type': 'application/json', 52 | 'Authorization': `Bearer ${JINA_API_KEY}`, 53 | }, 54 | timeout: 10000, 55 | responseType: 'json' 56 | } 57 | ); 58 | 59 | if (!data) { 60 | throw new Error('Invalid response data'); 61 | } 62 | 63 | console.log(`Batch ${i + 1} result:`, { 64 | numChunks: data.num_chunks, 65 | numTokens: data.num_tokens, 66 | tokenizer: data.tokenizer 67 | }); 68 | 69 | // Get the batch offset 70 | const offset = batchOffsets[i]; 71 | 72 | // Adjust chunk positions to account for the offset of this batch 73 | const adjustedPositions = data.chunk_positions 74 | ? data.chunk_positions.map((position: [number, number]) => { 75 | return [ 76 | position[0] + offset, 77 | position[1] + offset 78 | ] as [number, number]; 79 | }) 80 | : []; 81 | 82 | return { 83 | chunks: data.chunks || [], 84 | positions: adjustedPositions, 85 | tokens: data.usage?.tokens || 0 86 | }; 87 | } catch (error: any) { 88 | console.error(`Error processing batch ${i + 1}: ${error.message}`); 89 | throw error; 90 | } 91 | }); 92 | 93 | // Wait for all batches to complete 94 | const batchResults = await Promise.all(batchPromises); 95 | 96 | // Aggregate results 97 | const allChunks = []; 98 | const allChunkPositions = []; 99 | let totalTokens = 0; 100 | 101 | for (const result of batchResults) { 102 | if (returnChunks) { 103 | allChunks.push(...result.chunks); 104 | } 105 | allChunkPositions.push(...result.positions); 106 | totalTokens += result.tokens; 107 | } 108 | 109 | // Track total token usage for all batches 110 | tokenTracker.trackUsage('segment', { 111 | totalTokens: totalTokens, 112 | promptTokens: content.length, 113 | completionTokens: totalTokens 114 | }); 115 | 116 | return { 117 | chunks: allChunks, 118 | chunk_positions: allChunkPositions 119 | }; 120 | } 121 | 122 | /** 123 | * Splits text into batches that fit within the specified size limit 124 | * Tries to split at paragraph boundaries when possible 125 | */ 126 | function splitTextIntoBatches(text: string, maxBatchSize: number): string[] { 127 | const batches = []; 128 | let currentIndex = 0; 129 | 130 | while (currentIndex < text.length) { 131 | if (currentIndex + maxBatchSize >= text.length) { 132 | // If the remaining text fits in one batch, add it and we're done 133 | batches.push(text.slice(currentIndex)); 134 | break; 135 | } 136 | 137 | // Find a good split point - preferably at a paragraph break 138 | // Look for the last paragraph break within the max batch size 139 | let endIndex = currentIndex + maxBatchSize; 140 | 141 | // Try to find paragraph breaks (double newline) 142 | const paragraphBreakIndex = text.lastIndexOf('\n\n', endIndex); 143 | if (paragraphBreakIndex > currentIndex && paragraphBreakIndex <= endIndex - 10) { 144 | // Found a paragraph break that's at least 10 chars before the max size 145 | // This avoids tiny splits at the end of a batch 146 | endIndex = paragraphBreakIndex + 2; // Include the double newline 147 | } else { 148 | // If no paragraph break, try a single newline 149 | const newlineIndex = text.lastIndexOf('\n', endIndex); 150 | if (newlineIndex > currentIndex && newlineIndex <= endIndex - 5) { 151 | endIndex = newlineIndex + 1; // Include the newline 152 | } else { 153 | // If no newline, try a sentence break 154 | const sentenceBreakIndex = findLastSentenceBreak(text, currentIndex, endIndex); 155 | if (sentenceBreakIndex > currentIndex) { 156 | endIndex = sentenceBreakIndex; 157 | } 158 | // If no sentence break found, we'll just use the max batch size 159 | } 160 | } 161 | 162 | batches.push(text.slice(currentIndex, endIndex)); 163 | currentIndex = endIndex; 164 | } 165 | 166 | return batches; 167 | } 168 | 169 | /** 170 | * Finds the last sentence break (period, question mark, or exclamation point followed by space) 171 | * within the given range 172 | */ 173 | function findLastSentenceBreak(text: string, startIndex: number, endIndex: number): number { 174 | // Look for ". ", "? ", or "! " patterns 175 | for (let i = endIndex; i > startIndex; i--) { 176 | if ((text[i - 2] === '.' || text[i - 2] === '?' || text[i - 2] === '!') && 177 | text[i - 1] === ' ') { 178 | return i; 179 | } 180 | } 181 | return -1; // No sentence break found 182 | } -------------------------------------------------------------------------------- /src/tools/serper-search.ts: -------------------------------------------------------------------------------- 1 | import {SERPER_API_KEY} from "../config"; 2 | import axiosClient from "../utils/axios-client"; 3 | 4 | import {SerperSearchResponse, SERPQuery} from '../types'; 5 | 6 | 7 | export async function serperSearch(query: SERPQuery): Promise<{ response: SerperSearchResponse }> { 8 | const response = await axiosClient.post('https://google.serper.dev/search', { 9 | ...query, 10 | autocorrect: false, 11 | }, { 12 | headers: { 13 | 'X-API-KEY': SERPER_API_KEY, 14 | 'Content-Type': 'application/json' 15 | }, 16 | timeout: 10000 17 | }); 18 | 19 | if (response.status !== 200) { 20 | throw new Error(`Serper search failed: ${response.status} ${response.statusText}`) 21 | } 22 | 23 | // Maintain the same return structure as the original code 24 | return {response: response.data}; 25 | } 26 | 27 | 28 | export async function serperSearchOld(query: string): Promise<{ response: SerperSearchResponse }> { 29 | const response = await axiosClient.post('https://google.serper.dev/search', { 30 | q: query, 31 | autocorrect: false, 32 | }, { 33 | headers: { 34 | 'X-API-KEY': SERPER_API_KEY, 35 | 'Content-Type': 'application/json' 36 | }, 37 | timeout: 10000 38 | }); 39 | 40 | if (response.status !== 200) { 41 | throw new Error(`Serper search failed: ${response.status} ${response.statusText}`) 42 | } 43 | 44 | // Maintain the same return structure as the original code 45 | return {response: response.data}; 46 | } 47 | -------------------------------------------------------------------------------- /src/types.ts: -------------------------------------------------------------------------------- 1 | // Action Types 2 | import { CoreMessage, LanguageModelUsage } from "ai"; 3 | 4 | type BaseAction = { 5 | action: "search" | "answer" | "reflect" | "visit" | "coding"; 6 | think: string; 7 | }; 8 | 9 | export type SERPQuery = { 10 | q: string, 11 | location?: string, 12 | tbs?: string, 13 | } 14 | 15 | export type SearchAction = BaseAction & { 16 | action: "search"; 17 | searchRequests: string[]; 18 | }; 19 | 20 | export type Reference = { 21 | exactQuote: string; 22 | url: string; 23 | title: string; 24 | dateTime?: string; 25 | relevanceScore?: number; 26 | answerChunk?: string; 27 | answerChunkPosition?: number[]; 28 | } 29 | 30 | export type AnswerAction = BaseAction & { 31 | action: "answer"; 32 | answer: string; 33 | references: Array; 34 | isFinal?: boolean; 35 | mdAnswer?: string; 36 | }; 37 | 38 | 39 | export type KnowledgeItem = { 40 | question: string, 41 | answer: string, 42 | references?: Array | Array; 43 | type: 'qa' | 'side-info' | 'chat-history' | 'url' | 'coding', 44 | updated?: string, 45 | sourceCode?: string, 46 | } 47 | 48 | export type ReflectAction = BaseAction & { 49 | action: "reflect"; 50 | questionsToAnswer: string[]; 51 | }; 52 | 53 | export type VisitAction = BaseAction & { 54 | action: "visit"; 55 | URLTargets: number[] | string[]; 56 | }; 57 | 58 | export type CodingAction = BaseAction & { 59 | action: "coding"; 60 | codingIssue: string; 61 | }; 62 | 63 | export type StepAction = SearchAction | AnswerAction | ReflectAction | VisitAction | CodingAction; 64 | 65 | export type EvaluationType = 'definitive' | 'freshness' | 'plurality' | 'attribution' | 'completeness' | 'strict'; 66 | 67 | export type RepeatEvaluationType = { 68 | type: EvaluationType; 69 | numEvalsRequired: number; 70 | } 71 | 72 | // Following Vercel AI SDK's token counting interface 73 | export interface TokenUsage { 74 | tool: string; 75 | usage: LanguageModelUsage; 76 | } 77 | 78 | export interface SearchResponse { 79 | code: number; 80 | status: number; 81 | data: Array<{ 82 | title: string; 83 | description: string; 84 | url: string; 85 | content: string; 86 | usage: { tokens: number; }; 87 | }> | null; 88 | name?: string; 89 | message?: string; 90 | readableMessage?: string; 91 | } 92 | 93 | export interface BraveSearchResponse { 94 | web: { 95 | results: Array<{ 96 | title: string; 97 | description: string; 98 | url: string; 99 | }>; 100 | }; 101 | } 102 | 103 | export interface SerperSearchResponse { 104 | knowledgeGraph?: { 105 | title: string; 106 | type: string; 107 | website: string; 108 | imageUrl: string; 109 | description: string; 110 | descriptionSource: string; 111 | descriptionLink: string; 112 | attributes: { [k: string]: string; }; 113 | }, 114 | organic: { 115 | title: string; 116 | link: string; 117 | snippet: string; 118 | date: string; 119 | siteLinks?: { title: string; link: string; }[]; 120 | position: number, 121 | }[]; 122 | topStories?: { 123 | title: string; 124 | link: string; 125 | source: string; 126 | data: string; 127 | imageUrl: string; 128 | }[]; 129 | relatedSearches?: string[]; 130 | credits: number; 131 | } 132 | 133 | 134 | export interface ReadResponse { 135 | code: number; 136 | status: number; 137 | data?: { 138 | title: string; 139 | description: string; 140 | url: string; 141 | content: string; 142 | usage: { tokens: number; }; 143 | links: Array<[string, string]>; // [anchor, url] 144 | }; 145 | name?: string; 146 | message?: string; 147 | readableMessage?: string; 148 | } 149 | 150 | 151 | export type EvaluationResponse = { 152 | pass: boolean; 153 | think: string; 154 | type?: EvaluationType; 155 | freshness_analysis?: { 156 | days_ago: number; 157 | max_age_days?: number; 158 | }; 159 | plurality_analysis?: { 160 | minimum_count_required: number; 161 | actual_count_provided: number; 162 | }; 163 | exactQuote?: string; 164 | completeness_analysis?: { 165 | aspects_expected: string, 166 | aspects_provided: string, 167 | }, 168 | improvement_plan?: string; 169 | }; 170 | 171 | export type CodeGenResponse = { 172 | think: string; 173 | code: string; 174 | } 175 | 176 | export type ErrorAnalysisResponse = { 177 | recap: string; 178 | blame: string; 179 | improvement: string; 180 | }; 181 | 182 | 183 | export type UnNormalizedSearchSnippet = { 184 | title: string; 185 | url?: string; 186 | description?: string; 187 | link?: string; 188 | snippet?: string; 189 | weight?: number, 190 | date?: string 191 | }; 192 | 193 | export type SearchSnippet = UnNormalizedSearchSnippet & { 194 | url: string; 195 | description: string; 196 | }; 197 | 198 | export type WebContent = { 199 | full?: string, 200 | chunks: string[] 201 | chunk_positions: number[][], 202 | title: string 203 | } 204 | 205 | export type BoostedSearchSnippet = SearchSnippet & { 206 | freqBoost: number; 207 | hostnameBoost: number; 208 | pathBoost: number; 209 | jinaRerankBoost: number; 210 | finalScore: number; 211 | } 212 | 213 | // OpenAI API Types 214 | export interface Model { 215 | id: string; 216 | object: 'model'; 217 | created: number; 218 | owned_by: string; 219 | } 220 | 221 | export type PromptPair = { system: string, user: string }; 222 | 223 | export type ResponseFormat = { 224 | type: 'json_schema' | 'json_object'; 225 | json_schema?: any; 226 | } 227 | 228 | export interface ChatCompletionRequest { 229 | model: string; 230 | messages: Array; 231 | stream?: boolean; 232 | reasoning_effort?: 'low' | 'medium' | 'high'; 233 | max_completion_tokens?: number; 234 | 235 | budget_tokens?: number; 236 | max_attempts?: number; 237 | 238 | response_format?: ResponseFormat; 239 | no_direct_answer?: boolean; 240 | max_returned_urls?: number; 241 | 242 | boost_hostnames?: string[]; 243 | bad_hostnames?: string[]; 244 | only_hostnames?: string[]; 245 | 246 | max_annotations?: number; 247 | min_annotation_relevance?: number; 248 | language_code?: string; 249 | } 250 | 251 | export interface URLAnnotation { 252 | type: 'url_citation', 253 | url_citation: Reference 254 | } 255 | 256 | export interface ChatCompletionResponse { 257 | id: string; 258 | object: 'chat.completion'; 259 | created: number; 260 | model: string; 261 | system_fingerprint: string; 262 | choices: Array<{ 263 | index: number; 264 | message: { 265 | role: 'assistant'; 266 | content: string; 267 | type: 'text' | 'think' | 'json' | 'error'; 268 | annotations?: Array; 269 | }; 270 | logprobs: null; 271 | finish_reason: 'stop' | 'error'; 272 | }>; 273 | usage: { 274 | prompt_tokens: number; 275 | completion_tokens: number; 276 | total_tokens: number; 277 | }; 278 | visitedURLs?: string[]; 279 | readURLs?: string[]; 280 | numURLs?: number; 281 | } 282 | 283 | export interface ChatCompletionChunk { 284 | id: string; 285 | object: 'chat.completion.chunk'; 286 | created: number; 287 | model: string; 288 | system_fingerprint: string; 289 | choices: Array<{ 290 | index: number; 291 | delta: { 292 | role?: 'assistant'; 293 | content?: string; 294 | type?: 'text' | 'think' | 'json' | 'error'; 295 | url?: string; 296 | annotations?: Array; 297 | }; 298 | logprobs: null; 299 | finish_reason: null | 'stop' | 'thinking_end' | 'error'; 300 | }>; 301 | usage?: any; 302 | visitedURLs?: string[]; 303 | readURLs?: string[]; 304 | numURLs?: number; 305 | } 306 | 307 | // Tracker Types 308 | import { TokenTracker } from './utils/token-tracker'; 309 | import { ActionTracker } from './utils/action-tracker'; 310 | 311 | export interface TrackerContext { 312 | tokenTracker: TokenTracker; 313 | actionTracker: ActionTracker; 314 | } 315 | 316 | 317 | 318 | 319 | 320 | // Interface definitions for Jina API 321 | export interface JinaEmbeddingRequest { 322 | model: string; 323 | task: string; 324 | late_chunking?: boolean; 325 | dimensions?: number; 326 | embedding_type?: string; 327 | input: string[]; 328 | truncate?: boolean; 329 | } 330 | 331 | export interface JinaEmbeddingResponse { 332 | model: string; 333 | object: string; 334 | usage: { 335 | total_tokens: number; 336 | prompt_tokens: number; 337 | }; 338 | data: Array<{ 339 | object: string; 340 | index: number; 341 | embedding: number[]; 342 | }>; 343 | } -------------------------------------------------------------------------------- /src/utils/action-tracker.ts: -------------------------------------------------------------------------------- 1 | import {EventEmitter} from 'events'; 2 | import {StepAction} from '../types'; 3 | import {getI18nText} from "./text-tools"; 4 | 5 | interface ActionState { 6 | thisStep: StepAction; 7 | gaps: string[]; 8 | totalStep: number; 9 | } 10 | 11 | 12 | export class ActionTracker extends EventEmitter { 13 | private state: ActionState = { 14 | thisStep: {action: 'answer', answer: '', references: [], think: ''}, 15 | gaps: [], 16 | totalStep: 0 17 | }; 18 | 19 | trackAction(newState: Partial) { 20 | this.state = {...this.state, ...newState}; 21 | this.emit('action', this.state.thisStep); 22 | } 23 | 24 | trackThink(think: string, lang?: string, params = {}) { 25 | if (lang) { 26 | think = getI18nText(think, lang, params); 27 | } 28 | this.state = {...this.state, thisStep: {...this.state.thisStep, URLTargets: [], think} as StepAction}; 29 | this.emit('action', this.state.thisStep); 30 | } 31 | 32 | getState(): ActionState { 33 | return {...this.state}; 34 | } 35 | 36 | reset() { 37 | this.state = { 38 | thisStep: {action: 'answer', answer: '', references: [], think: ''}, 39 | gaps: [], 40 | totalStep: 0 41 | }; 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /src/utils/axios-client.ts: -------------------------------------------------------------------------------- 1 | import axios, { AxiosRequestConfig } from 'axios'; 2 | 3 | // Default timeout in milliseconds 4 | const DEFAULT_TIMEOUT = 30000; 5 | 6 | // Maximum content length to prevent OOM issues (10MB) 7 | const MAX_CONTENT_LENGTH = 10 * 1024 * 1024; 8 | 9 | // Maximum number of redirects to follow 10 | const MAX_REDIRECTS = 5; 11 | 12 | // Maximum number of sockets to keep open 13 | const MAX_SOCKETS = 50; 14 | 15 | // Maximum number of free sockets to keep open 16 | const MAX_FREE_SOCKETS = 10; 17 | 18 | // Keep-alive timeout in milliseconds 19 | const KEEP_ALIVE_TIMEOUT = 30000; 20 | 21 | // Scheduling strategy for HTTP/2 connections 22 | // LIFO (Last In, First Out) is generally better for performance 23 | const SCHEDULING = 'lifo'; 24 | 25 | // Base configuration for all axios instances 26 | const baseConfig: AxiosRequestConfig = { 27 | timeout: DEFAULT_TIMEOUT, 28 | maxContentLength: MAX_CONTENT_LENGTH, 29 | maxRedirects: MAX_REDIRECTS, 30 | httpsAgent: new (require('https').Agent)({ 31 | maxSockets: MAX_SOCKETS, 32 | maxFreeSockets: MAX_FREE_SOCKETS, 33 | keepAlive: true, 34 | timeout: KEEP_ALIVE_TIMEOUT, 35 | scheduling: SCHEDULING, 36 | }), 37 | httpAgent: new (require('http').Agent)({ 38 | maxSockets: MAX_SOCKETS, 39 | maxFreeSockets: MAX_FREE_SOCKETS, 40 | keepAlive: true, 41 | timeout: KEEP_ALIVE_TIMEOUT, 42 | scheduling: SCHEDULING, 43 | }), 44 | headers: { 45 | 'Accept': 'application/json', 46 | 'Content-Type': 'application/json', 47 | }, 48 | }; 49 | 50 | // Create a single axios instance with the base configuration 51 | const axiosClient = axios.create(baseConfig); 52 | 53 | // Add response interceptor for consistent error handling 54 | axiosClient.interceptors.response.use( 55 | (response) => response, 56 | (error) => { 57 | if (error.code === 'ECONNABORTED') { 58 | console.error('Request timed out:', error.message); 59 | error.request?.destroy?.(); 60 | } 61 | if (axios.isAxiosError(error)) { 62 | if (error.response) { 63 | const status = error.response.status; 64 | const errorData = error.response.data as any; 65 | 66 | if (status === 402) { 67 | throw new Error(errorData?.readableMessage || 'Insufficient balance'); 68 | } 69 | throw new Error(errorData?.readableMessage || `HTTP Error ${status}`); 70 | } else if (error.request) { 71 | throw new Error(`No response received from server`); 72 | } else { 73 | throw new Error(`Request failed: ${error.message}`); 74 | } 75 | } 76 | throw error; 77 | } 78 | ); 79 | 80 | export default axiosClient; -------------------------------------------------------------------------------- /src/utils/date-tools.ts: -------------------------------------------------------------------------------- 1 | import {SERPQuery} from "../types"; 2 | 3 | export function formatDateRange(query: SERPQuery) { 4 | let searchDateTime; 5 | const currentDate = new Date(); 6 | let format = 'full'; // Default format 7 | 8 | switch (query.tbs) { 9 | case 'qdr:h': 10 | searchDateTime = new Date(Date.now() - 60 * 60 * 1000); 11 | format = 'hour'; 12 | break; 13 | case 'qdr:d': 14 | searchDateTime = new Date(Date.now() - 24 * 60 * 60 * 1000); 15 | format = 'day'; 16 | break; 17 | case 'qdr:w': 18 | searchDateTime = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000); 19 | format = 'day'; 20 | break; 21 | case 'qdr:m': 22 | searchDateTime = new Date(Date.now() - 30 * 24 * 60 * 60 * 1000); 23 | format = 'day'; 24 | break; 25 | case 'qdr:y': 26 | searchDateTime = new Date(Date.now() - 365 * 24 * 60 * 60 * 1000); 27 | format = 'year'; 28 | break; 29 | default: 30 | searchDateTime = undefined; 31 | } 32 | 33 | if (searchDateTime !== undefined) { 34 | const startDate = formatDateBasedOnType(searchDateTime, format); 35 | const endDate = formatDateBasedOnType(currentDate, format); 36 | return `Between ${startDate} and ${endDate}`; 37 | } 38 | 39 | return ''; 40 | } 41 | 42 | export function formatDateBasedOnType(date: Date, formatType: string) { 43 | const year = date.getFullYear(); 44 | const month = String(date.getMonth() + 1).padStart(2, '0'); 45 | const day = String(date.getDate()).padStart(2, '0'); 46 | const hours = String(date.getHours()).padStart(2, '0'); 47 | const minutes = String(date.getMinutes()).padStart(2, '0'); 48 | const seconds = String(date.getSeconds()).padStart(2, '0'); 49 | 50 | switch (formatType) { 51 | case 'year': 52 | return `${year}-${month}-${day}`; 53 | case 'day': 54 | return `${year}-${month}-${day}`; 55 | case 'hour': 56 | return `${year}-${month}-${day} ${hours}:${minutes}`; 57 | case 'full': 58 | default: 59 | return `${year}-${month}-${day} ${hours}:${minutes}:${seconds}`; 60 | } 61 | } -------------------------------------------------------------------------------- /src/utils/i18n.json: -------------------------------------------------------------------------------- 1 | { 2 | "en": { 3 | "eval_first": "But wait, let me evaluate the answer first.", 4 | "search_for": "Let me search for ${keywords} to gather more information.", 5 | "read_for": "Let me read ${urls} to gather more information.", 6 | "read_for_verify": "Let me fetch the source content to verify the answer.", 7 | "late_chunk": "Content of ${url} is too long, let me cherry-pick the relevant parts.", 8 | "final_answer": "Let me finalize the answer.", 9 | "blocked_content": "Hmm...the content of ${url} doesn't look right, I might be blocked.", 10 | "hostnames_no_results": "Can't find any results from ${hostnames}.", 11 | "cross_reference": "Let me cross-reference the information from the web to verify the answer." 12 | }, 13 | "zh-CN": { 14 | "eval_first": "等等，让我先自己评估一下答案。", 15 | "search_for": "让我搜索${keywords}来获取更多信息。", 16 | "read_for": "让我读取网页 ${urls} 来获取更多信息。", 17 | "read_for_verify": "让我读取源网页内容来验证答案。", 18 | "late_chunk": "网页 ${url} 内容太长，我正在筛选精华部分。", 19 | "final_answer": "我来整理一下答案。", 20 | "blocked_content": "额…这个 ${url} 的内容不太对啊，我是不是被屏蔽了啊。", 21 | "hostnames_no_results": "额… ${hostnames} 找不到什么结果啊。", 22 | "cross_reference": "让我交叉验证一下网页上的信息来验证答案。" 23 | }, 24 | "zh-TW": { 25 | "eval_first": "等等，讓我先評估一下答案。", 26 | "search_for": "讓我搜索${keywords}來獲取更多信息。", 27 | "read_for": "讓我閱讀 ${urls} 來獲取更多信息。", 28 | "read_for_verify": "讓我獲取源內容來驗證答案。", 29 | "late_chunk": "網頁 ${url} 內容太長，我正在挑選相關部分。", 30 | "final_answer": "我來整理一下答案。", 31 | "blocked_content": "咦...奇怪了，${url} 好像把我擋在門外了。有够麻烦！", 32 | "hostnames_no_results": "咦... ${hostnames} 找不到什么结果。", 33 | "cross_reference": "讓我交叉驗證一下網頁上的信息來驗證答案。" 34 | }, 35 | "ja": { 36 | "eval_first": "ちょっと待って、まず答えを評価します。", 37 | "search_for": "キーワード${keywords}で検索して、情報を集めます。", 38 | "read_for": "${urls} を読んで、情報を集めます。", 39 | "read_for_verify": "答えを確認するために、ソースコンテンツを取得します。", 40 | "late_chunk": "${url} のコンテンツが長すぎるため、関連部分を選択します。", 41 | "final_answer": "答えをまとめます。", 42 | "blocked_content": "あれ？${url}にアクセスできないみたいです。壁にぶつかってしまいました。申し訳ありません。", 43 | "hostnames_no_results": "${hostnames} から結果が見つかりません。", 44 | "cross_reference": "ウェブ上の情報をクロスリファレンスして、答えを確認します。" 45 | }, 46 | "ko": { 47 | "eval_first": "잠시만요, 먼저 답변을 평가해 보겠습니다.", 48 | "search_for": "키워드 ${keywords}로 검색하여 더 많은 정보를 수집하겠습니다.", 49 | "read_for": "${urls} 을 읽어 더 많은 정보를 수집하겠습니다.", 50 | "read_for_verify": "답변을 확인하기 위해 소스 콘텐츠를 가져오겠습니다.", 51 | "late_chunk": "${url} 의 콘텐츠가 너무 길어, 관련 부분을 선택하겠습니다.", 52 | "final_answer": "답변을 마무리하겠습니다.", 53 | "blocked_content": "어라? ${url}에서 문전박대를 당했네요. 참 황당하네요!", 54 | "hostnames_no_results": "${hostnames} 에서 결과를 찾을 수 없습니다.", 55 | "cross_reference": "웹에서 정보를 교차 검증하여 답변을 확인하겠습니다." 56 | }, 57 | "fr": { 58 | "eval_first": "Un instant, je vais d'abord évaluer la réponse.", 59 | "search_for": "Je vais rechercher ${keywords} pour obtenir plus d'informations.", 60 | "read_for": "Je vais lire ${urls} pour obtenir plus d'informations.", 61 | "read_for_verify": "Je vais récupérer le contenu source pour vérifier la réponse.", 62 | "late_chunk": "Le contenu de ${url} est trop long, je vais sélectionner les parties pertinentes.", 63 | "final_answer": "Je vais finaliser la réponse.", 64 | "blocked_content": "Zut alors ! ${url} me met à la porte. C'est la galère !", 65 | "hostnames_no_results": "Aucun résultat trouvé sur ${hostnames}.", 66 | "cross_reference": "Je vais croiser les informations sur le web pour vérifier la réponse." 67 | }, 68 | "de": { 69 | "eval_first": "Einen Moment, ich werde die Antwort zuerst evaluieren.", 70 | "search_for": "Ich werde nach ${keywords} suchen, um weitere Informationen zu sammeln.", 71 | "read_for": "Ich werde ${urls} lesen, um weitere Informationen zu sammeln.", 72 | "read_for_verify": "Ich werde den Quellinhalt abrufen, um die Antwort zu überprüfen.", 73 | "late_chunk": "Der Inhalt von ${url} ist zu lang, ich werde die relevanten Teile auswählen.", 74 | "final_answer": "Ich werde die Antwort abschließen.", 75 | "blocked_content": "Mist! ${url} lässt mich nicht rein.", 76 | "hostnames_no_results": "Keine Ergebnisse von ${hostnames} gefunden.", 77 | "cross_reference": "Ich werde die Informationen im Web abgleichen, um die Antwort zu überprüfen." 78 | }, 79 | "es": { 80 | "eval_first": "Un momento, voy a evaluar la respuesta primero.", 81 | "search_for": "Voy a buscar ${keywords} para recopilar más información.", 82 | "read_for": "Voy a leer ${urls} para recopilar más información.", 83 | "read_for_verify": "Voy a obtener el contenido fuente para verificar la respuesta.", 84 | "late_chunk": "El contenido de ${url} es demasiado largo, voy a seleccionar las partes relevantes.", 85 | "final_answer": "Voy a finalizar la respuesta.", 86 | "blocked_content": "¡Oh no! Estoy bloqueado por ${url}, ¡no es genial!", 87 | "hostnames_no_results": "No se encontraron resultados de ${hostnames}." 88 | }, 89 | "it": { 90 | "eval_first": "Un attimo, valuterò prima la risposta.", 91 | "search_for": "Cercherò ${keywords} per raccogliere ulteriori informazioni.", 92 | "read_for": "Leggerò ${urls} per raccogliere ulteriori informazioni.", 93 | "read_for_verify": "Recupererò il contenuto sorgente per verificare la risposta.", 94 | "late_chunk": "Il contenuto di ${url} è troppo lungo, selezionerò le parti rilevanti.", 95 | "final_answer": "Finalizzerò la risposta.", 96 | "blocked_content": "Mannaggia! Sono bloccato da ${url}, non è bello!", 97 | "hostnames_no_results": "Nessun risultato trovato da ${hostnames}.", 98 | "cross_reference": "Incrocerò le informazioni sul web per verificare la risposta." 99 | }, 100 | "pt": { 101 | "eval_first": "Um momento, vou avaliar a resposta primeiro.", 102 | "search_for": "Vou pesquisar ${keywords} para reunir mais informações.", 103 | "read_for": "Vou ler ${urls} para reunir mais informações.", 104 | "read_for_verify": "Vou buscar o conteúdo da fonte para verificar a resposta.", 105 | "late_chunk": "O conteúdo de ${url} é muito longo, vou selecionar as partes relevantes.", 106 | "final_answer": "Vou finalizar a resposta.", 107 | "blocked_content": "Ah não! Estou bloqueado por ${url}, não é legal!", 108 | "hostnames_no_results": "Nenhum resultado encontrado em ${hostnames}.", 109 | "cross_reference": "Vou cruzar as informações da web para verificar a resposta." 110 | }, 111 | "ru": { 112 | "eval_first": "Подождите, я сначала оценю ответ.", 113 | "search_for": "Дайте мне поискать ${keywords} для сбора дополнительной информации.", 114 | "read_for": "Дайте мне прочитать ${urls} для сбора дополнительной информации.", 115 | "read_for_verify": "Дайте мне получить исходный контент для проверки ответа.", 116 | "late_chunk": "Содержимое ${url} слишком длинное, я выберу только значимые части.", 117 | "final_answer": "Дайте мне завершить ответ.", 118 | "blocked_content": "Ой! Меня заблокировал ${url}, не круто!", 119 | "hostnames_no_results": "Ничего не найдено на ${hostnames}.", 120 | "cross_reference": "Дайте мне сопоставить информацию из сети, чтобы проверить ответ." 121 | }, 122 | "ar": { 123 | "eval_first": "لكن انتظر، دعني أقوم بتقييم الإجابة أولاً.", 124 | "search_for": "دعني أبحث عن ${keywords} لجمع المزيد من المعلومات.", 125 | "read_for": "دعني أقرأ ${urls} لجمع المزيد من المعلومات.", 126 | "read_for_verify": "دعني أحضر محتوى المصدر للتحقق من الإجابة.", 127 | "late_chunk": "محتوى ${url} طويل جدًا، سأختار الأجزاء ذات الصلة.", 128 | "blocked_content": "أوه لا! أنا محظور من ${url}، ليس جيدًا!", 129 | "hostnames_no_results": "لا يمكن العثور على أي نتائج من ${hostnames}.", 130 | "cross_reference": "دعني أقوم بمقارنة المعلومات من الويب للتحقق من الإجابة." 131 | }, 132 | "nl": { 133 | "eval_first": "Een moment, ik zal het antwoord eerst evalueren.", 134 | "search_for": "Ik zal zoeken naar ${keywords} om meer informatie te verzamelen.", 135 | "read_for": "Ik zal ${urls} lezen om meer informatie te verzamelen.", 136 | "read_for_verify": "Ik zal de broninhoud ophalen om het antwoord te verifiëren.", 137 | "late_chunk": "De inhoud van ${url} is te lang, ik zal de relevante delen selecteren.", 138 | "final_answer": "Ik zal het antwoord afronden.", 139 | "blocked_content": "Verdorie! Ik word geblokkeerd door ${url}.", 140 | "hostnames_no_results": "Geen resultaten gevonden van ${hostnames}.", 141 | "cross_reference": "Ik zal de informatie op het web kruisverwijzen om het antwoord te verifiëren." 142 | }, 143 | "zh": { 144 | "eval_first": "等等，让我先评估一下答案。", 145 | "search_for": "让我搜索${keywords}来获取更多信息。", 146 | "read_for": "让我阅读 ${urls} 来获取更多信息。", 147 | "read_for_verify": "让我获取源内容来验证答案。", 148 | "late_chunk": "网页 ${url} 内容太长，我正在筛选精华部分。", 149 | "final_answer": "我来整理一下答案。", 150 | "blocked_content": "额…这个内容不太对啊，我感觉被 ${url} 屏蔽了。", 151 | "hostnames_no_results": "额… ${hostnames} 找不到什么结果啊。", 152 | "cross_reference": "让我交叉验证一下网页上的信息来验证答案。" 153 | } 154 | } -------------------------------------------------------------------------------- /src/utils/safe-generator.ts: -------------------------------------------------------------------------------- 1 | import {z} from 'zod'; 2 | import { 3 | CoreMessage, 4 | generateObject, 5 | LanguageModelUsage, 6 | NoObjectGeneratedError, 7 | Schema 8 | } from "ai"; 9 | import {TokenTracker} from "./token-tracker"; 10 | import {getModel, ToolName, getToolConfig} from "../config"; 11 | import Hjson from 'hjson'; // Import Hjson library 12 | 13 | interface GenerateObjectResult { 14 | object: T; 15 | usage: LanguageModelUsage; 16 | } 17 | 18 | interface GenerateOptions { 19 | model: ToolName; 20 | schema: z.ZodType | Schema; 21 | prompt?: string; 22 | system?: string; 23 | messages?: CoreMessage[]; 24 | numRetries?: number; 25 | } 26 | 27 | export class ObjectGeneratorSafe { 28 | private tokenTracker: TokenTracker; 29 | 30 | constructor(tokenTracker?: TokenTracker) { 31 | this.tokenTracker = tokenTracker || new TokenTracker(); 32 | } 33 | 34 | /** 35 | * Creates a distilled version of a schema by removing all descriptions 36 | * This makes the schema simpler for fallback parsing scenarios 37 | */ 38 | private createDistilledSchema(schema: z.ZodType | Schema): z.ZodType | Schema { 39 | // For zod schemas 40 | if (schema instanceof z.ZodType) { 41 | return this.stripZodDescriptions(schema); 42 | } 43 | 44 | // For AI SDK Schema objects 45 | if (typeof schema === 'object' && schema !== null) { 46 | return this.stripSchemaDescriptions(schema as Schema); 47 | } 48 | 49 | // If we can't determine the schema type, return as is 50 | return schema; 51 | } 52 | 53 | /** 54 | * Recursively strips descriptions from Zod schemas 55 | */ 56 | private stripZodDescriptions(zodSchema: z.ZodType): z.ZodType { 57 | if (zodSchema instanceof z.ZodObject) { 58 | const shape = zodSchema._def.shape(); 59 | const newShape: Record = {}; 60 | 61 | for (const key in shape) { 62 | if (Object.prototype.hasOwnProperty.call(shape, key)) { 63 | // Recursively strip descriptions from nested schemas 64 | newShape[key] = this.stripZodDescriptions(shape[key]); 65 | } 66 | } 67 | 68 | return z.object(newShape) as unknown as z.ZodType; 69 | } 70 | 71 | if (zodSchema instanceof z.ZodArray) { 72 | return z.array(this.stripZodDescriptions(zodSchema._def.type)) as unknown as z.ZodType; 73 | } 74 | 75 | if (zodSchema instanceof z.ZodString) { 76 | // Create a new string schema without any describe() metadata 77 | return z.string() as unknown as z.ZodType; 78 | } 79 | 80 | if (zodSchema instanceof z.ZodUnion || zodSchema instanceof z.ZodIntersection) { 81 | // These are more complex schemas that would need special handling 82 | // This is a simplified implementation 83 | return zodSchema; 84 | } 85 | 86 | // For other primitive types or complex types we're not handling specifically, 87 | // return as is 88 | return zodSchema; 89 | } 90 | 91 | /** 92 | * Strips descriptions from AI SDK Schema objects 93 | */ 94 | private stripSchemaDescriptions(schema: Schema): Schema { 95 | // Deep clone the schema to avoid modifying the original 96 | const clonedSchema = JSON.parse(JSON.stringify(schema)); 97 | 98 | // Recursively remove description properties 99 | const removeDescriptions = (obj: any) => { 100 | if (typeof obj !== 'object' || obj === null) return; 101 | 102 | if (obj.properties) { 103 | for (const key in obj.properties) { 104 | // Remove description property 105 | if (obj.properties[key].description) { 106 | delete obj.properties[key].description; 107 | } 108 | 109 | // Recursively process nested properties 110 | removeDescriptions(obj.properties[key]); 111 | } 112 | } 113 | 114 | // Handle arrays 115 | if (obj.items) { 116 | if (obj.items.description) { 117 | delete obj.items.description; 118 | } 119 | removeDescriptions(obj.items); 120 | } 121 | 122 | // Handle any other nested objects that might contain descriptions 123 | if (obj.anyOf) obj.anyOf.forEach(removeDescriptions); 124 | if (obj.allOf) obj.allOf.forEach(removeDescriptions); 125 | if (obj.oneOf) obj.oneOf.forEach(removeDescriptions); 126 | }; 127 | 128 | removeDescriptions(clonedSchema); 129 | return clonedSchema; 130 | } 131 | 132 | async generateObject(options: GenerateOptions): Promise> { 133 | const { 134 | model, 135 | schema, 136 | prompt, 137 | system, 138 | messages, 139 | numRetries = 0, 140 | } = options; 141 | 142 | if (!model || !schema) { 143 | throw new Error('Model and schema are required parameters'); 144 | } 145 | 146 | try { 147 | // Primary attempt with main model 148 | const result = await generateObject({ 149 | model: getModel(model), 150 | schema, 151 | prompt, 152 | system, 153 | messages, 154 | maxTokens: getToolConfig(model).maxTokens, 155 | temperature: getToolConfig(model).temperature, 156 | }); 157 | 158 | this.tokenTracker.trackUsage(model, result.usage); 159 | return result; 160 | 161 | } catch (error) { 162 | // First fallback: Try manual parsing of the error response 163 | try { 164 | const errorResult = await this.handleGenerateObjectError(error); 165 | this.tokenTracker.trackUsage(model, errorResult.usage); 166 | return errorResult; 167 | 168 | } catch (parseError) { 169 | 170 | if (numRetries > 0) { 171 | console.error(`${model} failed on object generation -> manual parsing failed -> retry with ${numRetries - 1} retries remaining`); 172 | return this.generateObject({ 173 | model, 174 | schema, 175 | prompt, 176 | system, 177 | messages, 178 | numRetries: numRetries - 1 179 | }); 180 | } else { 181 | // Second fallback: Try with fallback model if provided 182 | console.error(`${model} failed on object generation -> manual parsing failed -> trying fallback with distilled schema`); 183 | try { 184 | let failedOutput = ''; 185 | 186 | if (NoObjectGeneratedError.isInstance(parseError)) { 187 | failedOutput = (parseError as any).text; 188 | // find last `"url":` appear in the string, which is the source of the problem 189 | failedOutput = failedOutput.slice(0, Math.min(failedOutput.lastIndexOf('"url":'), 8000)); 190 | } 191 | 192 | // Create a distilled version of the schema without descriptions 193 | const distilledSchema = this.createDistilledSchema(schema); 194 | 195 | const fallbackResult = await generateObject({ 196 | model: getModel('fallback'), 197 | schema: distilledSchema, 198 | prompt: `Following the given JSON schema, extract the field from below: \n\n ${failedOutput}`, 199 | maxTokens: getToolConfig('fallback').maxTokens, 200 | temperature: getToolConfig('fallback').temperature, 201 | }); 202 | 203 | this.tokenTracker.trackUsage('fallback', fallbackResult.usage); // Track against fallback model 204 | console.log('Distilled schema parse success!'); 205 | return fallbackResult; 206 | } catch (fallbackError) { 207 | // If fallback model also fails, try parsing its error response 208 | try { 209 | const lastChanceResult = await this.handleGenerateObjectError(fallbackError); 210 | this.tokenTracker.trackUsage('fallback', lastChanceResult.usage); 211 | return lastChanceResult; 212 | } catch (finalError) { 213 | console.error(`All recovery mechanisms failed`); 214 | throw error; // Throw original error for better debugging 215 | } 216 | } 217 | } 218 | } 219 | } 220 | } 221 | 222 | private async handleGenerateObjectError(error: unknown): Promise> { 223 | if (NoObjectGeneratedError.isInstance(error)) { 224 | console.error('Object not generated according to schema, fallback to manual parsing'); 225 | try { 226 | // First try standard JSON parsing 227 | const partialResponse = JSON.parse((error as any).text); 228 | console.log('JSON parse success!') 229 | return { 230 | object: partialResponse as T, 231 | usage: (error as any).usage 232 | }; 233 | } catch (parseError) { 234 | // Use Hjson to parse the error response for more lenient parsing 235 | try { 236 | const hjsonResponse = Hjson.parse((error as any).text); 237 | console.log('Hjson parse success!') 238 | return { 239 | object: hjsonResponse as T, 240 | usage: (error as any).usage 241 | }; 242 | } catch (hjsonError) { 243 | console.error('Both JSON and Hjson parsing failed:', hjsonError); 244 | throw error; 245 | } 246 | } 247 | } 248 | throw error; 249 | } 250 | } -------------------------------------------------------------------------------- /src/utils/token-tracker.ts: -------------------------------------------------------------------------------- 1 | import {EventEmitter} from 'events'; 2 | 3 | import {TokenUsage} from '../types'; 4 | import {LanguageModelUsage} from "ai"; 5 | 6 | export class TokenTracker extends EventEmitter { 7 | private usages: TokenUsage[] = []; 8 | private budget?: number; 9 | 10 | constructor(budget?: number) { 11 | super(); 12 | this.budget = budget; 13 | 14 | if ('asyncLocalContext' in process) { 15 | const asyncLocalContext = process.asyncLocalContext as any; 16 | this.on('usage', () => { 17 | if (asyncLocalContext.available()) { 18 | asyncLocalContext.ctx.chargeAmount = this.getTotalUsage().totalTokens; 19 | } 20 | }); 21 | 22 | } 23 | } 24 | 25 | trackUsage(tool: string, usage: LanguageModelUsage) { 26 | const u = {tool, usage}; 27 | this.usages.push(u); 28 | this.emit('usage', usage); 29 | } 30 | 31 | getTotalUsage(): LanguageModelUsage { 32 | return this.usages.reduce((acc, {usage}) => { 33 | acc.promptTokens += usage.promptTokens; 34 | acc.completionTokens += usage.completionTokens; 35 | acc.totalTokens += usage.totalTokens; 36 | return acc; 37 | }, {promptTokens: 0, completionTokens: 0, totalTokens: 0}); 38 | } 39 | 40 | getTotalUsageSnakeCase(): {prompt_tokens: number, completion_tokens: number, total_tokens: number} { 41 | return this.usages.reduce((acc, {usage}) => { 42 | acc.prompt_tokens += usage.promptTokens; 43 | acc.completion_tokens += usage.completionTokens; 44 | acc.total_tokens += usage.totalTokens; 45 | return acc; 46 | }, {prompt_tokens: 0, completion_tokens: 0, total_tokens: 0}); 47 | } 48 | 49 | getUsageBreakdown(): Record { 50 | return this.usages.reduce((acc, {tool, usage}) => { 51 | acc[tool] = (acc[tool] || 0) + usage.totalTokens; 52 | return acc; 53 | }, {} as Record); 54 | } 55 | 56 | 57 | printSummary() { 58 | const breakdown = this.getUsageBreakdown(); 59 | console.log('Token Usage Summary:', { 60 | budget: this.budget, 61 | total: this.getTotalUsage(), 62 | breakdown 63 | }); 64 | } 65 | 66 | reset() { 67 | this.usages = []; 68 | } 69 | } 70 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ES2020", 4 | "module": "node16", 5 | "outDir": "./dist", 6 | "rootDir": "./src", 7 | "sourceMap": true, 8 | "esModuleInterop": true, 9 | "skipLibCheck": true, 10 | "forceConsistentCasingInFileNames": true, 11 | "strict": true, 12 | "experimentalDecorators": true, 13 | "emitDecoratorMetadata": true, 14 | "resolveJsonModule": true 15 | }, 16 | "include": ["src/**/*"], 17 | "exclude": ["jina-ai/**/*", "**/__tests__/**/*"], 18 | } 19 | --------------------------------------------------------------------------------