feat: remove search rerank
This commit is contained in:
@@ -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(),
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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)));
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -67,11 +67,3 @@ export async function getProviderDetails(provider: ProviderDrizzle, providerApiK
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
export function isRerankingProvider(gateway: any): gateway is { reranking: (modelId: string) => any } {
|
||||
return (
|
||||
gateway !== null &&
|
||||
'reranking' in gateway &&
|
||||
typeof (gateway as any).reranking === 'function'
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
import type { Result } from "../../types/result";
|
||||
import { z } from "zod";
|
||||
import { Err, Ok } from "../../types/result";
|
||||
|
||||
const searchResponseSchema = z.object({
|
||||
results: z.array(z.object({
|
||||
title: z.string(),
|
||||
url: z.string(),
|
||||
content: z.string().optional().default(""),
|
||||
})),
|
||||
unresponsive_engines: z.array(z.tuple([z.string(), z.string()])).optional().default([]),
|
||||
});
|
||||
|
||||
export const searchInputSchema = z.object({
|
||||
query: z.string(),
|
||||
results: z.number().int().positive().describe("Number of search results to return; choose this for every search."),
|
||||
});
|
||||
|
||||
export const searchOutputSchema = z.object({
|
||||
results: z.array(z.object({
|
||||
title: z.string(),
|
||||
link: z.string(),
|
||||
snippet: z.string(),
|
||||
})),
|
||||
warnings: z.array(z.object({ engine: z.string(), reason: z.string() })),
|
||||
});
|
||||
|
||||
type SearchOutput = z.infer<typeof searchOutputSchema>;
|
||||
|
||||
export const limitSearchResults = (output: SearchOutput, results: number): SearchOutput => ({
|
||||
results: output.results.slice(0, results),
|
||||
warnings: output.warnings,
|
||||
});
|
||||
|
||||
export const readSearchResponse = (body: unknown): Result<SearchOutput, string> => {
|
||||
const parsed = searchResponseSchema.safeParse(body);
|
||||
if (!parsed.success) {
|
||||
return Err("Search returned an unexpected response format.");
|
||||
}
|
||||
const warnings = parsed.data.unresponsive_engines.map(([engine, reason]) => ({ engine, reason }));
|
||||
if (parsed.data.results.length === 0 && warnings.length > 0) {
|
||||
return Err(`Search returned no results because engines failed: ${warnings.map(warning => `${warning.engine}: ${warning.reason}`).join("; ")}`);
|
||||
}
|
||||
return Ok({
|
||||
results: parsed.data.results.map(item => ({
|
||||
title: item.title,
|
||||
link: item.url,
|
||||
snippet: item.content,
|
||||
})),
|
||||
warnings,
|
||||
});
|
||||
};
|
||||
Reference in New Issue
Block a user