diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index ca40ec2c8c..1d86646629 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -912,6 +912,7 @@ export enum ScraperProviders { export enum RerankerTypes { JINA = 'jina', COHERE = 'cohere', + CUSTOM = 'custom', } export enum SafeSearchTypes { @@ -930,6 +931,9 @@ export const webSearchSchema = z.object({ jinaApiKey: z.string().optional().default('${JINA_API_KEY}'), jinaApiUrl: z.string().optional().default('${JINA_API_URL}'), cohereApiKey: z.string().optional().default('${COHERE_API_KEY}'), + customRerankerApiUrl: z.string().optional().default('${CUSTOM_RERANKER_API_URL}'), + customRerankerApiKey: z.string().optional().default('${CUSTOM_RERANKER_API_KEY}'), + customRerankerModel: z.string().optional().default('${CUSTOM_RERANKER_MODEL}'), searchProvider: z.nativeEnum(SearchProviders).optional(), scraperProvider: z.nativeEnum(ScraperProviders).optional(), rerankerType: z.nativeEnum(RerankerTypes).optional(), diff --git a/packages/data-schemas/src/app/web.ts b/packages/data-schemas/src/app/web.ts index a61e1f1611..5e46d84322 100644 --- a/packages/data-schemas/src/app/web.ts +++ b/packages/data-schemas/src/app/web.ts @@ -31,6 +31,12 @@ export const webSearchAuth = { jinaApiUrl: 0 as const, }, cohere: { cohereApiKey: 1 as const }, + custom: { + customRerankerApiUrl: 1 as const, + customRerankerModel: 1 as const, + /** Optional (0) */ + customRerankerApiKey: 0 as const, + }, }, }; @@ -72,6 +78,9 @@ export function loadWebSearchConfig( const jinaApiKey = config?.jinaApiKey ?? '${JINA_API_KEY}'; const jinaApiUrl = config?.jinaApiUrl ?? '${JINA_API_URL}'; const cohereApiKey = config?.cohereApiKey ?? '${COHERE_API_KEY}'; + const customRerankerApiUrl = config?.customRerankerApiUrl ?? '${CUSTOM_RERANKER_API_URL}'; + const customRerankerApiKey = config?.customRerankerApiKey ?? '${CUSTOM_RERANKER_API_KEY}'; + const customRerankerModel = config?.customRerankerModel ?? '${CUSTOM_RERANKER_MODEL}'; const safeSearch = config?.safeSearch ?? SafeSearchTypes.MODERATE; return { @@ -80,6 +89,9 @@ export function loadWebSearchConfig( jinaApiKey, jinaApiUrl, cohereApiKey, + customRerankerApiUrl, + customRerankerApiKey, + customRerankerModel, serperApiKey, searxngApiKey, firecrawlApiKey, diff --git a/packages/data-schemas/src/types/web.ts b/packages/data-schemas/src/types/web.ts index a9cc1f0cc6..dcd0dcd320 100644 --- a/packages/data-schemas/src/types/web.ts +++ b/packages/data-schemas/src/types/web.ts @@ -9,7 +9,10 @@ export type TWebSearchKeys = | 'firecrawlVersion' | 'jinaApiKey' | 'jinaApiUrl' - | 'cohereApiKey'; + | 'cohereApiKey' + | 'customRerankerApiUrl' + | 'customRerankerApiKey' + | 'customRerankerModel'; export type TWebSearchCategories = | SearchCategories.PROVIDERS