feat: add chat message

This commit is contained in:
2024-10-01 00:47:14 +08:00
parent e05c042827
commit bd38ad2eaa
20 changed files with 556 additions and 10 deletions

View File

@@ -0,0 +1,184 @@
import { app } from '@/app.ts';
import { AiAgent } from '@/models/agent.ts';
import { ChatPrompt } from '@/models/chat-prompt.ts';
import { ChatSession } from '@/models/chat-session.ts';
import { Prompt } from '@/models/prompt.ts';
import { agentManger } from '@kevisual/ai-lang';
import { PromptTemplate } from '@kevisual/ai-graph';
import { v4 } from 'uuid';
import { ChatHistory } from '@/models/chat-history.ts';
import { User } from '@/models/user.ts';
const clients = [];
export const compotedToken = () => {
// 计算token消耗
};
export const getConfigByKey = async (key) => {
const chatPrompt = await ChatPrompt.findOne({ where: { key } });
const { promptId, aiAgentId } = chatPrompt.data;
const prompt = await Prompt.findByPk(promptId);
let aiAgent = agentManger.getAgent(aiAgentId);
if (!aiAgent) {
// throw new Error('aiAgent not found');
const aiAgnetModel = await AiAgent.findByPk(aiAgentId);
aiAgent = agentManger.createAgent({
id: aiAgnetModel.id,
type: aiAgnetModel.type as any,
model: aiAgnetModel.model as any,
baseUrl: aiAgnetModel.baseUrl,
apiKey: aiAgnetModel.apiKey,
temperature: aiAgnetModel.temperature,
cache: aiAgnetModel.cache as any,
cacheName: aiAgnetModel.cacheName,
});
}
return { chatPrompt, prompt, aiAgent };
};
export const getTemplate = async ({ data, inputs }) => {
const promptTemplate = new PromptTemplate({
prompt: data.prompt,
inputVariables: inputs.map((item) => {
return {
key: item.key,
value: item.value,
};
}),
localVariables: [],
}); // 传入参数
return await promptTemplate.getTemplate();
};
const onMessage = async ({ data, end, ws }) => {
// messages的 data
const client = clients.find((client) => client.ws === ws);
if (!client) {
end({ code: 404, data: {}, message: 'client not found' });
return;
}
const { uid, id, key } = client.data;
const { inputs, message: sendMessage } = data;
let root = data.root || false;
let chatSession = await ChatSession.findByPk(id);
const config = await getConfigByKey(key);
const { prompt, aiAgent, chatPrompt } = config;
if (!chatSession) {
chatSession = await ChatSession.create({ key, id, data: {}, uid, chatPromptId: chatPrompt.id });
root = true;
} else {
root = false;
const chatHistory = await ChatHistory.findAll({ where: { chatId: id }, logging: false });
chatHistory.forEach((item) => {
end({ code: 200, data: item, message: 'success', type: 'messages' });
});
return;
}
const template = await getTemplate({ data: prompt.presetData.data, inputs });
if (!template) {
end({ code: 404, data: {}, message: 'template not found' });
return;
}
const userQuestion = template || sendMessage;
// 保存到数据库
const roleUser = await ChatHistory.create({
data: {
message: userQuestion,
},
chatId: id,
chatPromptId: chatPrompt.id,
root: root,
uid: uid,
show: true,
role: 'user',
});
end({ code: 200, data: roleUser, message: 'success', type: 'messages' });
const result = await aiAgent.sendHumanMessage(template, { thread_id: id });
const lastMessage = result.messages[result.messages.length - 1];
const message = result.messages[result.messages.length - 1].content;
// 根据key找到对应的prompt
// 保存到数据库
const roleAi = await ChatHistory.create({
data: {
message,
result: lastMessage,
},
chatId: id,
chatPromptId: chatPrompt.id,
root: false,
uid: uid,
show: true,
role: 'ai',
});
end({ code: 200, data: roleAi, message: 'success', type: 'messages' });
};
const getHistory = async (id: string, { data, end, ws }) => {
const chatHistory = await ChatHistory.findAll({ where: { chatId: id }, logging: false });
chatHistory.forEach((item) => {
end({ code: 200, data: item, message: 'success', type: 'messages' });
});
};
app.io.addListener('chat', async ({ data, end, ws }) => {
const { type } = data || {};
if (type === 'subscribe') {
const token = data?.token;
if (!token) {
end({ code: 401, data: {}, message: 'need token' });
return;
}
let tokenUesr;
try {
tokenUesr = await User.verifyToken(token);
} catch (e) {
end({ code: 401, data: {}, message: 'token is invaild' });
return;
}
const uid = tokenUesr.id;
const id = v4();
const clientData = { ...data?.data, uid };
if (!clientData.id) {
clientData.id = id;
}
const client = clients.find((client) => client.ws === ws);
if (!client) {
clients.push({ ws, data: clientData }); // 拆包里面包含的type信息去掉
}
end({ code: 200, data: clientData, message: 'subscribe success' });
} else if (type === 'unsubscribe') {
// 需要手动取消订阅
const index = clients.findIndex((client) => client.ws === ws);
if (index > -1) {
const data = clients[index]?.data;
clients.splice(index, 1);
end({ code: 200, data, message: 'unsubscribe success' });
return;
}
end({ code: 200, data: {}, message: 'unsubscribe success' });
return;
} else if (type === 'messages') {
try {
await onMessage({ data: data.data, end, ws });
} catch (e) {
console.error('onMessage error', e);
end({ code: 500, data: {}, message: 'onMessage error' });
}
return;
} else if (type === 'changeSession') {
// 修改client的session的id
const client = clients.find((client) => client.ws === ws);
if (!client) {
end({ code: 404, data: {}, message: 'client not found' });
return;
}
const { id } = data?.data;
client.data.id = id || v4();
// 返回修改后的history的内容
end({ code: 200, data: client.data, message: 'changeSession success' });
getHistory(id, { data, end, ws });
return;
} else {
end({ code: 404, data: {}, message: 'subscribe fail' });
return;
}
ws.on('close', () => {
const index = clients.findIndex((client) => client.ws === ws);
if (index > -1) clients.splice(index, 1);
});
});

