github/node-DeepResearch
1
0
Fork 0
mirror of https://github.com/jina-ai/node-DeepResearch.git synced 2026-03-22 07:29:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: fix mock implementation for token tracking tests

Browse Source 
Co-Authored-By: sha.zhou@jina.ai <sha.zhou@jina.ai>
This commit is contained in:
Devin AI
2025-02-11 12:19:41 +00:00
parent 05b628e61c
commit 1309d14630
1 changed files with 11 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
src/__tests__/server.test.ts
View File

@@ -4,38 +4,17 @@ import type { Express } from 'express';
// Mock OpenAI API responses // Mock OpenAI API responses
jest.mock('@ai-sdk/openai', () => { jest.mock('@ai-sdk/openai', () => {
const generateObjectFn = jest.fn().mockImplementation((options) => {
const responseText = 'This is a test response';
const completionTokens = Math.ceil(Buffer.byteLength(responseText, 'utf-8') / 4);
const promptTokens = Math.ceil(Buffer.byteLength(options.prompt, 'utf-8') / 4);
const response: any = {
usage: {
prompt_tokens: promptTokens,
completion_tokens: completionTokens,
total_tokens: promptTokens + completionTokens,
completion_tokens_details: {
reasoning_tokens: Math.ceil(completionTokens * 0.25),
accepted_prediction_tokens: Math.ceil(completionTokens * 0.5),
rejected_prediction_tokens: Math.ceil(completionTokens * 0.25)
}
}
};
response.object = {
type: 'answer',
think: 'Thinking about the response',
answer: responseText,
references: []
};
return Promise.resolve(response);
});
const mockModel = { const mockModel = {
doGenerate: jest.fn().mockImplementation((messages, options = {}) => { doGenerate: jest.fn().mockImplementation((messages, options = {}) => {
const promptTokens = Math.ceil(Buffer.byteLength(JSON.stringify(messages), 'utf-8') / 4);
const responseText = 'This is a test response'; const responseText = 'This is a test response';
const promptTokens = Math.ceil(Buffer.byteLength(JSON.stringify(messages), 'utf-8') / 4);
const completionTokens = Math.ceil(Buffer.byteLength(responseText, 'utf-8') / 4); const completionTokens = Math.ceil(Buffer.byteLength(responseText, 'utf-8') / 4);
const response: any = { return Promise.resolve({
choices: [{
message: {
content: responseText
}
}],
usage: { usage: {
prompt_tokens: promptTokens, prompt_tokens: promptTokens,
completion_tokens: completionTokens, completion_tokens: completionTokens,
@@ -46,34 +25,15 @@ jest.mock('@ai-sdk/openai', () => {
rejected_prediction_tokens: Math.ceil(completionTokens * 0.25) rejected_prediction_tokens: Math.ceil(completionTokens * 0.25)
} }
} }
}; });
})
response.object = {
type: 'answer',
think: 'Thinking about the response',
answer: responseText,
references: []
};
return Promise.resolve(response);
}),
generateObject: generateObjectFn
}; };
return { return {
createOpenAI: jest.fn().mockImplementation(() => { createOpenAI: jest.fn().mockImplementation(() => {
const model = () => ({ const model = () => mockModel;
...mockModel,
generateObject: generateObjectFn,
defaultObjectGenerationMode: 'object'
});
model.defaultObjectGenerationMode = 'object';
return model; return model;
}), }),
OpenAIChatLanguageModel: jest.fn().mockImplementation(() => ({ OpenAIChatLanguageModel: jest.fn().mockImplementation(() => mockModel)
...mockModel,
generateObject: generateObjectFn,
defaultObjectGenerationMode: 'object'
}))
}; };
}); });
const TEST_SECRET = 'test-secret'; const TEST_SECRET = 'test-secret';