gongxue-base/apps/admin/src/components/AiChat/provider.ts

import {
  AbstractChatProvider,
  XRequest,
  type TransformMessage,
  type XRequestOptions,
} from '@ant-design/x-sdk';
import type {
  AiAttachment,
  AiChatInput,
  AiChatMessage,
  AiSseChunk,
  AiToolRun,
} from './types';

interface AiSsePayload {
  messageId?: number;
  userMessageId?: number;
  assistantMessageId?: number;
  delta?: string;
  content?: string;
  reasoningContent?: string | null;
  toolCallId?: string;
  toolName?: string;
  skillKey?: string | null;
  status?: string;
  summary?: string | null;
  durationMs?: number | null;
  attachment?: AiAttachment;
  message?:
    | string
    | {
        id?: number;
        content?: string;
        reasoningContent?: string | null;
        status?: string;
        toolRuns?: AiToolRun[];
        attachments?: AiAttachment[];
        replyToMessageId?: number | null;
        feedback?: 'like' | 'dislike' | null;
        feedbackReason?: string | null;
        metadata?: Record<string, unknown> | null;
      };
  error?: string;
}

function emptyAssistant(): AiChatMessage {
  return {
    role: 'assistant',
    content: '',
    reasoningContent: '',
    toolRuns: [],
    attachments: [],
  };
}

export function parseSsePayload(chunk?: AiSseChunk): {
  event: string;
  payload: AiSsePayload;
} {
  if (!chunk) return { event: '', payload: {} };
  const event = chunk.event?.trim() || 'message';
  if (!chunk.data || chunk.data === '[DONE]') return { event, payload: {} };
  try {
    const parsed: unknown = JSON.parse(chunk.data);
    return {
      event,
      payload: parsed && typeof parsed === 'object' ? (parsed as AiSsePayload) : {},
    };
  } catch {
    return { event, payload: { delta: chunk.data } };
  }
}

function upsertToolRun(
  toolRuns: AiToolRun[],
  payload: AiSsePayload,
  fallbackStatus: AiToolRun['status'],
): AiToolRun[] {
  const toolCallId = payload.toolCallId || `${payload.toolName || 'tool'}-${toolRuns.length}`;
  const next: AiToolRun = {
    toolCallId,
    toolName: payload.toolName || '查询工具',
    skillKey: payload.skillKey,
    status: (payload.status as AiToolRun['status']) || fallbackStatus,
    summary: payload.summary,
    resultSummary: fallbackStatus === 'running' ? undefined : payload.summary,
    argumentsSummary: fallbackStatus === 'running' ? payload.summary : undefined,
    durationMs: payload.durationMs,
  };
  const index = toolRuns.findIndex((item) => item.toolCallId === toolCallId);
  if (index === -1) return [...toolRuns, next];
  return toolRuns.map((item, itemIndex) => (itemIndex === index ? { ...item, ...next } : item));
}

function normalizeToolRuns(toolRuns: AiToolRun[] | undefined, fallback: AiToolRun[]): AiToolRun[] {
  if (!toolRuns) return fallback;
  return toolRuns.map((tool) => ({
    ...tool,
    status: tool.status === 'error' ? 'failed' : tool.status,
    summary: tool.resultSummary ?? tool.argumentsSummary ?? tool.summary,
  }));
}

