feat: remove search rerank

This commit is contained in:
Zoe
2026-09-27 09:45:19 -05:00
parent bb5795e00d
commit d48d8163a7
17 changed files with 243 additions and 336 deletions
-2
View File
@@ -18,8 +18,6 @@ export default defineEventHandler(async (event) => {
config: z.object({
search: z.object({
enabled: z.boolean(),
rerank: z.boolean(),
maxResults: z.number().min(1).max(50),
}).optional(),
tools: z.object({
python: z.boolean().optional(),
+2 -7
View File
@@ -22,21 +22,16 @@ export default defineEventHandler(async (event) => {
rename: {
enabled: false,
modelId: null,
},
rerank: {
enabled: false,
modelId: null,
}
};
if (existing) {
const missingSystemAssistants = Object.keys(defaultSystemAssistants).filter(key => !(key in existing!.systemAssistants));
if (missingSystemAssistants.length > 0) {
if (!("rename" in existing.systemAssistants) || "rerank" in existing.systemAssistants) {
const [updated] = await db.update(settings)
.set({
systemAssistants: {
...defaultSystemAssistants,
...existing.systemAssistants,
rename: existing.systemAssistants.rename ?? defaultSystemAssistants.rename,
},
})
.where(eq(settings.id, existing.id)).returning();
+13 -14
View File
@@ -12,7 +12,12 @@ const appearanceSchema = z.object({
});
const updateSettingsSchema = z.object({
systemAssistants: z.any().optional(),
systemAssistants: z.object({
rename: z.object({
enabled: z.boolean(),
modelId: z.string().nullable(),
}).partial().optional(),
}).optional(),
appearance: appearanceSchema.optional(),
});
@@ -53,28 +58,22 @@ export default defineEventHandler(async (event) => {
return { ...base, ...incoming };
};
const mergeSystemAssistants = (
existingSA: Record<string, unknown> | null,
incoming: typeof updates.systemAssistants,
): Record<string, unknown> => {
const base = existingSA || {};
if (!incoming) return base;
return { ...base, ...incoming };
const newSystemAssistants = {
rename: {
enabled: updates.systemAssistants?.rename?.enabled ?? existing.systemAssistants.rename?.enabled ?? false,
modelId: updates.systemAssistants?.rename?.modelId !== undefined
? updates.systemAssistants.rename.modelId
: existing.systemAssistants.rename?.modelId ?? null,
},
};
const newAppearance = mergeAppearance(
existing.appearance as Record<string, unknown> | null,
updates.appearance,
);
const newSystemAssistants = mergeSystemAssistants(
existing.systemAssistants as Record<string, unknown> | null,
updates.systemAssistants,
);
await db.update(settings)
.set({
appearance: newAppearance,
// @ts-ignore - I dont have the energy to type this correctly right now
systemAssistants: newSystemAssistants,
})
.where(and(eq(settings.id, existingId)));
+25 -161
View File
@@ -2,7 +2,7 @@ import { db } from "~~/server/lib/db";
import * as z from 'zod';
import { type MessageEntity } from '~/composables/useChat';
import { promises as fs } from 'fs';
import { isLoopFinished, type ModelMessage, rerank, type RerankingModel, streamText, type StreamTextTransform, type Tool, tool } from "ai";
import { isLoopFinished, type ModelMessage, streamText, type StreamTextTransform, type Tool, tool } from "ai";
import { generations, messageParts, messages, toolCalls, ToolCallType } from "~~/drizzle/schema";
import { topicEvents } from "~~/server/utils/events";
import { nanoid } from "nanoid";
@@ -11,10 +11,10 @@ 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";
import { generateFileToken } from "~~/server/utils/file-token";
import { serializeToolError } from "~~/utils/tool-errors";
import { readFirecrawlResponse } from "~~/server/utils/firecrawl";
import { limitSearchResults, readSearchResponse, searchInputSchema, searchOutputSchema } from "~~/server/utils/search";
export default defineEventHandler(async (event) => {
await protectRoute(event);
@@ -25,10 +25,6 @@ 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);
@@ -40,7 +36,7 @@ export default defineEventHandler(async (event) => {
});
}
const { parentMessageId, modelId, rerank: rerankConfig, providerApiKey, args } = result.data;
const { parentMessageId, modelId, providerApiKey, args } = result.data;
const topic = await db.query.topics.findFirst({
where: {
@@ -116,90 +112,7 @@ 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,
}
});
}
}
}
const searchEnabled = topic.agent.config?.search?.enabled === true;
let agentmessage = await db.transaction(async tx => {
const generationId = nanoid();
@@ -332,7 +245,7 @@ export default defineEventHandler(async (event) => {
generateResponse(
agentmessage as MessageEntity,
{ gateway: gateway.gateway, model, parameters: args },
searchParam,
searchEnabled,
{
python: topic.agent.config?.tools?.python ?? false,
},
@@ -418,61 +331,23 @@ const evalPython = async (code: string): Promise<string> => {
}
};
interface SearchTheWebRerankedConfig {
rerank: true;
maxResults: number;
model: RerankingModel;
}
export const searchTheWeb = async (query: string, results: number) => {
const searchUrl = new URL(process.env.SEARCH_API_URL ?? process.env.SEARXNG_URL ?? "");
searchUrl.pathname = `${searchUrl.pathname.replace(/\/+$/, "")}/search`;
interface SearchTheWebUnrankedConfig {
rerank?: false;
maxResults: number;
}
type SearchTheWebConfig = SearchTheWebRerankedConfig | SearchTheWebUnrankedConfig;
export const searchTheWeb = (config: SearchTheWebConfig) => {
return async (query: string) => {
const searchUrl = new URL(process.env.SEARCH_API_URL ?? process.env.SEARXNG_URL ?? "");
searchUrl.pathname = `${searchUrl.pathname.replace(/\/+$/, "")}/search`;
const results = await $fetch<{
results: Array<{
title: string;
url: string;
content: string;
}>;
}>(searchUrl.toString(), {
query: {
q: query,
format: 'json',
}
});
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);
const response = await $fetch<unknown>(searchUrl.toString(), {
query: {
q: query,
format: 'json',
}
});
return sites.slice(0, config.maxResults || 10);
};
const result = readSearchResponse(response);
if (!result.ok) {
// Reject execution so the SDK emits a tool-error rather than an empty success.
return Promise.reject(new Error(result.error));
}
return limitSearchResults(result.data, results);
};
const fetchUrlTool = tool({
@@ -529,9 +404,7 @@ async function generateResponse(
model: Model,
parameters?: Record<string, any>,
},
search: false | {
config: SearchTheWebConfig,
},
search: boolean,
enabledTools: {
python?: boolean;
},
@@ -560,20 +433,11 @@ async function generateResponse(
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);
description: 'Searches the web. Choose how many results to return on every call using the required results argument.',
inputSchema: searchInputSchema,
outputSchema: searchOutputSchema,
execute: async ({ query, results }) => {
return await searchTheWeb(query, results);
},
});
tools.fetchUrl = fetchUrlTool;