KnowledgeBase/web/services/chatService.ts

export interface ChatMessage {
  role: 'user' | 'assistant';
  content: string;
}

export interface ChatSource {
  fileName: string;
  content: string;
  score: number;
  chunkIndex: number;
  pageNumber?: number; // 追加
  fileId?: string;
}

export class ChatService {
  async *streamChat(
    message: string,
    history: ChatMessage[],
    authToken: string,
    userLanguage: string = 'ja',
    selectedEmbeddingId?: string,
    selectedLLMId?: string, // 新增：选中的 LLM ID
    selectedGroups?: string[], // 新增：选中的分组
    selectedFiles?: string[], // 新增：选中的文件
    historyId?: string, // 新增：对话历史ID
    enableRerank?: boolean, // 新增：启用 Rerank
    selectedRerankId?: string, // 新增：Rerank 模型 ID
    temperature?: number, // 新增：temperature 参数
    maxTokens?: number, // 新增：maxTokens 参数
    topK?: number, // 新增：topK 参数
    similarityThreshold?: number, // 新增：similarityThreshold 参数
    enableQueryExpansion?: boolean, // 新增
    enableHyDE?: boolean, // 新增
    scoreThreshold?: number // 新增
  ): AsyncGenerator<{ type: 'content' | 'sources' | 'error'; data: any }> {
    try {
      const response = await fetch('/api/chat/stream', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${authToken}`,
        },
        body: JSON.stringify({
          message,
          history,
          userLanguage,
          selectedEmbeddingId,
          selectedLLMId, // Pass LLM ID
          selectedGroups, // 传递分组过滤参数
          selectedFiles, // 传递文件过滤参数
          historyId, // 传递历史记录ID
          enableRerank,
          selectedRerankId,
          temperature, // 传递 temperature 参数
          maxTokens, // 传递 maxTokens 参数
          topK, // 传递 topK 参数
          similarityThreshold, // 传递 similarityThreshold 参数
          enableQueryExpansion,
          enableHyDE,
          scoreThreshold
        }),
      });

      if (!response.ok) {
        let errorMessage = '请求失败';
        try {
          const error = await response.json();
          errorMessage = error.error || error.message || '请求失败';
        } catch {
          errorMessage = `服务器错误: ${response.status}`;
        }
        yield { type: 'error', data: errorMessage };
        return;
      }

      const reader = response.body?.getReader();
      if (!reader) {
        yield { type: 'error', data: '无法读取响应流' };
        return;
      }

      const decoder = new TextDecoder();
      let buffer = '';

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = line.slice(6);
            if (data === '[DONE]') {
              return;
            }

            try {
              const parsed = JSON.parse(data);
              yield parsed;
            } catch (e) {
              console.warn('Failed to parse SSE data:', data);
            }
          }
        }
      }
    } catch (error) {
      yield { type: 'error', data: error.message || '网络错误' };
    }
  }

  async *streamAssist(
    instruction: string,
    context: string,
    authToken: string
  ): AsyncGenerator<{ type: 'content' | 'error'; data: any }> {
    try {
      const response = await fetch('/api/chat/assist', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${authToken}`,
        },
        body: JSON.stringify({ instruction, context }),
      });

      if (!response.ok) {
        yield { type: 'error', data: '请求失败' };
        return;
      }

      const reader = response.body?.getReader();
      if (!reader) return;

      const decoder = new TextDecoder();
      let buffer = '';

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = line.slice(6);
            if (data === '[DONE]') return;
            try {
              yield JSON.parse(data);
            } catch (e) { console.warn(e) }
          }
        }
      }
    } catch (error: any) {
      yield { type: 'error', data: error.message };
    }
  }
}

export const chatService = new ChatService();