| import { Context } from 'koa'; |
| import { Rag } from './rag'; |
| import platform from './provider'; |
| import { DefaultQuery, Models } from './utils/constant'; |
| import { searchWithSogou } from './service'; |
| import { ESearchEngine, IChatInputMessage, Provider, TMode } from './interface'; |
| import { getFromCache, setToCache } from './cache'; |
| import { ESearXNGCategory } from './search/searxng'; |
|
|
| const CACHE_ENABLED = process.env.CACHE_ENABLE; |
|
|
| export const searchController = async (ctx: Context) => { |
| const stream = ctx.request.body.stream ?? true; |
| const q = ctx.request.query.q || DefaultQuery; |
| const model: string = ctx.request.body.model; |
| const reload: boolean = ctx.request.body.reload ?? false; |
| const engine: ESearchEngine = ctx.request.body.engine; |
| const locally: boolean = ctx.request.body.locally ?? false; |
| const categories: ESearXNGCategory[] = ctx.request.body.categories ?? []; |
| const mode: TMode = ctx.request.body.mode ?? 'simple'; |
| const language: string = ctx.request.body.language || 'all'; |
| const provider: Provider = ctx.request.body.provider || 'ollama'; |
|
|
| ctx.res.setHeader('Content-Type', 'text/event-stream'); |
| ctx.res.setHeader('Cache-Control', 'no-cache'); |
| ctx.res.setHeader('Connection', 'keep-alive'); |
| ctx.res.statusCode = 200; |
|
|
| |
| if (!reload) { |
| const cached = await getFromCache(q as string, mode, categories); |
| if (cached) { |
| ctx.body = cached; |
| ctx.res.write(cached, 'utf-8'); |
| ctx.res.end(); |
| return; |
| } |
| } |
|
|
| const rag = new Rag({ |
| stream, |
| model, |
| engine, |
| locally, |
| provider |
| }); |
|
|
| if (!stream) { |
| const res = await rag.query(q as string); |
| ctx.body = res; |
| return; |
| } |
|
|
| let result = ''; |
|
|
| await rag.query(q as string, categories, mode, language, (json: string) => { |
| const eventData = `data:${JSON.stringify({ data: json })}\n\n`; |
| result += eventData; |
| ctx.res.write(eventData, 'utf-8'); |
| }); |
|
|
| ctx.res.end(); |
| |
| if (CACHE_ENABLED === '1') { |
| setToCache(q as string, result, mode, categories); |
| } |
| }; |
|
|
| export const sogouSearchController = async (ctx: Context) => { |
| const q = ctx.request.query.q || DefaultQuery; |
| const res = await searchWithSogou(q as string); |
| ctx.body = res; |
| }; |
|
|
| export const chatStreamController = async (ctx: Context) => { |
| const messages: IChatInputMessage[] = ctx.request.body.messages || []; |
| const system: string | undefined = ctx.request.body.system; |
| const model: string | undefined = ctx.request.body.model; |
| const locally: boolean = ctx.request.body.locally ?? false; |
| const provider: Provider = ctx.request.body.provider ?? 'ollama'; |
|
|
| if (!model) throw new Error('model is required'); |
|
|
| ctx.res.setHeader('Content-Type', 'text/event-stream'); |
| ctx.res.setHeader('Cache-Control', 'no-cache'); |
| ctx.res.setHeader('Connection', 'keep-alive'); |
| const handler = locally ? platform[provider].chatStream.bind(platform[provider]) : processModel(model); |
| ctx.res.statusCode = 200; |
|
|
| await handler?.(messages, (data: any) => { |
| const eventData = `data: ${JSON.stringify({ text: data || '' })}\n\n`; |
| ctx.res.write(eventData); |
| }, model, system); |
|
|
| ctx.res.end(); |
| }; |
|
|
| export const modelsController = async (ctx: Context) => { |
| const { GOOGLE_KEY, ALIYUN_KEY, OPENAI_KEY, BAIDU_KEY, TENCENT_KEY, YI_KEY, MOONSHOT_KEY, LEPTON_KEY, DEEPSEEK_KEY, GLM_KEY } = process.env; |
| const keys: Record<string, string | undefined> = { |
| google: GOOGLE_KEY, |
| aliyun: ALIYUN_KEY, |
| openai: OPENAI_KEY, |
| baidu: BAIDU_KEY, |
| tencent: TENCENT_KEY, |
| yi: YI_KEY, |
| moonshot: MOONSHOT_KEY, |
| lepton: LEPTON_KEY, |
| chatglm: GLM_KEY, |
| deepseek: DEEPSEEK_KEY |
| }; |
| const models = Models.filter(item => keys[item.platform] !== undefined); |
| const enabledModels: Record<string, string[]> = {}; |
| for (const model of models) { |
| if (keys[model.platform]) enabledModels[model.platform] = model.models; |
| } |
| ctx.body = enabledModels; |
| }; |
|
|
| |
| export const localModelsController = async (ctx: Context) => { |
| const provider: Provider = ctx.URL.searchParams.get('provider') as Provider ?? 'ollama'; |
|
|
| const list = await platform[provider].list(); |
| ctx.body = list; |
| }; |
|
|
|
|
| |
| export const localChatStreamController = async (ctx: Context) => { |
| const messages: IChatInputMessage[] = ctx.request.body.messages || []; |
| const model: string | undefined = ctx.request.body.model; |
| const provider: Provider = ctx.request.body.provider; |
|
|
| ctx.res.setHeader('Content-Type', 'text/event-stream'); |
| ctx.res.setHeader('Cache-Control', 'no-cache'); |
| ctx.res.setHeader('Connection', 'keep-alive'); |
| ctx.res.statusCode = 200; |
| await platform[provider].chatStream(messages, (data) => { |
| const eventData = `data: ${JSON.stringify({ text: data || '' })}\n\n`; |
| ctx.res.write(eventData); |
| }, model); |
| ctx.res.end(); |
| }; |
|
|
|
|
| function processModel(model: string) { |
| const targetModel = Models.find(item => { |
| return item.models.includes(model); |
| }); |
| if (targetModel?.platform) { |
| const target = platform[targetModel.platform]; |
| return target.chatStream.bind(target); |
| } |
| } |
|
|
