wip

.dockerignore

@@ -0,0 +1 @@
+node_modules

jina-ai/.dockerignore

@@ -0,0 +1 @@
+node_modules

jina-ai/Dockerfile

@@ -0,0 +1,50 @@
+# ---- BUILD STAGE ----
+FROM node:20-slim AS builder
+
+# Set working directory
+WORKDIR /app
+
+# Copy package.json and package-lock.json
+COPY ./package*.json ./
+COPY ./jina-ai/package*.json ./jina-ai/
+
+# Install dependencies
+RUN npm ci
+WORKDIR /app/jina-ai
+RUN npm ci
+
+WORKDIR /app
+
+# Copy application code
+COPY ./src ./src
+COPY ./config.json ./
+COPY ./tsconfig.json ./tsconfig.json
+COPY ./jina-ai/src ./jina-ai/src
+COPY ./jina-ai/tsconfig.json ./jina-ai/tsconfig.json
+
+RUN npm run build
+WORKDIR /app/jina-ai
+RUN npm run build
+
+# ---- PRODUCTION STAGE ----
+FROM node:20 AS production
+
+# Set working directory
+WORKDIR /app
+
+COPY --from=builder /app ./
+# Copy config.json and built files from builder
+
+WORKDIR /app/jina-ai
+
+# Set environment variables (Recommended to set at runtime, avoid hardcoding)
+ENV GEMINI_API_KEY=${GEMINI_API_KEY}
+ENV OPENAI_API_KEY=${OPENAI_API_KEY}
+ENV JINA_API_KEY=${JINA_API_KEY}
+ENV BRAVE_API_KEY=${BRAVE_API_KEY}
+
+# Expose the port the app runs on
+EXPOSE 3000
+
+# Set startup command
+CMD ["node", "./dist/server.js"]

jina-ai/package.json

@@ -0,0 +1,38 @@
+{
+  "name": "@jina-ai/node-deepresearch",
+  "version": "1.0.0",
+  "main": "dist/app.js",
+  "files": [
+    "dist",
+    "README.md",
+    "LICENSE"
+  ],
+  "scripts": {
+    "build": "tsc",
+    "dev": "npx ts-node src/agent.ts",
+    "search": "npx ts-node src/test-duck.ts",
+    "rewrite": "npx ts-node src/tools/query-rewriter.ts",
+    "lint": "eslint . --ext .ts",
+    "lint:fix": "eslint . --ext .ts --fix",
+    "serve": "ts-node src/server.ts",
+    "eval": "ts-node src/evals/batch-evals.ts",
+    "test": "jest --testTimeout=30000",
+    "test:watch": "jest --watch"
+  },
+  "keywords": [],
+  "author": "Jina AI",
+  "license": "Apache-2.0",
+  "description": "",
+  "dependencies": {
+    "@google-cloud/firestore": "^7.11.0",
+    "civkit": "^0.8.3-3e69606",
+    "dayjs": "^1.11.13",
+    "lodash": "^4.17.21",
+    "reflect-metadata": "^0.2.2",
+    "tsyringe": "^4.8.0"
+  },
+  "devDependencies": {
+    "@types/lodash": "^4.17.15",
+    "pino-pretty": "^13.0.0"
+  }
+}

jina-ai/src/lib/async-context.ts

@@ -5,4 +5,5 @@ import { container, singleton } from 'tsyringe';
 export class AsyncLocalContext extends GlobalAsyncContext {}
 
 const instance = container.resolve(AsyncLocalContext);
+Reflect.set(process, 'asyncLocalContext', instance);
 export default instance;

jina-ai/src/patch-express.ts

