marsview-enterprise-node/controller/ai.controller.js

const aiService = require('../service/ai/ai.service.js');
const util = require('../utils/util');
const { model: openaiModel } = require('../service/ai/openai.service');
const { findRelevantContent } = require('../service/ai/embedding.service');
const { getSystemPrompt } = require('../service/ai/prompt.service');
const { streamText } = require('ai');

function formatMessages(messages) {
  return messages.map((message) => ({
    ...message,
    content: Array.isArray(message.content)
      ? message.content.map((content) => ({
          ...content,
          ...(content.type === 'image_url' && {
            type: 'image',
            image: content.image_url.url.replace(/^data:image\/\w+;base64,/, ''),
          }),
        }))
      : message.content,
  }));
}

module.exports = {
  async codeGenerate(ctx) {
    const { message } = ctx.request.body;
    if (!message) {
      return util.fail(ctx, '请输入提示词');
    }
    try {
      const result = await aiService.ApiZhipuAi(message);
      util.success(ctx, {
        jsx: result[0],
        config: result[1],
      });
    } catch (error) {
      util.fail(ctx, error);
    }
  },
  async chatStream(ctx) {
    try {
      const { messages } = ctx.request.body;
      const lastMessage = messages[messages.length - 1];

      // 处理最后一条消息
      const lastMessageContent = Array.isArray(lastMessage.content)
        ? lastMessage.content
            .filter((c) => c.type === 'text')
            .map((c) => c.text)
            .join('')
        : lastMessage.content;

      // 获取相关上下文
      const relevantContent = await findRelevantContent(lastMessageContent);
      const systemPrompt = getSystemPrompt(relevantContent.map((c) => c.content).join('\n'));

      const { textStream } = streamText({
        model: openaiModel,
        system: systemPrompt,
        messages: formatMessages(messages),
      });

      ctx.set('Access-Control-Allow-Origin', '*');
      ctx.set('Access-Control-Allow-Methods', 'POST, GET, OPTIONS');
      ctx.set('Access-Control-Allow-Headers', 'Content-Type, Authorization, X-Requested-With');

      ctx.status = 200;

      for await (const textPart of textStream) {
        // 返回消息
        ctx.res.write(textPart);
      }
      ctx.res.end();
      util.success(ctx, '对话结束');
    } catch (error) {
      ctx.status = 500;
      console.error('Chat stream error:', error);
      ctx.body = { error: error.message || 'Internal Server Error' };
    }
  },
};