File size: 5,137 Bytes
755dd12 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
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;
// get from cache, skip if enable reload
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();
// caching
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;
};
// locally llm models
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;
};
// chat with locally models
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);
}
}
|