feat: add web server wrapper for getResponse (#6)

* feat: add web server wrapper for getResponse

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

* chore: update package-lock.json

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

* fix: add proper types to Express route handlers

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

* refactor: update API endpoint path and request parameters

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

* refactor: centralize model configurations with per-tool settings

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

* refactor: update agent.ts to use centralized model configuration

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

---------

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

README.md

@@ -17,6 +17,54 @@ npm run dev "who will be president of US in 2028?"
 npm run dev "what should be jina ai strategy for 2025?"
 ```
 
+## Web Server API
+
+Start the server:
+```bash
+npm run serve
+```
+
+The server will start on http://localhost:3000 with the following endpoints:
+
+### POST /api/v1/query
+Submit a query to be answered:
+```bash
+curl -X POST http://localhost:3000/api/v1/query \
+  -H "Content-Type: application/json" \
+  -d '{
+    "q": "what is the capital of France?",
+    "budget": 1000000,
+    "maxBadAttempt": 3
+  }'
+```
+
+Response:
+```json
+{
+  "requestId": "1234567890"
+}
+```
+
+### GET /api/v1/stream/:requestId
+Connect to the Server-Sent Events stream to receive progress updates and the final answer:
+```bash
+curl -N http://localhost:3000/api/v1/stream/1234567890
+```
+
+The server will emit the following event types:
+- Progress updates: Step number and budget usage
+- Final answer with complete response data
+- Error messages if something goes wrong
+
+Example events:
+```
+data: {"type":"progress","data":"Step 1 / Budget used 10%"}
+data: {"type":"progress","data":"Step 2 / Budget used 25%"}
+data: {"type":"answer","data":{"action":"answer","answer":"Paris is the capital of France","references":[]}}
+```
+
+```
+
 ```mermaid
 flowchart TD
     subgraph Inputs[System Inputs]
@@ -84,4 +132,4 @@ flowchart TD
     class GapQueue,ContextStore,BadStore,QuestionStore,KeywordStore,KnowledgeStore,URLStore state
     class OrigQuestion,TokenBudget input
     class FinalAnswer output
-```
+```

package.json

@@ -8,7 +8,8 @@
     "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"
+    "lint:fix": "eslint . --ext .ts --fix",
+    "serve": "ts-node src/server.ts"
   },
   "keywords": [],
   "author": "",
@@ -16,10 +17,14 @@
   "description": "",
   "dependencies": {
     "@google/generative-ai": "^0.21.0",
+    "@types/cors": "^2.8.17",
+    "@types/express": "^5.0.0",
     "@types/node-fetch": "^2.6.12",
     "axios": "^1.7.9",
+    "cors": "^2.8.5",
     "dotenv": "^16.4.7",
     "duck-duck-scrape": "^2.2.7",
+    "express": "^4.21.2",
     "node-fetch": "^3.3.2",
     "undici": "^7.3.0"
   },

src/agent.ts

@@ -7,7 +7,7 @@ import {rewriteQuery} from "./tools/query-rewriter";
 import {dedupQueries} from "./tools/dedup";
 import {evaluateAnswer} from "./tools/evaluator";
 import {analyzeSteps} from "./tools/error-analyzer";
-import {GEMINI_API_KEY, JINA_API_KEY, MODEL_NAME, SEARCH_PROVIDER, STEP_SLEEP} from "./config";
+import {GEMINI_API_KEY, JINA_API_KEY, SEARCH_PROVIDER, STEP_SLEEP, modelConfigs} from "./config";
 import {tokenTracker} from "./utils/token-tracker";
 import {StepAction, SchemaProperty, ResponseSchema} from "./types";
 
@@ -241,7 +241,7 @@ function removeAllLineBreaks(text: string) {
   return text.replace(/(\r\n|\n|\r)/gm, " ");
 }
 
-async function getResponse(question: string, tokenBudget: number = 1_000_000, maxBadAttempts: number = 3) {
+export async function getResponse(question: string, tokenBudget: number = 1_000_000, maxBadAttempts: number = 3) {
   let step = 0;
   let totalStep = 0;
   let badAttempts = 0;
@@ -291,9 +291,9 @@ async function getResponse(question: string, tokenBudget: number = 1_000_000, ma
       );
 
     const model = genAI.getGenerativeModel({
-      model: MODEL_NAME,
+      model: modelConfigs.agent.model,
       generationConfig: {
-        temperature: 0.7,
+        temperature: modelConfigs.agent.temperature,
         responseMimeType: "application/json",
         responseSchema: getSchema(allowReflect, allowRead, allowAnswer, allowSearch)
       }
@@ -610,9 +610,9 @@ You decided to think out of the box or cut from a completely different angle.`);
       );
 
     const model = genAI.getGenerativeModel({
-      model: MODEL_NAME,
+      model: modelConfigs.agentBeastMode.model,
       generationConfig: {
-        temperature: 0.7,
+        temperature: modelConfigs.agentBeastMode.temperature,
         responseMimeType: "application/json",
         responseSchema: getSchema(false, false, allowAnswer, false)
       }

src/config.ts

@@ -1,6 +1,21 @@
 import dotenv from 'dotenv';
 import { ProxyAgent, setGlobalDispatcher } from 'undici';
 
+interface ModelConfig {
+  model: string;
+  temperature: number;
+}
+
+interface ToolConfigs {
+  dedup: ModelConfig;
+  evaluator: ModelConfig;
+  errorAnalyzer: ModelConfig;
+  queryRewriter: ModelConfig;
+  agent: ModelConfig;
+  agentBeastMode: ModelConfig;
+}
+
+
 dotenv.config();
 
 // Setup the proxy globally if present
@@ -19,7 +34,38 @@ export const JINA_API_KEY = process.env.JINA_API_KEY as string;
 export const BRAVE_API_KEY = process.env.BRAVE_API_KEY as string;
 export const SEARCH_PROVIDER = BRAVE_API_KEY ? 'brave' : 'duck';
 
-export const MODEL_NAME = 'gemini-1.5-flash';
+const DEFAULT_MODEL = 'gemini-1.5-flash';
+
+const defaultConfig: ModelConfig = {
+  model: DEFAULT_MODEL,
+  temperature: 0
+};
+
+export const modelConfigs: ToolConfigs = {
+  dedup: {
+    ...defaultConfig,
+    temperature: 0.1
+  },
+  evaluator: {
+    ...defaultConfig
+  },
+  errorAnalyzer: {
+    ...defaultConfig
+  },
+  queryRewriter: {
+    ...defaultConfig,
+    temperature: 0.1
+  },
+  agent: {
+    ...defaultConfig,
+    temperature: 0.7
+  },
+  agentBeastMode: {
+    ...defaultConfig,
+    temperature: 0.7
+  }
+};
+
 export const STEP_SLEEP = 1000;
 
 if (!GEMINI_API_KEY) throw new Error("GEMINI_API_KEY not found");

src/server.ts

@@ -0,0 +1,84 @@
+import express, { Request, Response, RequestHandler } from 'express';
+import cors from 'cors';
+import { EventEmitter } from 'events';
+import { getResponse } from './agent';
+
+const app = express();
+const port = process.env.PORT || 3000;
+
+// Enable CORS for localhost debugging
+app.use(cors());
+app.use(express.json());
+
+// Create event emitter for SSE
+const eventEmitter = new EventEmitter();
+
+// Type definitions
+interface QueryRequest extends Request {
+  body: {
+    q: string;
+    budget?: number;
+    maxBadAttempt?: number;
+  };
+}
+
+interface StreamResponse extends Response {
+  write: (chunk: string) => boolean;
+}
+
+// SSE endpoint for progress updates
+app.get('/api/v1/stream/:requestId', ((req: Request, res: StreamResponse) => {
+  const requestId = req.params.requestId;
+  
+  res.setHeader('Content-Type', 'text/event-stream');
+  res.setHeader('Cache-Control', 'no-cache');
+  res.setHeader('Connection', 'keep-alive');
+
+  const listener = (data: any) => {
+    res.write(`data: ${JSON.stringify(data)}\n\n`);
+  };
+
+  eventEmitter.on(`progress-${requestId}`, listener);
+
+  req.on('close', () => {
+    eventEmitter.removeListener(`progress-${requestId}`, listener);
+  });
+}) as RequestHandler);
+
+// POST endpoint to handle questions
+app.post('/api/v1/query', (async (req: QueryRequest, res: Response) => {
+  const { q, budget, maxBadAttempt } = req.body;
+  if (!q) {
+    return res.status(400).json({ error: 'Query (q) is required' });
+  }
+
+  const requestId = Date.now().toString();
+  res.json({ requestId });
+
+  // Store original console.log
+  const originalConsoleLog: typeof console.log = console.log;
+
+  try {
+    // Wrap getResponse to emit progress
+    console.log = (...args: any[]) => {
+      originalConsoleLog(...args);
+      const message = args.join(' ');
+      if (message.includes('Step') || message.includes('Budget used')) {
+        eventEmitter.emit(`progress-${requestId}`, { type: 'progress', data: message });
+      }
+    };
+
+    const result = await getResponse(q, budget, maxBadAttempt);
+    eventEmitter.emit(`progress-${requestId}`, { type: 'answer', data: result });
+  } catch (error: any) {
+    eventEmitter.emit(`progress-${requestId}`, { type: 'error', data: error?.message || 'Unknown error' });
+  } finally {
+    console.log = originalConsoleLog;
+  }
+}) as RequestHandler);
+
+app.listen(port, () => {
+  console.log(`Server running at http://localhost:${port}`);
+});
+
+export default app;

src/tools/dedup.ts

@@ -1,5 +1,5 @@
 import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
-import { GEMINI_API_KEY, MODEL_NAME } from "../config";
+import { GEMINI_API_KEY, modelConfigs } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 
 import { DedupResponse } from '../types';
@@ -24,9 +24,9 @@ const responseSchema = {
 
 const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
 const model = genAI.getGenerativeModel({
-  model: MODEL_NAME,
+  model: modelConfigs.dedup.model,
   generationConfig: {
-    temperature: 0.1,
+    temperature: modelConfigs.dedup.temperature,
     responseMimeType: "application/json",
     responseSchema: responseSchema
   }

src/tools/error-analyzer.ts

@@ -1,5 +1,5 @@
 import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
-import { GEMINI_API_KEY, MODEL_NAME } from "../config";
+import { GEMINI_API_KEY, modelConfigs } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 
 import { ErrorAnalysisResponse } from '../types';
@@ -25,9 +25,9 @@ const responseSchema = {
 
 const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
 const model = genAI.getGenerativeModel({
-  model: MODEL_NAME,
+  model: modelConfigs.errorAnalyzer.model,
   generationConfig: {
-    temperature: 0,
+    temperature: modelConfigs.errorAnalyzer.temperature,
     responseMimeType: "application/json",
     responseSchema: responseSchema
   }

src/tools/evaluator.ts

@@ -1,5 +1,5 @@
 import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
-import { GEMINI_API_KEY, MODEL_NAME } from "../config";
+import { GEMINI_API_KEY, modelConfigs } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 
 import { EvaluationResponse } from '../types';
@@ -21,9 +21,9 @@ const responseSchema = {
 
 const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
 const model = genAI.getGenerativeModel({
-  model: MODEL_NAME,
+  model: modelConfigs.evaluator.model,
   generationConfig: {
-    temperature: 0,
+    temperature: modelConfigs.evaluator.temperature,
     responseMimeType: "application/json",
     responseSchema: responseSchema
   }

src/tools/query-rewriter.ts

@@ -1,5 +1,5 @@
 import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
-import { GEMINI_API_KEY, MODEL_NAME } from "../config";
+import { GEMINI_API_KEY, modelConfigs } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 import { SearchAction } from "../types";
 
@@ -28,9 +28,9 @@ const responseSchema = {
 
 const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
 const model = genAI.getGenerativeModel({
-  model: MODEL_NAME,
+  model: modelConfigs.queryRewriter.model,
   generationConfig: {
-    temperature: 0.1,
+    temperature: modelConfigs.queryRewriter.temperature,
     responseMimeType: "application/json",
     responseSchema: responseSchema
   }