feat: add chat message

src/routes/chat-history/chat-io.ts

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

src/routes/chat-history/index.ts

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

src/routes/chat-history/list.ts

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

src/routes/chat-history/session-list.ts

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