mirror of
https://github.com/jina-ai/node-DeepResearch.git
synced 2026-03-22 15:39:06 +08:00
485 lines
16 KiB
TypeScript
485 lines
16 KiB
TypeScript
import express, {Request, Response, RequestHandler} from 'express';
|
|
import cors from 'cors';
|
|
import {EventEmitter} from 'events';
|
|
import {getResponse} from './agent';
|
|
import {
|
|
StepAction,
|
|
StreamMessage,
|
|
TrackerContext,
|
|
ChatCompletionRequest,
|
|
ChatCompletionResponse,
|
|
ChatCompletionChunk,
|
|
AnswerAction,
|
|
TOKEN_CATEGORIES
|
|
} from './types';
|
|
import fs from 'fs/promises';
|
|
import path from 'path';
|
|
import {TokenTracker} from "./utils/token-tracker";
|
|
import {ActionTracker} from "./utils/action-tracker";
|
|
|
|
const app = express();
|
|
const port = process.env.PORT || 3000;
|
|
|
|
// Get secret from command line args for optional authentication
|
|
const secret = process.argv.find(arg => arg.startsWith('--secret='))?.split('=')[1];
|
|
|
|
app.use(cors());
|
|
app.use(express.json());
|
|
|
|
const eventEmitter = new EventEmitter();
|
|
|
|
interface QueryRequest extends Request {
|
|
body: {
|
|
q: string;
|
|
budget?: number;
|
|
maxBadAttempt?: number;
|
|
};
|
|
}
|
|
|
|
// OpenAI-compatible chat completions endpoint
|
|
app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
|
|
// Check authentication only if secret is set
|
|
if (secret) {
|
|
const authHeader = req.headers.authorization;
|
|
if (!authHeader || !authHeader.startsWith('Bearer ') || authHeader.split(' ')[1] !== secret) {
|
|
console.log('[chat/completions] Unauthorized request');
|
|
res.status(401).json({ error: 'Unauthorized' });
|
|
return;
|
|
}
|
|
}
|
|
|
|
// Log request details (excluding sensitive data)
|
|
console.log('[chat/completions] Request:', {
|
|
model: req.body.model,
|
|
stream: req.body.stream,
|
|
messageCount: req.body.messages?.length,
|
|
hasAuth: !!req.headers.authorization,
|
|
requestId: Date.now().toString()
|
|
});
|
|
|
|
const body = req.body as ChatCompletionRequest;
|
|
if (!body.messages?.length) {
|
|
return res.status(400).json({ error: 'Messages array is required and must not be empty' });
|
|
}
|
|
const lastMessage = body.messages[body.messages.length - 1];
|
|
if (lastMessage.role !== 'user') {
|
|
return res.status(400).json({ error: 'Last message must be from user' });
|
|
}
|
|
|
|
const requestId = Date.now().toString();
|
|
const context: TrackerContext = {
|
|
tokenTracker: new TokenTracker(),
|
|
actionTracker: new ActionTracker()
|
|
};
|
|
|
|
// Track prompt tokens for the initial message
|
|
// Use Vercel's token counting convention - 1 token per message
|
|
const messageTokens = body.messages.length;
|
|
context.tokenTracker.trackUsage('agent', messageTokens, TOKEN_CATEGORIES.PROMPT);
|
|
|
|
if (body.stream) {
|
|
res.setHeader('Content-Type', 'text/event-stream');
|
|
res.setHeader('Cache-Control', 'no-cache');
|
|
res.setHeader('Connection', 'keep-alive');
|
|
|
|
// Send initial chunk with opening think tag
|
|
const initialChunk: ChatCompletionChunk = {
|
|
id: requestId,
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
delta: { role: 'assistant', content: '<think>' },
|
|
logprobs: null,
|
|
finish_reason: null
|
|
}]
|
|
};
|
|
res.write(`data: ${JSON.stringify(initialChunk)}\n\n`);
|
|
|
|
// Set up progress listener with cleanup
|
|
const actionListener = (action: any) => {
|
|
// Track reasoning tokens for each chunk using Vercel's convention
|
|
const chunkTokens = 1; // Default to 1 token per chunk
|
|
context.tokenTracker.trackUsage('evaluator', chunkTokens, TOKEN_CATEGORIES.REASONING);
|
|
|
|
// Only send chunk if there's content to send
|
|
if (action.think) {
|
|
const chunk: ChatCompletionChunk = {
|
|
id: requestId,
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
delta: { content: `<think>${action.think}</think>` },
|
|
logprobs: null,
|
|
finish_reason: null
|
|
}]
|
|
};
|
|
const chunkStr = `data: ${JSON.stringify(chunk)}\n\n`;
|
|
console.log('[chat/completions] Sending chunk:', {
|
|
id: chunk.id,
|
|
content: chunk.choices[0].delta.content,
|
|
finish_reason: chunk.choices[0].finish_reason
|
|
});
|
|
res.write(chunkStr);
|
|
}
|
|
};
|
|
context.actionTracker.on('action', actionListener);
|
|
|
|
// Clean up listener on response finish
|
|
res.on('finish', () => {
|
|
context.actionTracker.removeListener('action', actionListener);
|
|
});
|
|
}
|
|
|
|
try {
|
|
// Track initial query tokens - already tracked above
|
|
// const queryTokens = Buffer.byteLength(lastMessage.content, 'utf-8');
|
|
// context.tokenTracker.trackUsage('agent', queryTokens, 'prompt');
|
|
|
|
let result;
|
|
try {
|
|
({ result } = await getResponse(lastMessage.content, undefined, undefined, context));
|
|
} catch (error: any) {
|
|
// If deduplication fails, retry without it
|
|
if (error?.response?.status === 402) {
|
|
// If deduplication fails, retry with maxBadAttempt=3 to skip dedup
|
|
({ result } = await getResponse(lastMessage.content, undefined, 3, context));
|
|
} else {
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
// Track tokens based on action type
|
|
if (result.action === 'answer') {
|
|
// Track accepted prediction tokens for the final answer using Vercel's convention
|
|
const answerTokens = 1; // Default to 1 token per answer
|
|
context.tokenTracker.trackUsage('evaluator', answerTokens, TOKEN_CATEGORIES.ACCEPTED);
|
|
} else {
|
|
// Track rejected prediction tokens for non-answer responses
|
|
const rejectedTokens = 1; // Default to 1 token per rejected response
|
|
context.tokenTracker.trackUsage('evaluator', rejectedTokens, TOKEN_CATEGORIES.REJECTED);
|
|
}
|
|
|
|
if (body.stream) {
|
|
// Send closing think tag
|
|
const closeThinkChunk: ChatCompletionChunk = {
|
|
id: requestId,
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
delta: { content: '</think>' },
|
|
logprobs: null,
|
|
finish_reason: null
|
|
}]
|
|
};
|
|
res.write(`data: ${JSON.stringify(closeThinkChunk)}\n\n`);
|
|
|
|
// Send final answer as separate chunk
|
|
const answerChunk: ChatCompletionChunk = {
|
|
id: requestId,
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
delta: { content: result.action === 'answer' ? (result as AnswerAction).answer : result.think },
|
|
logprobs: null,
|
|
finish_reason: 'stop'
|
|
}]
|
|
};
|
|
res.write(`data: ${JSON.stringify(answerChunk)}\n\n`);
|
|
res.end();
|
|
} else {
|
|
const usage = context.tokenTracker.getUsageDetails();
|
|
const response: ChatCompletionResponse = {
|
|
id: requestId,
|
|
object: 'chat.completion',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
message: {
|
|
role: 'assistant',
|
|
content: result.action === 'answer' ? (result as AnswerAction).answer : result.think
|
|
},
|
|
logprobs: null,
|
|
finish_reason: 'stop'
|
|
}],
|
|
usage
|
|
};
|
|
|
|
// Log final response (excluding full content for brevity)
|
|
console.log('[chat/completions] Response:', {
|
|
id: response.id,
|
|
status: 200,
|
|
contentLength: response.choices[0].message.content.length,
|
|
usage: response.usage
|
|
});
|
|
|
|
res.json(response);
|
|
}
|
|
} catch (error: any) {
|
|
// Log error details
|
|
console.error('[chat/completions] Error:', {
|
|
message: error?.message || 'An error occurred',
|
|
stack: error?.stack,
|
|
type: error?.constructor?.name,
|
|
requestId
|
|
});
|
|
|
|
// Track error as rejected tokens with Vercel token counting
|
|
const errorMessage = error?.message || 'An error occurred';
|
|
// Default to 1 token for errors as per Vercel AI SDK convention
|
|
const errorTokens = 1;
|
|
context.tokenTracker.trackUsage('evaluator', errorTokens, TOKEN_CATEGORIES.REJECTED);
|
|
|
|
// Clean up event listeners
|
|
context.actionTracker.removeAllListeners('action');
|
|
|
|
// Get token usage in OpenAI API format
|
|
const usage = context.tokenTracker.getUsageDetails();
|
|
|
|
if (body.stream && res.headersSent) {
|
|
// For streaming responses that have already started, send error as a chunk
|
|
// First send closing think tag if we're in the middle of thinking
|
|
const closeThinkChunk: ChatCompletionChunk = {
|
|
id: requestId,
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
delta: { content: '</think>' },
|
|
logprobs: null,
|
|
finish_reason: null
|
|
}]
|
|
};
|
|
res.write(`data: ${JSON.stringify(closeThinkChunk)}\n\n`);
|
|
|
|
// Track error token and send error message
|
|
context.tokenTracker.trackUsage('evaluator', 1, TOKEN_CATEGORIES.REJECTED);
|
|
const errorChunk: ChatCompletionChunk = {
|
|
id: requestId,
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
delta: { content: errorMessage },
|
|
logprobs: null,
|
|
finish_reason: 'stop'
|
|
}]
|
|
};
|
|
res.write(`data: ${JSON.stringify(errorChunk)}\n\n`);
|
|
res.end();
|
|
} else {
|
|
// For non-streaming or not-yet-started responses, send error as JSON
|
|
const response: ChatCompletionResponse = {
|
|
id: requestId,
|
|
object: 'chat.completion',
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: body.model,
|
|
system_fingerprint: 'fp_' + requestId,
|
|
choices: [{
|
|
index: 0,
|
|
message: {
|
|
role: 'assistant',
|
|
content: `Error: ${errorMessage}`
|
|
},
|
|
logprobs: null,
|
|
finish_reason: 'stop'
|
|
}],
|
|
usage
|
|
};
|
|
res.json(response);
|
|
}
|
|
}
|
|
}) as RequestHandler);
|
|
|
|
interface StreamResponse extends Response {
|
|
write: (chunk: string) => boolean;
|
|
}
|
|
|
|
function createProgressEmitter(requestId: string, budget: number | undefined, context: TrackerContext) {
|
|
return () => {
|
|
const state = context.actionTracker.getState();
|
|
const budgetInfo = {
|
|
used: context.tokenTracker.getTotalUsage(),
|
|
total: budget || 1_000_000,
|
|
percentage: ((context.tokenTracker.getTotalUsage() / (budget || 1_000_000)) * 100).toFixed(2)
|
|
};
|
|
|
|
eventEmitter.emit(`progress-${requestId}`, {
|
|
type: 'progress',
|
|
data: {...state.thisStep, totalStep: state.totalStep},
|
|
step: state.totalStep,
|
|
budget: budgetInfo,
|
|
trackers: {
|
|
tokenUsage: context.tokenTracker.getTotalUsage(),
|
|
actionState: context.actionTracker.getState()
|
|
}
|
|
});
|
|
};
|
|
}
|
|
|
|
function cleanup(requestId: string) {
|
|
const context = trackers.get(requestId);
|
|
if (context) {
|
|
context.actionTracker.removeAllListeners();
|
|
context.tokenTracker.removeAllListeners();
|
|
trackers.delete(requestId);
|
|
}
|
|
}
|
|
|
|
function emitTrackerUpdate(requestId: string, context: TrackerContext) {
|
|
const trackerData = {
|
|
tokenUsage: context.tokenTracker.getTotalUsage(),
|
|
tokenBreakdown: context.tokenTracker.getUsageBreakdown(),
|
|
actionState: context.actionTracker.getState().thisStep,
|
|
step: context.actionTracker.getState().totalStep,
|
|
badAttempts: context.actionTracker.getState().badAttempts,
|
|
gaps: context.actionTracker.getState().gaps
|
|
};
|
|
|
|
eventEmitter.emit(`progress-${requestId}`, {
|
|
type: 'progress',
|
|
trackers: trackerData
|
|
});
|
|
}
|
|
|
|
// Store the trackers for each request
|
|
const trackers = new Map<string, TrackerContext>();
|
|
|
|
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();
|
|
|
|
// Create new trackers for this request
|
|
const context: TrackerContext = {
|
|
tokenTracker: new TokenTracker(),
|
|
actionTracker: new ActionTracker()
|
|
};
|
|
trackers.set(requestId, context);
|
|
|
|
// Set up listeners immediately for both trackers
|
|
context.actionTracker.on('action', () => emitTrackerUpdate(requestId, context));
|
|
// context.tokenTracker.on('usage', () => emitTrackerUpdate(requestId, context));
|
|
|
|
res.json({requestId});
|
|
|
|
try {
|
|
const {result} = await getResponse(q, budget, maxBadAttempt, context);
|
|
const emitProgress = createProgressEmitter(requestId, budget, context);
|
|
context.actionTracker.on('action', emitProgress);
|
|
await storeTaskResult(requestId, result);
|
|
eventEmitter.emit(`progress-${requestId}`, {
|
|
type: 'answer',
|
|
data: result,
|
|
trackers: {
|
|
tokenUsage: context.tokenTracker.getTotalUsage(),
|
|
actionState: context.actionTracker.getState()
|
|
}
|
|
});
|
|
cleanup(requestId);
|
|
} catch (error: any) {
|
|
eventEmitter.emit(`progress-${requestId}`, {
|
|
type: 'error',
|
|
data: error?.message || 'Unknown error',
|
|
status: 500,
|
|
trackers: {
|
|
tokenUsage: context.tokenTracker.getTotalUsage(),
|
|
actionState: context.actionTracker.getState()
|
|
}
|
|
});
|
|
cleanup(requestId);
|
|
}
|
|
}) as RequestHandler);
|
|
|
|
app.get('/api/v1/stream/:requestId', (async (req: Request, res: StreamResponse) => {
|
|
const requestId = req.params.requestId;
|
|
const context = trackers.get(requestId);
|
|
|
|
res.setHeader('Content-Type', 'text/event-stream');
|
|
res.setHeader('Cache-Control', 'no-cache');
|
|
res.setHeader('Connection', 'keep-alive');
|
|
|
|
const listener = (data: StreamMessage) => {
|
|
// The trackers are now included in all event types
|
|
// We don't need to add them here as they're already part of the data
|
|
res.write(`data: ${JSON.stringify(data)}\n\n`);
|
|
};
|
|
|
|
eventEmitter.on(`progress-${requestId}`, listener);
|
|
|
|
// Handle client disconnection
|
|
req.on('close', () => {
|
|
eventEmitter.removeListener(`progress-${requestId}`, listener);
|
|
});
|
|
|
|
// Send initial connection confirmation with tracker state
|
|
const initialData = {
|
|
type: 'connected',
|
|
requestId,
|
|
trackers: context ? {
|
|
tokenUsage: context.tokenTracker.getTotalUsage(),
|
|
actionState: context.actionTracker.getState()
|
|
} : null
|
|
};
|
|
res.write(`data: ${JSON.stringify(initialData)}\n\n`);
|
|
}) as RequestHandler);
|
|
|
|
async function storeTaskResult(requestId: string, result: StepAction) {
|
|
try {
|
|
const taskDir = path.join(process.cwd(), 'tasks');
|
|
await fs.mkdir(taskDir, {recursive: true});
|
|
await fs.writeFile(
|
|
path.join(taskDir, `${requestId}.json`),
|
|
JSON.stringify(result, null, 2)
|
|
);
|
|
} catch (error) {
|
|
console.error('Task storage failed:', error);
|
|
throw new Error('Failed to store task result');
|
|
}
|
|
}
|
|
|
|
app.get('/api/v1/task/:requestId', (async (req: Request, res: Response) => {
|
|
const requestId = req.params.requestId;
|
|
try {
|
|
const taskPath = path.join(process.cwd(), 'tasks', `${requestId}.json`);
|
|
const taskData = await fs.readFile(taskPath, 'utf-8');
|
|
res.json(JSON.parse(taskData));
|
|
} catch (error) {
|
|
res.status(404).json({error: 'Task not found'});
|
|
}
|
|
}) as RequestHandler);
|
|
|
|
// Export server startup function for better testing
|
|
export function startServer() {
|
|
return app.listen(port, () => {
|
|
console.log(`Server running at http://localhost:${port}`);
|
|
});
|
|
}
|
|
|
|
// Start server if running directly
|
|
if (process.env.NODE_ENV !== 'test') {
|
|
startServer();
|
|
}
|
|
|
|
export default app;
|