zoeissleeping/veridian
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add web search tool with reranking support

Browse Source 
Integrates SearXNG web search as a tool available during chat when the
agent has search enabled. Supports optional reranking of results via a
configurable reranking model. Replaces Python subprocess evaluation with
@pydantic/monty WASM runtime. Introduces stable tool call ID mapping to
avoid exposing provider-native IDs to the database.
This commit is contained in:
Zoe
2026-04-27 12:05:47 -05:00
parent 90d5698e76
commit 9914c043cb
6 changed files with 490 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/api/topic/[topicId]/chat/index.post.ts
+258 -96
View File
@@ -3,15 +3,16 @@ import * as z from 'zod';
import { type MessageEntity } from '~/composables/useChat';
import { promises as fs } from 'fs';
import { glob } from 'glob';
import { type ModelMessage, streamText, type StreamTextTransform, tool } from "ai";
import { isLoopFinished, type ModelMessage, rerank, type RerankingModel, streamText, type StreamTextTransform, Tool, tool } from "ai";
import { generations, messageParts, messages, toolCalls, ToolCallType } from "~~/drizzle/schema";
import { topicEvents } from "~~/server/utils/events";
import { nanoid } from "nanoid";
import { spawn } from "child_process";
import { Monty, MontyRuntimeError, MontySyntaxError, MontyTypingError } from '@pydantic/monty';
import { type Model } from "~/composables/useModels";
import { eq } from "drizzle-orm";
import { buildFocusedMessageTree, buildMessageTree, marshallMessages } from "~~/utils/message";
import path from "path";
import { isRerankingProvider } from "~~/server/utils/ai-provider";
export default defineEventHandler(async (event) => {
await protectRoute(event);
@@ -22,6 +23,10 @@ export default defineEventHandler(async (event) => {
const result = await readValidatedBody(event, z.object({
parentMessageId: z.string().optional(),
modelId: z.string(),
rerank: z.object({
modelId: z.string(),
providerApiKey: z.string().optional(),
}).optional(),
args: z.record(z.string(), z.any()).optional(),
providerApiKey: z.string().optional(),
}).safeParse);
@@ -33,7 +38,7 @@ export default defineEventHandler(async (event) => {
});
}
const { parentMessageId, modelId, providerApiKey, args } = result.data;
const { parentMessageId, modelId, rerank: rerankConfig, providerApiKey, args } = result.data;
const topic = await db.query.topics.findFirst({
where: {
@@ -109,6 +114,91 @@ export default defineEventHandler(async (event) => {
}
}
const agentSearchConfig = topic.agent.config?.search;
let searchParam: false | { config: SearchTheWebConfig } = false;
if (agentSearchConfig?.enabled) {
searchParam = { config: { rerank: false, maxResults: agentSearchConfig.maxResults ?? 10 } };
if (agentSearchConfig.rerank && rerankConfig) {
const rerankModel = await db.query.models.findFirst({
where: {
id: rerankConfig.modelId,
userId,
},
with: {
provider: true,
}
});
if (rerankModel === undefined) {
throw createError({
statusCode: 400,
statusMessage: 'Invalid rerank model',
data: {
code: 'INVALID_RERANK_MODEL',
ok: false,
}
});
}
const rerankProviderDetails = await getProviderDetails(rerankModel.provider, rerankConfig.providerApiKey, rerankModel);
if (!rerankProviderDetails.ok) {
switch (rerankProviderDetails.error) {
case GatewayFetchError.NoProviderApiKey: {
throw createError({
statusCode: 400,
statusMessage: `${rerankModel.provider.type} provider requires an API key`,
data: {
code: 'NO_RERANK_PROVIDER_API_KEY',
ok: false,
}
});
}
case GatewayFetchError.NoProviderBaseUrl: {
throw createError({
statusCode: 400,
statusMessage: 'Invalid provider URL',
data: {
code: 'BAD_RERANK_PROVIDER_URL',
ok: false,
}
});
}
}
}
const { gateway: rerankGateway } = rerankProviderDetails.data;
if (rerankGateway === null) {
throw createError({
statusCode: 500,
statusMessage: 'Invalid gateway',
data: {
code: 'INVALID_RERANK_GATEWAY',
ok: false,
}
});
}
if (isRerankingProvider(rerankGateway.gateway)) {
searchParam.config = {
rerank: true,
maxResults: agentSearchConfig.maxResults ?? 10,
model: rerankGateway.gateway.reranking(rerankModel.externalId)
};
} else {
throw createError({
statusCode: 500,
statusMessage: 'Invalid rerank model',
data: {
code: 'INVALID_RERANK_MODEL',
ok: false,
}
});
}
}
}
let agentmessage = await db.transaction(async tx => {
const generationId = nanoid();
@@ -234,6 +324,7 @@ export default defineEventHandler(async (event) => {
generateResponse(
agentmessage as MessageEntity,
{ gateway: gateway.gateway, model, parameters: args },
searchParam,
agentmessage.generationId!,
userId,
topicId,
@@ -255,54 +346,83 @@ const INTERNAL_ERROR = 'An internal error occurred';
// todo message takes in variadics like console.log
const todo = (...args: any[]) => {
console.error('TODO', ...args);
throw new Error('TODO');
};
const formatPartId = (partType: string, existingId: string) => {
return `veridian__part-${partType}-${existingId}-${nanoid()}`;
};
const formatToolCallId = (nativeId: string) => {
return `veridian__tool-${nativeId}-${nanoid()}`;
};
const evalPython = async (code: string) => {
// This is the wrapper logic from above, minified or stored as a string
// Or you can save the wrapper script to a file and call that.
const wrapper = `
import ast
import sys
code = sys.stdin.read()
tree = ast.parse(code)
last_node = tree.body[-1] if tree.body else None
namespace = {}
if len(tree.body) > 1:
exec(compile(ast.Module(tree.body[:-1], []), "<ast>", "exec"), namespace)
if isinstance(last_node, ast.Expr):
res = eval(compile(ast.Expression(last_node.value), "<ast>", "eval"), namespace)
if res is not None: print(res)
elif last_node:
exec(compile(ast.Module([last_node], []), "<ast>", "exec"), namespace)
`.trim();
try {
let stdout = ''
const printCallback = (_: string, text: string) => {
stdout += text
}
const m = new Monty(code)
await m.run({ printCallback })
return stdout
} catch (error) {
if (error instanceof MontySyntaxError) {
console.log('Syntax error:', error.message)
} else if (error instanceof MontyRuntimeError) {
console.log('Runtime error:', error.message)
console.log('Traceback:', error.traceback())
} else if (error instanceof MontyTypingError) {
console.log('Type error:', error.displayDiagnostics())
}
}
};
return new Promise<string>((resolve, reject) => {
const child = spawn('python3', ['-c', wrapper]);
interface SearchTheWebRerankedConfig {
rerank: true;
maxResults: number;
model: RerankingModel;
}
let output = '';
let errorOutput = '';
interface SearchTheWebUnrankedConfig {
rerank?: false;
maxResults: number;
}
child.stdout.on('data', (data) => {
output += data.toString();
});
type SearchTheWebConfig = SearchTheWebRerankedConfig | SearchTheWebUnrankedConfig;
child.stderr.on('data', (data) => {
errorOutput += data.toString();
});
child.on('close', (exitCode) => {
if (exitCode !== 0) {
reject(errorOutput || `Exit code ${exitCode}`);
} else {
resolve(output.trim());
export const searchTheWeb = (config: SearchTheWebConfig) => {
return async (query: string) => {
const results = await $fetch<any>(`${process.env.SEARXNG_URL}/search`, {
query: {
q: query,
format: 'json',
}
});
// Send the agent's code to the wrapper via stdin
child.stdin.write(code);
child.stdin.end();
});
const sites = results.results.map((item: any) => ({
title: item.title,
link: item.url,
snippet: item.content,
}));
if (config.rerank) {
const { ranking } = await rerank({
model: config.model,
query,
documents: sites.map(site => site.snippet),
});
const ranked_sites = [];
for (let i = 0; i < ranking.length; i++) {
ranked_sites.push(sites[ranking[i]!.originalIndex]);
}
return ranked_sites.slice(0, config.maxResults || 10);
}
return sites.slice(0, config.maxResults || 10);
};
};
// TODO: obviously come up with a better way for the user to define their own tools
@@ -462,6 +582,9 @@ async function generateResponse(
model: Model,
parameters?: Record<string, any>,
},
search: false | {
config: SearchTheWebConfig,
},
generationId: string,
userId: string,
topicId: string,
@@ -477,9 +600,10 @@ async function generateResponse(
let ttft = undefined;
const activeParts = new Map<string, { id: string; accumulatedContent: string; providerOptions?: any }>();
const activeToolCalls = new Set<string>();
const nativeToDbToolCallId = new Map<string, string>();
// TODO: somehow let the user turn on and off tools
const tools = {
const tools: Record<string, Tool> = {
// writeFile: tool({
// inputSchema: z.object({
// path: z.string(),
@@ -504,7 +628,25 @@ async function generateResponse(
bash: bashTool,
};
console.log({ messages });
if (search) {
tools.search = tool({
description: 'Searches the web',
inputSchema: z.object({
query: z.string(),
}),
outputSchema: z.array(
z.object({
title: z.string(),
link: z.string(),
snippet: z.string(),
engine: z.string(),
})
),
execute: async ({ query }) => {
return await searchTheWeb(search.config)(query);
},
});
}
const response = streamText({
model: model.gateway(model.model.externalId),
@@ -518,8 +660,7 @@ async function generateResponse(
}
},
experimental_transform: streamTransoforms,
// a little trick that makes it so that the stream doesnt stop because of tool calls, and will continue an unbounded amount of time and steps
stopWhen: [],
stopWhen: isLoopFinished(),
tools: model.model.capabilities.includes('tools') ? tools : undefined,
onError: async (error: any) => {
// TODO: the docs say "The stream processing will pause until the callback promise is resolved." Suggesting that this error might not be fatal?
@@ -646,8 +787,10 @@ async function generateResponse(
case 'reasoning-start': {
const type = token.type.split('-')[0] as 'text' | 'reasoning';
const key = `${type}-${curStepIdx}`;
const nativeId = nanoid();
const [part] = await db.insert(messageParts).values({
id: formatPartId(type, nativeId),
userId,
topicId,
messageId: message.id,
@@ -772,9 +915,11 @@ async function generateResponse(
key = `tool-call-${curStepIdx}`;
const toolCallId = token.id;
const dbToolCallId = formatToolCallId(toolCallId);
nativeToDbToolCallId.set(toolCallId, dbToolCallId);
const [toolCall] = await db.insert(toolCalls).values({
id: toolCallId,
id: dbToolCallId,
userId: userId,
toolName: token.toolName,
status: 'pending',
@@ -789,10 +934,11 @@ async function generateResponse(
}
const [part] = await db.insert(messageParts).values({
id: formatPartId('tool-call', dbToolCallId),
userId,
topicId,
messageId: message.id,
toolCallId,
toolCallId: dbToolCallId,
type: 'tool-call',
content: null,
finished: false,
@@ -816,7 +962,7 @@ async function generateResponse(
}
})
activeToolCalls.add(toolCallId);
activeToolCalls.add(dbToolCallId);
activeParts.set(key, { id: part.id, accumulatedContent: '' });
} break;
@@ -838,7 +984,9 @@ async function generateResponse(
break;
}
if (activeToolCalls.has(token.toolCallId)) {
const dbToolCallIdFromMap = nativeToDbToolCallId.get(token.toolCallId);
if (dbToolCallIdFromMap && activeToolCalls.has(dbToolCallIdFromMap)) {
const dbToolCallId = dbToolCallIdFromMap;
await db.update(toolCalls)
.set({
status: 'pending',
@@ -847,13 +995,13 @@ async function generateResponse(
value: inputValue,
},
})
.where(eq(toolCalls.id, token.toolCallId));
.where(eq(toolCalls.id, dbToolCallId));
await topicEvents.emit(topicId, {
type: 'tool-call-delta',
payload: {
messageId: message.id,
toolCallId: token.toolCallId,
toolCallId: dbToolCallId,
toolName: token.toolName,
input: {
type: inputType,
@@ -865,9 +1013,11 @@ async function generateResponse(
key = `tool-call-${curStepIdx}`;
const toolCallId = token.toolCallId;
const dbToolCallId = formatToolCallId(toolCallId);
nativeToDbToolCallId.set(toolCallId, dbToolCallId);
const [toolCall] = await db.insert(toolCalls).values({
id: toolCallId,
id: dbToolCallId,
userId: userId,
toolName: token.toolName,
status: 'pending',
@@ -885,10 +1035,11 @@ async function generateResponse(
}
const [part] = await db.insert(messageParts).values({
id: formatPartId('tool-call', dbToolCallId),
userId,
topicId,
messageId: message.id,
toolCallId: token.toolCallId,
toolCallId: dbToolCallId,
type: 'tool-call',
content: null,
finished: false,
@@ -911,7 +1062,7 @@ async function generateResponse(
}
})
activeToolCalls.add(toolCallId);
activeToolCalls.add(dbToolCallId);
activeParts.set(key, { id: part.id, accumulatedContent: '' });
}
@@ -919,6 +1070,7 @@ async function generateResponse(
case 'tool-result': {
let outputType: ToolCallType = ToolCallType.Text;
let outputValue: string = '';
const dbToolCallId = nativeToDbToolCallId.get(token.toolCallId);
switch (typeof token.output) {
case 'string':
@@ -931,52 +1083,56 @@ async function generateResponse(
break;
default:
console.error('Unknown output type', token.output);
await db.update(toolCalls).set({
status: 'failed',
error: { type: ToolCallType.Text, value: 'Tool returned invalid output' }
}).where(eq(toolCalls.id, token.toolCallId));
await topicEvents.emit(topicId, {
type: 'tool-call-delta',
payload: {
messageId: message.id,
toolCallId: token.toolCallId,
toolName: token.toolName,
output: {
type: outputType,
value: outputValue,
},
if (dbToolCallId) {
await db.update(toolCalls).set({
status: 'failed',
}
})
error: { type: ToolCallType.Text, value: 'Tool returned invalid output' }
}).where(eq(toolCalls.id, dbToolCallId));
activeToolCalls.delete(token.toolCallId);
await topicEvents.emit(topicId, {
type: 'tool-call-delta',
payload: {
messageId: message.id,
toolCallId: dbToolCallId,
toolName: token.toolName,
output: {
type: outputType,
value: outputValue,
},
status: 'failed',
}
})
activeToolCalls.delete(dbToolCallId);
}
break;
}
await db.update(toolCalls).set({
status: 'completed',
output: {
type: outputType,
value: outputValue,
},
}).where(eq(toolCalls.id, token.toolCallId));
await topicEvents.emit(topicId, {
type: 'tool-call-delta',
payload: {
messageId: message.id,
toolCallId: token.toolCallId,
toolName: token.toolName,
if (dbToolCallId) {
await db.update(toolCalls).set({
status: 'completed',
output: {
type: outputType,
value: outputValue,
},
status: 'completed',
}
})
}).where(eq(toolCalls.id, dbToolCallId));
activeToolCalls.delete(token.toolCallId);
await topicEvents.emit(topicId, {
type: 'tool-call-delta',
payload: {
messageId: message.id,
toolCallId: dbToolCallId,
toolName: token.toolName,
output: {
type: outputType,
value: outputValue,
},
status: 'completed',
}
})
activeToolCalls.delete(dbToolCallId);
}
break;
}
@@ -1002,20 +1158,21 @@ async function generateResponse(
break;
}
if (activeToolCalls.has(token.toolCallId)) {
const existingDbToolCallId = nativeToDbToolCallId.get(token.toolCallId);
if (existingDbToolCallId && activeToolCalls.has(existingDbToolCallId)) {
await db.update(toolCalls).set({
status: 'failed',
error: {
type: outputType as ToolCallType,
value: outputValue as string,
}
}).where(eq(toolCalls.id, token.toolCallId));
}).where(eq(toolCalls.id, existingDbToolCallId));
await topicEvents.emit(topicId, {
type: 'tool-call-delta',
payload: {
messageId: message.id,
toolCallId: token.toolCallId,
toolCallId: existingDbToolCallId,
toolName: token.toolName,
error: {
type: outputType as ToolCallType,
@@ -1027,9 +1184,11 @@ async function generateResponse(
key = `tool-call-${curStepIdx}`;
const toolCallId = token.toolCallId;
const dbToolCallId = formatToolCallId(toolCallId);
nativeToDbToolCallId.set(toolCallId, dbToolCallId);
const [toolCall] = await db.insert(toolCalls).values({
id: toolCallId,
id: dbToolCallId,
userId: userId,
toolName: token.toolName,
status: 'failed',
@@ -1047,10 +1206,11 @@ async function generateResponse(
}
const [part] = await db.insert(messageParts).values({
id: formatPartId('tool-call', dbToolCallId),
userId,
topicId,
messageId: message.id,
toolCallId: token.toolCallId,
toolCallId: dbToolCallId,
type: 'tool-call',
content: null,
finished: false,
@@ -1073,7 +1233,9 @@ async function generateResponse(
activeParts.set(key, { id: part.id, accumulatedContent: '' });
}
activeToolCalls.delete(token.toolCallId);
if (existingDbToolCallId) {
activeToolCalls.delete(existingDbToolCallId);
}
break;
}