View File

@@ -0,0 +1,3 @@
import './list.ts'
import './session-list.ts'
import './chat-io.ts'

View File

@@ -0,0 +1,33 @@
import { app } from '@/app.ts';
import { ChatHistory } from '@/models/chat-history.ts';
import { CustomError } from '@abearxiong/router';
// Admin only
app
.route({
path: 'chat-history',
key: 'list',
})
.define(async (ctx) => {
const chatPrompt = await ChatHistory.findAll({
order: [['updatedAt', 'DESC']],
});
ctx.body = chatPrompt;
})
.addTo(app);
app
.route({
path: 'chat-history',
key: 'delete',
})
.define(async (ctx) => {
const { id } = ctx.query;
const chatHistory = await ChatHistory.findByPk(id);
if (!chatHistory) {
throw new CustomError('ChatHistory not found');
}
await chatHistory.destroy();
ctx.body = chatHistory;
})
.addTo(app);

View File

@@ -0,0 +1,83 @@
import { app } from '@/app.ts';
import { ChatSession } from '@/models/chat-session.ts';
import { ChatPrompt } from '@/models/chat-prompt.ts';
import { CustomError } from '@abearxiong/router';
app
.route({
path: 'chat-session',
key: 'list',
})
.define(async (ctx) => {
const chatSession = await ChatSession.findAll({
order: [['updatedAt', 'DESC']],
});
ctx.body = chatSession;
})
.addTo(app);
// Admin only
app
.route({
path: 'chat-session',
key: 'list-history',
})
.define(async (ctx) => {
const data = ctx.query.data;
const chatPrompt = await ChatPrompt.findOne({
where: {
key: data.key,
},
});
if (!chatPrompt) {
throw new CustomError('ChatPrompt not found');
}
console.log('chatPrompt', chatPrompt.id);
const chatSession = await ChatSession.findAll({
order: [['updatedAt', 'DESC']],
where: {
chatPromptId: chatPrompt.id,
},
limit: data.limit || 10,
});
ctx.body = chatSession;
})
.addTo(app);
app
.route({
path: 'chat-session',
key: 'update',
middleware: ['auth'],
})
.define(async (ctx) => {
const tokenUser = ctx.state.tokenUser;
const uid = tokenUser.id;
const { id, ...data } = ctx.query.data;
if (id) {
const session = await ChatSession.findByPk(id);
if (session) {
await session.update(data);
} else {
throw new CustomError('Session not found');
}
ctx.body = session;
return;
}
const session = await ChatSession.create({ ...data, uid });
ctx.body = session;
})
.addTo(app);
app
.route({
path: 'chat-session',
key: 'delete',
})
.define(async (ctx) => {
const { id } = ctx.query;
const session = await ChatSession.findByPk(id);
if (!session) {
throw new CustomError('Session not found');
}
await session.destroy();
ctx.body = session;
})
.addTo(app);

