feat: add searchProvider param to executeSearchQueries and getResponse

This commit is contained in:
Han Xiao 2025-06-09 14:31:53 -07:00
parent d5a31bce52
commit 9fa0fe3b22
3 changed files with 58 additions and 37 deletions

View File

@ -276,7 +276,8 @@ async function executeSearchQueries(
allURLs: Record<string, SearchSnippet>, allURLs: Record<string, SearchSnippet>,
SchemaGen: Schemas, SchemaGen: Schemas,
webContents: Record<string, WebContent>, webContents: Record<string, WebContent>,
onlyHostnames?: string[] onlyHostnames?: string[],
searchProvider?: string
): Promise<{ ): Promise<{
newKnowledge: KnowledgeItem[], newKnowledge: KnowledgeItem[],
searchedQueries: string[] searchedQueries: string[]
@ -295,7 +296,7 @@ async function executeSearchQueries(
try { try {
console.log('Search query:', query); console.log('Search query:', query);
switch (SEARCH_PROVIDER) { switch (searchProvider || SEARCH_PROVIDER) {
case 'jina': case 'jina':
results = (await search(query, context.tokenTracker)).response?.data || []; results = (await search(query, context.tokenTracker)).response?.data || [];
break; break;
@ -392,7 +393,8 @@ export async function getResponse(question?: string,
onlyHostnames: string[] = [], onlyHostnames: string[] = [],
maxRef: number = 10, maxRef: number = 10,
minRelScore: number = 0.75, minRelScore: number = 0.75,
languageCode: string | undefined = undefined languageCode: string | undefined = undefined,
searchProvider?: string
): Promise<{ result: StepAction; context: TrackerContext; visitedURLs: string[], readURLs: string[], allURLs: string[] }> { ): Promise<{ result: StepAction; context: TrackerContext; visitedURLs: string[], readURLs: string[], allURLs: string[] }> {
let step = 0; let step = 0;
@ -765,7 +767,9 @@ But then you realized you have asked them before. You decided to to think out of
context, context,
allURLs, allURLs,
SchemaGen, SchemaGen,
allWebContents allWebContents,
undefined,
searchProvider
); );
allKeywords.push(...searchedQueries); allKeywords.push(...searchedQueries);
@ -794,7 +798,8 @@ But then you realized you have asked them before. You decided to to think out of
allURLs, allURLs,
SchemaGen, SchemaGen,
allWebContents, allWebContents,
onlyHostnames onlyHostnames,
searchProvider
); );
if (searchedQueries.length > 0) { if (searchedQueries.length > 0) {

View File

@ -1,6 +1,6 @@
import express, {Request, Response, RequestHandler} from 'express'; import express, { Request, Response, RequestHandler } from 'express';
import cors from 'cors'; import cors from 'cors';
import {getResponse} from './agent'; import { getResponse } from './agent';
import { import {
TrackerContext, TrackerContext,
ChatCompletionRequest, ChatCompletionRequest,
@ -9,11 +9,11 @@ import {
AnswerAction, AnswerAction,
Model, StepAction, VisitAction Model, StepAction, VisitAction
} from './types'; } from './types';
import {TokenTracker} from "./utils/token-tracker"; import { TokenTracker } from "./utils/token-tracker";
import {ActionTracker} from "./utils/action-tracker"; import { ActionTracker } from "./utils/action-tracker";
import {ObjectGeneratorSafe} from "./utils/safe-generator"; import { ObjectGeneratorSafe } from "./utils/safe-generator";
import {jsonSchema} from "ai"; // or another converter library import { jsonSchema } from "ai"; // or another converter library
import {normalizeHostName} from "./utils/url-tools"; import { normalizeHostName } from "./utils/url-tools";
const app = express(); const app = express();
@ -28,7 +28,7 @@ app.use(express.json({
// Add health check endpoint for Docker container verification // Add health check endpoint for Docker container verification
app.get('/health', (req, res) => { app.get('/health', (req, res) => {
res.json({status: 'ok'}); res.json({ status: 'ok' });
}); });
async function* streamTextNaturally(text: string, streamingState: StreamingState) { async function* streamTextNaturally(text: string, streamingState: StreamingState) {
@ -192,7 +192,7 @@ async function emitRemainingContent(
system_fingerprint: 'fp_' + requestId, system_fingerprint: 'fp_' + requestId,
choices: [{ choices: [{
index: 0, index: 0,
delta: {content, type: "think"}, delta: { content, type: "think" },
logprobs: null, logprobs: null,
finish_reason: null finish_reason: null
}], }],
@ -222,12 +222,12 @@ function getTokenBudgetAndMaxAttempts(
switch (reasoningEffort) { switch (reasoningEffort) {
case 'low': case 'low':
return {tokenBudget: 100000, maxBadAttempts: 1}; return { tokenBudget: 100000, maxBadAttempts: 1 };
case 'high': case 'high':
return {tokenBudget: 1000000, maxBadAttempts: 2}; return { tokenBudget: 1000000, maxBadAttempts: 2 };
case 'medium': case 'medium':
default: default:
return {tokenBudget: 500000, maxBadAttempts: 1}; return { tokenBudget: 500000, maxBadAttempts: 1 };
} }
} }
@ -263,7 +263,14 @@ app.get('/v1/models', (async (_req: Request, res: Response) => {
object: 'model', object: 'model',
created: 1686935002, created: 1686935002,
owned_by: 'jina-ai' owned_by: 'jina-ai'
}]; },
{
id: 'jina-deepsearch-v2',
object: 'model',
created: 1717987200,
owned_by: 'jina-ai'
}
];
res.json({ res.json({
object: 'list', object: 'list',
@ -281,6 +288,13 @@ app.get('/v1/models/:model', (async (req: Request, res: Response) => {
created: 1686935002, created: 1686935002,
owned_by: 'jina-ai' owned_by: 'jina-ai'
}); });
} else if (modelId === 'jina-deepsearch-v2') {
res.json({
id: 'jina-deepsearch-v2',
object: 'model',
created: 1717987200,
owned_by: 'jina-ai'
});
} else { } else {
res.status(404).json({ res.status(404).json({
error: { error: {
@ -299,7 +313,7 @@ if (secret) {
const authHeader = req.headers.authorization; const authHeader = req.headers.authorization;
if (!authHeader || !authHeader.startsWith('Bearer ') || authHeader.split(' ')[1] !== secret) { if (!authHeader || !authHeader.startsWith('Bearer ') || authHeader.split(' ')[1] !== secret) {
console.log('[chat/completions] Unauthorized request'); console.log('[chat/completions] Unauthorized request');
res.status(401).json({error: 'Unauthorized'}); res.status(401).json({ error: 'Unauthorized' });
return; return;
} }
@ -338,7 +352,7 @@ async function processQueue(streamingState: StreamingState, res: Response, reque
system_fingerprint: 'fp_' + requestId, system_fingerprint: 'fp_' + requestId,
choices: [{ choices: [{
index: 0, index: 0,
delta: {content: word, type: 'think'}, delta: { content: word, type: 'think' },
logprobs: null, logprobs: null,
finish_reason: null finish_reason: null
}] }]
@ -366,16 +380,16 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
const authHeader = req.headers.authorization; const authHeader = req.headers.authorization;
if (!authHeader || !authHeader.startsWith('Bearer ') || authHeader.split(' ')[1] !== secret) { if (!authHeader || !authHeader.startsWith('Bearer ') || authHeader.split(' ')[1] !== secret) {
console.log('[chat/completions] Unauthorized request'); console.log('[chat/completions] Unauthorized request');
res.status(401).json({error: 'Unauthorized'}); res.status(401).json({ error: 'Unauthorized' });
return; return;
} }
} }
const clientIp = req.headers['cf-connecting-ip'] || const clientIp = req.headers['cf-connecting-ip'] ||
req.headers['x-forwarded-for'] || req.headers['x-forwarded-for'] ||
req.ip || req.ip ||
req.socket.remoteAddress || req.socket.remoteAddress ||
'unknown'; 'unknown';
// Log request details (excluding sensitive data) // Log request details (excluding sensitive data)
console.log('[chat/completions] Request:', { console.log('[chat/completions] Request:', {
model: req.body.model, model: req.body.model,
@ -388,11 +402,11 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
const body = req.body as ChatCompletionRequest; const body = req.body as ChatCompletionRequest;
if (!body.messages?.length) { if (!body.messages?.length) {
return res.status(400).json({error: 'Messages array is required and must not be empty'}); return res.status(400).json({ error: 'Messages array is required and must not be empty' });
} }
const lastMessage = body.messages[body.messages.length - 1]; const lastMessage = body.messages[body.messages.length - 1];
if (lastMessage.role !== 'user') { if (lastMessage.role !== 'user') {
return res.status(400).json({error: 'Last message must be from user'}); return res.status(400).json({ error: 'Last message must be from user' });
} }
console.log('messages', JSON.stringify(body.messages)); console.log('messages', JSON.stringify(body.messages));
@ -441,7 +455,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
return true; // Keep other messages return true; // Keep other messages
}); });
let {tokenBudget, maxBadAttempts} = getTokenBudgetAndMaxAttempts( let { tokenBudget, maxBadAttempts } = getTokenBudgetAndMaxAttempts(
body.reasoning_effort, body.reasoning_effort,
body.max_completion_tokens body.max_completion_tokens
); );
@ -460,7 +474,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
responseSchema = jsonSchema(body.response_format.json_schema); responseSchema = jsonSchema(body.response_format.json_schema);
console.log(responseSchema) console.log(responseSchema)
} catch (error: any) { } catch (error: any) {
return res.status(400).json({error: `Invalid JSON schema: ${error.message}`}); return res.status(400).json({ error: `Invalid JSON schema: ${error.message}` });
} }
} }
@ -496,7 +510,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
system_fingerprint: 'fp_' + requestId, system_fingerprint: 'fp_' + requestId,
choices: [{ choices: [{
index: 0, index: 0,
delta: {role: 'assistant', content: '<think>', type: 'think'}, delta: { role: 'assistant', content: '<think>', type: 'think' },
logprobs: null, logprobs: null,
finish_reason: null finish_reason: null
}] }]
@ -517,7 +531,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
system_fingerprint: 'fp_' + requestId, system_fingerprint: 'fp_' + requestId,
choices: [{ choices: [{
index: 0, index: 0,
delta: {type: 'think', url}, delta: { type: 'think', url },
logprobs: null, logprobs: null,
finish_reason: null, finish_reason: null,
}] }]
@ -567,8 +581,9 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
body.only_hostnames?.map(i => normalizeHostName(i)), body.only_hostnames?.map(i => normalizeHostName(i)),
body.max_annotations, body.max_annotations,
body.min_annotation_relevance, body.min_annotation_relevance,
body.language_code body.language_code,
) body.search_provider
)
let finalAnswer = (finalStep as AnswerAction).mdAnswer; let finalAnswer = (finalStep as AnswerAction).mdAnswer;
const annotations = (finalStep as AnswerAction).references?.map(ref => ({ const annotations = (finalStep as AnswerAction).references?.map(ref => ({
@ -613,7 +628,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
system_fingerprint: 'fp_' + requestId, system_fingerprint: 'fp_' + requestId,
choices: [{ choices: [{
index: 0, index: 0,
delta: {content: `</think>\n\n`, type: 'think'}, delta: { content: `</think>\n\n`, type: 'think' },
logprobs: null, logprobs: null,
finish_reason: 'thinking_end' finish_reason: 'thinking_end'
}] }]
@ -711,7 +726,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
system_fingerprint: 'fp_' + requestId, system_fingerprint: 'fp_' + requestId,
choices: [{ choices: [{
index: 0, index: 0,
delta: {content: '</think>', type: 'think'}, delta: { content: '</think>', type: 'think' },
logprobs: null, logprobs: null,
finish_reason: 'error' finish_reason: 'error'
}], }],
@ -728,7 +743,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
system_fingerprint: 'fp_' + requestId, system_fingerprint: 'fp_' + requestId,
choices: [{ choices: [{
index: 0, index: 0,
delta: {content: errorMessage, type: 'error'}, delta: { content: errorMessage, type: 'error' },
logprobs: null, logprobs: null,
finish_reason: 'error' finish_reason: 'error'
}], }],

View File

@ -260,6 +260,7 @@ export interface ChatCompletionRequest {
max_annotations?: number; max_annotations?: number;
min_annotation_relevance?: number; min_annotation_relevance?: number;
language_code?: string; language_code?: string;
search_provider?: string;
} }
export interface URLAnnotation { export interface URLAnnotation {