|
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); |
|
} |
|
} |
|
|