export function reduceAiSseMessage(
  originMessage: AiChatMessage | undefined,
  chunk?: AiSseChunk,
): AiChatMessage {
  const message = originMessage ? { ...originMessage } : emptyAssistant();
  const { event, payload } = parseSsePayload(chunk);

  if (event === 'message.created') {
    const nested = typeof payload.message === 'object' ? payload.message : undefined;
    message.id = nested?.id ?? payload.assistantMessageId ?? payload.messageId ?? message.id;
    message.content = nested?.content ?? message.content;
    message.reasoningContent = nested?.reasoningContent ?? message.reasoningContent;
    message.toolRuns = normalizeToolRuns(nested?.toolRuns, message.toolRuns);
    message.attachments = nested?.attachments ?? message.attachments;
    message.replyToMessageId = nested?.replyToMessageId ?? message.replyToMessageId;
    message.feedback = nested?.feedback ?? message.feedback;
    message.feedbackReason = nested?.feedbackReason ?? message.feedbackReason;
    message.metadata = nested?.metadata ?? message.metadata;
  } else if (event === 'reasoning.delta') {
    message.reasoningContent += payload.delta ?? payload.reasoningContent ?? '';
  } else if (event === 'content.delta') {
    message.content += payload.delta ?? payload.content ?? '';
  } else if (event === 'tool.started') {
    message.toolRuns = upsertToolRun(message.toolRuns, payload, 'running');
  } else if (event === 'tool.completed') {
    message.toolRuns = upsertToolRun(message.toolRuns, payload, 'success');
  } else if (event === 'tool.failed') {
    message.toolRuns = upsertToolRun(message.toolRuns, payload, 'failed');
  } else if (event === 'attachment.processed' && payload.attachment) {
    if (!message.attachments.some((item) => item.id === payload.attachment?.id)) {
      message.attachments = [...message.attachments, payload.attachment];
    }
  } else if (event === 'message.completed') {
    const nested = typeof payload.message === 'object' ? payload.message : undefined;
    message.id = nested?.id ?? payload.messageId ?? message.id;
    message.content = nested?.content ?? payload.content ?? message.content;
    message.reasoningContent =
      nested?.reasoningContent ?? payload.reasoningContent ?? message.reasoningContent;
    message.toolRuns = normalizeToolRuns(nested?.toolRuns, message.toolRuns);
    message.attachments = nested?.attachments ?? message.attachments;
    message.replyToMessageId = nested?.replyToMessageId ?? message.replyToMessageId;
    message.feedback = nested?.feedback ?? message.feedback;
    message.feedbackReason = nested?.feedbackReason ?? message.feedbackReason;
    message.metadata = nested?.metadata ?? message.metadata;
  } else if (event === 'message.cancelled') {
    message.id = payload.messageId ?? message.id;
    message.cancelled = true;
  } else if (event === 'error') {
    message.error =
      (typeof payload.message === 'string' ? payload.message : undefined) ||
      payload.error ||
      'AI 回答生成失败';
  }
  return message;
}

async function authenticatedFetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
  const headers = new Headers(init?.headers);
  const token = localStorage.getItem('token');
  if (token) headers.set('Authorization', `Bearer ${token}`);
  headers.set('Accept', 'text/event-stream');
  let requestInput = input;
  let requestInit = init;
  if (typeof init?.body === 'string') {
    try {
      const body = JSON.parse(init.body) as AiChatInput;
      if (body.regenerateMessageId) {
        requestInput = `${String(input).replace(/\/stream$/, '')}/messages/${body.regenerateMessageId}/regenerate/stream`;
        requestInit = {
          ...init,
          body: JSON.stringify({ clientRequestId: body.clientRequestId }),
        };
      } else {
        const {
          localAttachments: _localAttachments,
          reloadMessage: _reloadMessage,
          regenerateMessageId: _regenerateMessageId,
          ...payload
        } = body;
        requestInit = { ...init, body: JSON.stringify(payload) };
      }
    } catch {
      requestInit = init;
    }
  }
  const response = await fetch(requestInput, { ...requestInit, headers });
  if (response.status === 401) {
    localStorage.removeItem('token');
    localStorage.removeItem('user');
    localStorage.removeItem('permissions');
    window.location.href = '/login';
  }
  return response;
}

export class GongxueAiChatProvider extends AbstractChatProvider<
  AiChatMessage,
  AiChatInput,
  AiSseChunk
> {
  constructor(url: string, onSettled?: () => void) {
    super({
      request: XRequest<AiChatInput, AiSseChunk, AiChatMessage>(url, {
        manual: true,
        fetch: authenticatedFetch,
        timeout: 15_000,
        streamTimeout: 120_000,
        callbacks: {
          onUpdate: () => undefined,
          onSuccess: () => onSettled?.(),
          onError: () => onSettled?.(),
        },
      }),
    });
  }

  transformParams(
    requestParams: Partial<AiChatInput>,
    options: XRequestOptions<AiChatInput, AiSseChunk, AiChatMessage>,
  ): AiChatInput {
    return {
      ...options.params,
      message: requestParams.message?.trim() || '',
      attachmentIds: requestParams.attachmentIds ?? [],
      skillKey: requestParams.skillKey ?? null,
      clientRequestId: requestParams.clientRequestId || crypto.randomUUID(),
      localAttachments: requestParams.localAttachments,
      regenerateMessageId: requestParams.regenerateMessageId,
      reloadMessage: requestParams.reloadMessage,
    };
  }

  transformLocalMessage(requestParams: Partial<AiChatInput>): AiChatMessage {
    return {
      role: 'user',
      content: requestParams.message?.trim() || '',
      reasoningContent: '',
      toolRuns: [],
      attachments: requestParams.localAttachments ?? [],
    };
  }

  transformMessage(info: TransformMessage<AiChatMessage, AiSseChunk>): AiChatMessage {
    return reduceAiSseMessage(info.originMessage, info.chunk);
  }
}