@@ -14,6 +14,14 @@ const logger = globalLogger.child({ service: 'BillingMiddleware' });
 
 const appName = 'DEEPRESEARCH';
 export const jinaAiBillingMiddleware = (req: Request, res: Response, next: NextFunction) => {
+    if (req.path === '/ping') {
+        res.status(200).end('pone');
+        return;
+    }
+    if (req.method !== 'POST' && req.method !== 'GET') {
+        next();
+        return;
+    }
     asyncLocalContext.run(async () => {
         const googleTraceId = req.get('x-cloud-trace-context')?.split('/')?.[0];
         const ctx = asyncLocalContext.ctx;

jina-ai/src/server.ts

@@ -0,0 +1,56 @@
+import 'reflect-metadata'
+import express from 'express';
+import { jinaAiBillingMiddleware } from "./patch-express";
+import { Server } from 'http';
+
+const app = require('../..').default;
+
+const rootApp = express();
+rootApp.use(jinaAiBillingMiddleware, app);
+
+
+const port = process.env.PORT || 3000;
+
+let server: Server | undefined;
+// Export server startup function for better testing
+export function startServer() {
+  return rootApp.listen(port, () => {
+    console.log(`Server running at http://localhost:${port}`);
+  });
+}
+
+// Start server if running directly
+if (process.env.NODE_ENV !== 'test') {
+  server = startServer();
+}
+
+process.on('unhandledRejection', (_err) => `Is false alarm`);
+
+process.on('uncaughtException', (err) => {
+  console.log('Uncaught exception', err);
+
+  // Looks like Firebase runtime does not handle error properly.
+  // Make sure to quit the process.
+  process.nextTick(() => process.exit(1));
+  console.error('Uncaught exception, process quit.');
+  throw err;
+});
+
+const sigHandler = (signal: string) => {
+  console.log(`Received ${signal}, exiting...`);
+  if (server && server.listening) {
+    console.log(`Shutting down gracefully...`);
+    console.log(`Waiting for the server to drain and close...`);
+    server.close((err) => {
+      if (err) {
+        console.error('Error while closing server', err);
+        return;
+      }
+      process.exit(0);
+    });
+    server.closeIdleConnections();
+  }
+
+}
+process.on('SIGTERM', sigHandler);
+process.on('SIGINT', sigHandler);

jina-ai/tsconfig.json

@@ -0,0 +1,17 @@
+{
+    "compilerOptions": {
+      "target": "ES2020",
+      "module": "node16",
+      "outDir": "./dist",
+      "rootDir": "./src",
+      "sourceMap": true,
+      "esModuleInterop": true,
+      "skipLibCheck": true,
+      "forceConsistentCasingInFileNames": true,
+      "strict": true,
+      "experimentalDecorators": true,
+      "emitDecoratorMetadata": true,
+      "resolveJsonModule": true
+    }
+  }
+  

package.json

@@ -1,8 +1,7 @@
 {
   "name": "node-deepresearch",
   "version": "1.0.0",
-  "main": "dist/index.js",
-  "types": "dist/index.d.ts",
+  "main": "dist/app.js",
   "files": [
     "dist",
     "README.md",
@@ -56,15 +55,5 @@
     "ts-jest": "^29.2.5",
     "ts-node": "^10.9.2",
     "typescript": "^5.7.3"
-  },
-  "optionalDependencies": {
-    "@google-cloud/firestore": "^7.11.0",
-    "@types/lodash": "^4.17.15",
-    "civkit": "^0.8.3-3e69606",
-    "dayjs": "^1.11.13",
-    "lodash": "^4.17.21",
-    "pino-pretty": "^13.0.0",
-    "reflect-metadata": "^0.2.2",
-    "tsyringe": "^4.8.0"
   }
 }

src/jina-ai/Dockerfile

@@ -1,46 +0,0 @@
-# ---- BUILD STAGE ----
-    FROM node:20-slim AS builder
-
-    # Set working directory
-    WORKDIR /app
-    
-    # Copy package.json and package-lock.json
-    COPY package*.json ./
-    
-    # Install dependencies
-    RUN npm install --ignore-scripts
-    
-    # Copy application code
-    COPY . .
-    
-    # Build the application
-    RUN npm run build --ignore-scripts
-    
-    # ---- PRODUCTION STAGE ----
-    FROM node:20-slim AS production
-    
-    # Set working directory
-    WORKDIR /app
-    
-    # Copy package.json and package-lock.json
-    COPY package*.json ./
-    
-    # Install production dependencies only
-    RUN npm install --production  --ignore-scripts
-    
-    # Copy config.json and built files from builder
-    COPY --from=builder /app/config.json ./
-    COPY --from=builder /app/dist ./dist
-    
-    # Set environment variables (Recommended to set at runtime, avoid hardcoding)
-    ENV GEMINI_API_KEY=${GEMINI_API_KEY}
-    ENV OPENAI_API_KEY=${OPENAI_API_KEY}
-    ENV JINA_API_KEY=${JINA_API_KEY}
-    ENV BRAVE_API_KEY=${BRAVE_API_KEY}
-    
-    # Expose the port the app runs on
-    EXPOSE 3000
-    
-    # Set startup command
-    CMD ["node", "./dist/server.js"]
-    

src/jina-ai/server.ts

@@ -1,23 +0,0 @@
-import 'reflect-metadata'
-import app from "../app";
-import express from 'express';
-import { jinaAiBillingMiddleware } from "./patch-express";
-
-
-const rootApp = express();
-rootApp.use(jinaAiBillingMiddleware, app);
-
-
-const port = process.env.PORT || 3000;
-
-// Export server startup function for better testing
-export function startServer() {
-  return rootApp.listen(port, () => {
-    console.log(`Server running at http://localhost:${port}`);
-  });
-}
-
-// Start server if running directly
-if (process.env.NODE_ENV !== 'test') {
-  startServer();
-}

src/utils/token-tracker.ts

@@ -1,7 +1,6 @@
 import { EventEmitter } from 'events';
 
 import { TokenUsage, TokenCategory } from '../types';
-import asyncLocalContext from '../jina-ai/lib/async-context';
 
 export class TokenTracker extends EventEmitter {
   private usages: TokenUsage[] = [];
@@ -11,9 +10,15 @@ export class TokenTracker extends EventEmitter {
     super();
     this.budget = budget;
 
+    if ('asyncLocalContext' in process) {
+      const asyncLocalContext = process.asyncLocalContext as any;
       this.on('usage', () => {
+        if (asyncLocalContext.hasContext()) {
           asyncLocalContext.ctx.chargeAmount = this.getTotalUsage();
+        }
       });
+      
+    }
   }
 
   trackUsage(tool: string, tokens: number, category?: TokenCategory) {