aurak/web/services/chatService.ts

import { apiClient } from './apiClient';

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

export interface ChatSource {
  fileName: string;
  title?: string;
  content: string;
  score: number;
  chunkIndex: number;
  fileId?: string;
}

export class ChatService {
  async *streamChat(
    message: string,
    history: ChatMessage[],
    authToken: string,
    userLanguage: string = 'zh',
    selectedEmbeddingId?: string,
    selectedLLMId?: string, // Added: Selected LLM ID
    selectedGroups?: string[], // Added: Selected groups
    selectedFiles?: string[], // Added: Selected files
    historyId?: string, // Added: Conversation history ID
    enableRerank?: boolean, // Added: Enable Rerank
    selectedRerankId?: string, // Added: Rerank model ID
    temperature?: number, // Added: temperature parameter
    maxTokens?: number, // Added: maxTokens parameter
    topK?: number, // Added: topK parameter
    similarityThreshold?: number, // Added: similarityThreshold parameter
    rerankSimilarityThreshold?: number, // Added: rerankSimilarityThreshold parameter
    enableQueryExpansion?: boolean, // Added
    enableHyDE?: boolean // Added
  ): AsyncGenerator<{ type: 'content' | 'sources' | 'error' | 'historyId'; data: any }> {
    try {
      const response = await apiClient.request('/chat/stream', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'x-user-language': userLanguage || localStorage.getItem('userLanguage') || 'zh',
        },
        body: JSON.stringify({
          message,
          history,
          userLanguage,
          selectedEmbeddingId,
          selectedLLMId,
          selectedGroups,
          selectedFiles,
          historyId,
          enableRerank,
          selectedRerankId,
          temperature,
          maxTokens,
          topK,
          similarityThreshold,
          rerankSimilarityThreshold,
          enableQueryExpansion,
          enableHyDE
        }),
      });

      if (!response.ok) {
        let errorMessage = 'Request failed';
        try {
          const error = await response.json();
          errorMessage = error.error || error.message || 'Request failed';
        } catch {
          errorMessage = `Server error: ${response.status}`;
        }
        yield { type: 'error', data: errorMessage };
        return;
      }

      const reader = response.body?.getReader();
      if (!reader) {
        yield { type: 'error', data: 'Cannot read response stream' };
        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: any) {
      yield { type: 'error', data: error.message || 'Network error' };
    }
  }

  async *streamAssist(
    instruction: string,
    context: string,
    authToken: string
  ): AsyncGenerator<{ type: 'content' | 'error'; data: any }> {
    try {
      const response = await apiClient.request('/chat/assist', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'x-user-language': localStorage.getItem('userLanguage') || 'zh',
        },
        body: JSON.stringify({ instruction, context }),
      });

      if (!response.ok) {
        yield { type: 'error', data: 'Request failed' };
        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();