From 9fa0fe3b220036aeee50b546a8e3799c7e33e68a Mon Sep 17 00:00:00 2001 From: Han Xiao Date: Mon, 9 Jun 2025 14:31:53 -0700 Subject: [PATCH] feat: add searchProvider param to executeSearchQueries and getResponse --- src/agent.ts | 15 ++++++---- src/app.ts | 79 +++++++++++++++++++++++++++++++--------------------- src/types.ts | 1 + 3 files changed, 58 insertions(+), 37 deletions(-) diff --git a/src/agent.ts b/src/agent.ts index 13134a0..60f5def 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -276,7 +276,8 @@ async function executeSearchQueries( allURLs: Record, SchemaGen: Schemas, webContents: Record, - onlyHostnames?: string[] + onlyHostnames?: string[], + searchProvider?: string ): Promise<{ newKnowledge: KnowledgeItem[], searchedQueries: string[] @@ -295,7 +296,7 @@ async function executeSearchQueries( try { console.log('Search query:', query); - switch (SEARCH_PROVIDER) { + switch (searchProvider || SEARCH_PROVIDER) { case 'jina': results = (await search(query, context.tokenTracker)).response?.data || []; break; @@ -392,7 +393,8 @@ export async function getResponse(question?: string, onlyHostnames: string[] = [], maxRef: number = 10, 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[] }> { let step = 0; @@ -765,7 +767,9 @@ But then you realized you have asked them before. You decided to to think out of context, allURLs, SchemaGen, - allWebContents + allWebContents, + undefined, + searchProvider ); allKeywords.push(...searchedQueries); @@ -794,7 +798,8 @@ But then you realized you have asked them before. You decided to to think out of allURLs, SchemaGen, allWebContents, - onlyHostnames + onlyHostnames, + searchProvider ); if (searchedQueries.length > 0) { diff --git a/src/app.ts b/src/app.ts index dc563d1..7fca88e 100644 --- a/src/app.ts +++ b/src/app.ts @@ -1,6 +1,6 @@ -import express, {Request, Response, RequestHandler} from 'express'; +import express, { Request, Response, RequestHandler } from 'express'; import cors from 'cors'; -import {getResponse} from './agent'; +import { getResponse } from './agent'; import { TrackerContext, ChatCompletionRequest, @@ -9,11 +9,11 @@ import { AnswerAction, Model, StepAction, VisitAction } from './types'; -import {TokenTracker} from "./utils/token-tracker"; -import {ActionTracker} from "./utils/action-tracker"; -import {ObjectGeneratorSafe} from "./utils/safe-generator"; -import {jsonSchema} from "ai"; // or another converter library -import {normalizeHostName} from "./utils/url-tools"; +import { TokenTracker } from "./utils/token-tracker"; +import { ActionTracker } from "./utils/action-tracker"; +import { ObjectGeneratorSafe } from "./utils/safe-generator"; +import { jsonSchema } from "ai"; // or another converter library +import { normalizeHostName } from "./utils/url-tools"; const app = express(); @@ -28,7 +28,7 @@ app.use(express.json({ // Add health check endpoint for Docker container verification app.get('/health', (req, res) => { - res.json({status: 'ok'}); + res.json({ status: 'ok' }); }); async function* streamTextNaturally(text: string, streamingState: StreamingState) { @@ -192,7 +192,7 @@ async function emitRemainingContent( system_fingerprint: 'fp_' + requestId, choices: [{ index: 0, - delta: {content, type: "think"}, + delta: { content, type: "think" }, logprobs: null, finish_reason: null }], @@ -222,12 +222,12 @@ function getTokenBudgetAndMaxAttempts( switch (reasoningEffort) { case 'low': - return {tokenBudget: 100000, maxBadAttempts: 1}; + return { tokenBudget: 100000, maxBadAttempts: 1 }; case 'high': - return {tokenBudget: 1000000, maxBadAttempts: 2}; + return { tokenBudget: 1000000, maxBadAttempts: 2 }; case 'medium': 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', created: 1686935002, owned_by: 'jina-ai' - }]; + }, + { + id: 'jina-deepsearch-v2', + object: 'model', + created: 1717987200, + owned_by: 'jina-ai' + } + ]; res.json({ object: 'list', @@ -281,6 +288,13 @@ app.get('/v1/models/:model', (async (req: Request, res: Response) => { created: 1686935002, 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 { res.status(404).json({ error: { @@ -299,7 +313,7 @@ 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'}); + res.status(401).json({ error: 'Unauthorized' }); return; } @@ -338,7 +352,7 @@ async function processQueue(streamingState: StreamingState, res: Response, reque system_fingerprint: 'fp_' + requestId, choices: [{ index: 0, - delta: {content: word, type: 'think'}, + delta: { content: word, type: 'think' }, logprobs: null, finish_reason: null }] @@ -366,16 +380,16 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => { 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'}); + res.status(401).json({ error: 'Unauthorized' }); return; } } - const clientIp = req.headers['cf-connecting-ip'] || - req.headers['x-forwarded-for'] || - req.ip || - req.socket.remoteAddress || - 'unknown'; + const clientIp = req.headers['cf-connecting-ip'] || + req.headers['x-forwarded-for'] || + req.ip || + req.socket.remoteAddress || + 'unknown'; // Log request details (excluding sensitive data) console.log('[chat/completions] Request:', { model: req.body.model, @@ -388,11 +402,11 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => { 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'}); + 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'}); + return res.status(400).json({ error: 'Last message must be from user' }); } 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 }); - let {tokenBudget, maxBadAttempts} = getTokenBudgetAndMaxAttempts( + let { tokenBudget, maxBadAttempts } = getTokenBudgetAndMaxAttempts( body.reasoning_effort, 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); console.log(responseSchema) } 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, choices: [{ index: 0, - delta: {role: 'assistant', content: '', type: 'think'}, + delta: { role: 'assistant', content: '', type: 'think' }, logprobs: null, finish_reason: null }] @@ -517,7 +531,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => { system_fingerprint: 'fp_' + requestId, choices: [{ index: 0, - delta: {type: 'think', url}, + delta: { type: 'think', url }, logprobs: 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.max_annotations, body.min_annotation_relevance, - body.language_code - ) + body.language_code, + body.search_provider + ) let finalAnswer = (finalStep as AnswerAction).mdAnswer; 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, choices: [{ index: 0, - delta: {content: `\n\n`, type: 'think'}, + delta: { content: `\n\n`, type: 'think' }, logprobs: null, finish_reason: 'thinking_end' }] @@ -711,7 +726,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => { system_fingerprint: 'fp_' + requestId, choices: [{ index: 0, - delta: {content: '', type: 'think'}, + delta: { content: '', type: 'think' }, logprobs: null, finish_reason: 'error' }], @@ -728,7 +743,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => { system_fingerprint: 'fp_' + requestId, choices: [{ index: 0, - delta: {content: errorMessage, type: 'error'}, + delta: { content: errorMessage, type: 'error' }, logprobs: null, finish_reason: 'error' }], diff --git a/src/types.ts b/src/types.ts index cbbe1ea..cde50c6 100644 --- a/src/types.ts +++ b/src/types.ts @@ -260,6 +260,7 @@ export interface ChatCompletionRequest { max_annotations?: number; min_annotation_relevance?: number; language_code?: string; + search_provider?: string; } export interface URLAnnotation {