Skip to content

Commit

Permalink
支持模型名称包含silent_search来静默搜索不输出搜索过程
Browse files Browse the repository at this point in the history
  • Loading branch information
@Vinlic
Vinlic committed Mar 19, 2024
1 parent b8134a6 commit 909796b
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 18 deletions.
6 changes: 6 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,8 @@ Authorization: Bearer [refresh_token]
请求数据：
```json
{
// 模型名称随意填写，如果不希望输出检索过程模型名称请包含silent_search
"model": "kimi",
"messages": [
{
"role": "user",
Expand Down Expand Up @@ -254,6 +256,8 @@ Authorization: Bearer [refresh_token]
请求数据：
```json
{
// 模型名称随意填写，如果不希望输出检索过程模型名称请包含silent_search
"model": "kimi",
"messages": [
{
"role": "user",
Expand Down Expand Up @@ -318,6 +322,8 @@ Authorization: Bearer [refresh_token]
请求数据：
```json
{
// 模型名称随意填写，如果不希望输出检索过程模型名称请包含silent_search
"model": "kimi",
"messages": [
{
"role": "user",
Expand Down
2 changes: 1 addition & 1 deletion package.json
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "kimi-free-api",
"version": "0.0.15",
"version": "0.0.16",
"description": "Kimi Free API Server",
"type": "module",
"main": "dist/index.js",
Expand Down
36 changes: 21 additions & 15 deletions src/api/controllers/chat.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,12 +164,13 @@ async function removeConversation(convId: string, refreshToken: string) {
/**
* 同步对话补全
*
* @param model 模型名称
* @param messages 参考gpt系列消息格式，多轮对话请完整提供上下文
* @param refreshToken 用于刷新access_token的refresh_token
* @param useSearch 是否开启联网搜索
* @param retryCount 重试次数
*/
async function createCompletion(messages: any[], refreshToken: string, useSearch = true, retryCount = 0) {
async function createCompletion(model = MODEL_NAME, messages: any[], refreshToken: string, useSearch = true, retryCount = 0) {
return (async () => {
logger.info(messages);

Expand Down Expand Up @@ -204,7 +205,7 @@ async function createCompletion(messages: any[], refreshToken: string, useSearch

const streamStartTime = util.timestamp();
// 接收流为输出文本
const answer = await receiveStream(convId, result.data);
const answer = await receiveStream(model, convId, result.data);
logger.success(`Stream has completed transfer ${util.timestamp() - streamStartTime}ms`);

// 异步移除会话，如果消息不合规，此操作可能会抛出数据库错误异常，请忽略
Expand All @@ -219,7 +220,7 @@ async function createCompletion(messages: any[], refreshToken: string, useSearch
logger.warn(`Try again after ${RETRY_DELAY / 1000}s...`);
return (async () => {
await new Promise(resolve => setTimeout(resolve, RETRY_DELAY));
return createCompletion(messages, refreshToken, useSearch, retryCount + 1);
return createCompletion(model, messages, refreshToken, useSearch, retryCount + 1);
})();
}
throw err;
Expand All @@ -229,12 +230,13 @@ async function createCompletion(messages: any[], refreshToken: string, useSearch
/**
* 流式对话补全
*
* @param model 模型名称
* @param messages 参考gpt系列消息格式，多轮对话请完整提供上下文
* @param refreshToken 用于刷新access_token的refresh_token
* @param useSearch 是否开启联网搜索
* @param retryCount 重试次数
*/
async function createCompletionStream(messages: any[], refreshToken: string, useSearch = true, retryCount = 0) {
async function createCompletionStream(model = MODEL_NAME, messages: any[], refreshToken: string, useSearch = true, retryCount = 0) {
return (async () => {
logger.info(messages);

Expand Down Expand Up @@ -268,7 +270,7 @@ async function createCompletionStream(messages: any[], refreshToken: string, use
});
const streamStartTime = util.timestamp();
// 创建转换流将消息格式转换为gpt兼容格式
return createTransStream(convId, result.data, () => {
return createTransStream(model, convId, result.data, () => {
logger.success(`Stream has completed transfer ${util.timestamp() - streamStartTime}ms`);
// 流传输结束后异步移除会话，如果消息不合规，此操作可能会抛出数据库错误异常，请忽略
removeConversation(convId, refreshToken)
Expand All @@ -281,7 +283,7 @@ async function createCompletionStream(messages: any[], refreshToken: string, use
logger.warn(`Try again after ${RETRY_DELAY / 1000}s...`);
return (async () => {
await new Promise(resolve => setTimeout(resolve, RETRY_DELAY));
return createCompletionStream(messages, refreshToken, useSearch, retryCount + 1);
return createCompletionStream(model, messages, refreshToken, useSearch, retryCount + 1);
})();
}
throw err;
Expand Down Expand Up @@ -541,15 +543,16 @@ function checkResult(result: AxiosResponse, refreshToken: string) {
/**
* 从流接收完整的消息内容
*
* @param model 模型名称
* @param convId 会话ID
* @param stream 消息流
*/
async function receiveStream(convId: string, stream: any) {
async function receiveStream(model: string, convId: string, stream: any) {
return new Promise((resolve, reject) => {
// 消息初始化
const data = {
id: convId,
model: MODEL_NAME,
model,
object: 'chat.completion',
choices: [
{ index: 0, message: { role: 'assistant', content: '' }, finish_reason: 'stop' }
Expand All @@ -558,6 +561,7 @@ async function receiveStream(convId: string, stream: any) {
created: util.unixTimestamp()
};
let refContent = '';
const silentSearch = model.indexOf('silent_search') != -1;
const parser = createParser(event => {
try {
if (event.type !== "event") return;
Expand All @@ -576,7 +580,7 @@ async function receiveStream(convId: string, stream: any) {
resolve(data);
}
// 处理联网搜索
else if (result.event == 'search_plus' && result.msg && result.msg.type == 'get_res')
else if (!silentSearch && result.event == 'search_plus' && result.msg && result.msg.type == 'get_res')
refContent += `${result.msg.title}(${result.msg.url})\n`;
// else
// logger.warn(result.event, result);
Expand All @@ -598,19 +602,21 @@ async function receiveStream(convId: string, stream: any) {
*
* 将流格式转换为gpt兼容流格式
*
* @param model 模型名称
* @param convId 会话ID
* @param stream 消息流
* @param endCallback 传输结束回调
*/
function createTransStream(convId: string, stream: any, endCallback?: Function) {
function createTransStream(model: string, convId: string, stream: any, endCallback?: Function) {
// 消息创建时间
const created = util.unixTimestamp();
// 创建转换流
const transStream = new PassThrough();
let searchFlag = false;
const silentSearch = model.indexOf('silent_search') != -1;
!transStream.closed && transStream.write(`data: ${JSON.stringify({
id: convId,
model: MODEL_NAME,
model,
object: 'chat.completion.chunk',
choices: [
{ index: 0, delta: { role: 'assistant', content: '' }, finish_reason: null }
Expand All @@ -628,7 +634,7 @@ function createTransStream(convId: string, stream: any, endCallback?: Function)
if (result.event == 'cmpl') {
const data = `data: ${JSON.stringify({
id: convId,
model: MODEL_NAME,
model,
object: 'chat.completion.chunk',
choices: [
{ index: 0, delta: { content: (searchFlag ? '\n' : '') + result.text }, finish_reason: null }
Expand All @@ -643,7 +649,7 @@ function createTransStream(convId: string, stream: any, endCallback?: Function)
else if (result.event == 'all_done' || result.event == 'error') {
const data = `data: ${JSON.stringify({
id: convId,
model: MODEL_NAME,
model,
object: 'chat.completion.chunk',
choices: [
{
Expand All @@ -660,12 +666,12 @@ function createTransStream(convId: string, stream: any, endCallback?: Function)
endCallback && endCallback();
}
// 处理联网搜索
else if (result.event == 'search_plus' && result.msg && result.msg.type == 'get_res') {
else if (!silentSearch && result.event == 'search_plus' && result.msg && result.msg.type == 'get_res') {
if (!searchFlag)
searchFlag = true;
const data = `data: ${JSON.stringify({
id: convId,
model: MODEL_NAME,
model,
object: 'chat.completion.chunk',
choices: [
{
Expand Down
5 changes: 3 additions & 2 deletions src/api/routes/chat.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,16 @@ export default {
const tokens = chat.tokenSplit(request.headers.authorization);
// 随机挑选一个refresh_token
const token = _.sample(tokens);
const model = request.body.model;
const messages = request.body.messages;
if (request.body.stream) {
const stream = await chat.createCompletionStream(request.body.messages, token, request.body.use_search);
const stream = await chat.createCompletionStream(model, messages, token, request.body.use_search);
return new Response(stream, {
type: "text/event-stream"
});
}
else
return await chat.createCompletion(messages, token, request.body.use_search);
return await chat.createCompletion(model, messages, token, request.body.use_search);
}

}
Expand Down

0 comments on commit 909796b

Please sign in to comment.