View File

@@ -0,0 +1 @@
import './list.ts';

View File

@@ -0,0 +1,126 @@
import { app } from '@/app.ts';
import { AiAgent } from '@/models/agent.ts';
import { ChatPrompt } from '@/models/chat-prompt.ts';
import { Prompt } from '@/models/prompt.ts';
import { CustomError } from '@abearxiong/router';
// Admin only
app
.route({
path: 'chat-prompt',
key: 'list',
})
.define(async (ctx) => {
const chatPrompt = await ChatPrompt.findAll({
order: [['updatedAt', 'DESC']],
});
ctx.body = chatPrompt;
})
.addTo(app);
app
.route({
path: 'chat-prompt',
key: 'get',
validator: {
id: {
type: 'string',
required: true,
},
},
})
.define(async (ctx) => {
const { id } = ctx.query;
const chatPrompt = await ChatPrompt.findByPk(id);
if (!chatPrompt) {
throw new CustomError('ChatPrompt not found');
}
ctx.body = chatPrompt;
})
.addTo(app);
app
.route({
path: 'chat-prompt',
key: 'update',
middleware: ['auth'],
})
.define(async (ctx) => {
const tokenUser = ctx.state.tokenUser;
const uid = tokenUser.id;
const { data } = ctx.query;
const { id, ...rest } = data;
if (id) {
const page = await ChatPrompt.findByPk(id);
if (page) {
if (rest.data) {
rest.data = { ...page.data, ...rest.data };
}
const newPage = await page.update(rest);
ctx.body = newPage;
} else {
throw new CustomError('page not found');
}
} else if (data) {
const page = await ChatPrompt.create({ ...rest, uid });
ctx.body = page;
}
})
.addTo(app);
app
.route({
path: 'chat-prompt',
key: 'delete',
validator: {
id: {
type: 'string',
required: true,
},
},
})
.define(async (ctx) => {
const id = ctx.query.id;
const chatPrompt = await ChatPrompt.findByPk(id);
if (!chatPrompt) {
throw new CustomError('chatPrompt not found');
}
await chatPrompt.destroy();
ctx.body = 'success';
})
.addTo(app);
app
.route({
path: 'chat-prompt',
key: 'getByKey',
})
.define(async (ctx) => {
const { key } = ctx.query.data || {};
if (!key) {
throw new CustomError('key is required');
}
const chatPrompt = await ChatPrompt.findOne({
where: { key },
});
if (!chatPrompt) {
throw new CustomError('chatPrompt not found');
}
const { promptId, aiAgentId } = chatPrompt.data;
if (!aiAgentId) {
throw new CustomError('promptId is required');
}
const aiAgent = await AiAgent.findByPk(aiAgentId, {
// 只获取 id 和description 字段
attributes: ['id', 'description', 'key'],
});
if (!aiAgent) {
throw new CustomError('aiAgent not found');
}
const prompt = await Prompt.findByPk(promptId);
ctx.body = {
chatPrompt: chatPrompt,
aiAgent,
prompt,
};
})
.addTo(app);

View File

@@ -79,6 +79,7 @@ ContainerModel.init(
{
sequelize,
tableName: 'kv_container',
paranoid: true,
},
);

View File

@@ -9,3 +9,7 @@ import './prompt-graph/index.ts';
import './agent/index.ts';
import './user/index.ts';
import './chat-prompt/index.ts';
import './chat-history/index.ts';

View File

@@ -68,6 +68,7 @@ PageModel.init(
{
sequelize,
tableName: 'kv_page',
paranoid: true,
},
);

View File

@@ -82,6 +82,7 @@ ResourceModel.init(
{
sequelize,
tableName: 'kv_resource',
paranoid: true,
},
);

View File

@@ -36,7 +36,7 @@ app
const { checkToken: token } = ctx.query;
try {
const result = await User.verifyToken(token);
ctx.body = result?.payload || {};
ctx.body = result || {};
} catch (e) {
new CustomError(401, 'Token InValid ');
}