AIRegulation-DocAnalysis/frontend/src/api/rag.ts

import type { QuickQuestionsResponse, SSEMessage } from './index';

const AGENT_API_BASE = '/api/v1';
const TOKEN_KEY = 'auth_token';
function getToken(): string | null { return localStorage.getItem(TOKEN_KEY); }

const _FALLBACK_QUESTIONS = [
  { id: '1', question: '请总结最新入库法规对电池安全的核心要求', category: '法规解读' },
  { id: '2', question: '我上传的制度文档与新能源法规有哪些潜在冲突？', category: '差距分析' },
  { id: '3', question: '请给出法规依据，并按条款列出整改建议', category: '整改建议' },
  { id: '4', question: '请解释 UN-ECE 与 GB 标准在网络安全方面的差异', category: '标准对比' },
];

export async function getQuickQuestions(): Promise<QuickQuestionsResponse> {
  try {
    const response = await fetch(`${AGENT_API_BASE}/rag/quick-questions`);
    if (!response.ok) throw new Error(`status ${response.status}`);
    return response.json() as Promise<QuickQuestionsResponse>;
  } catch {
    return { questions: _FALLBACK_QUESTIONS };
  }
}

function parseSSEChunk(raw: string, onMessage: (data: SSEMessage) => void) {
  const blocks = raw.split('\n\n');
  for (const block of blocks) {
    if (!block.trim()) continue;

    let eventName = 'message';
    const dataLines: string[] = [];

    for (const line of block.split('\n')) {
      if (line.startsWith('event:')) {
        eventName = line.slice(6).trim();
      } else if (line.startsWith('data:')) {
        dataLines.push(line.slice(5).trim());
      }
    }

    const joined = dataLines.join('\n');
    if (!joined) continue;

    // /agent/chat/stream uses named events (sources, content, done, error, session, status)
    // /rag/chat wraps everything in event:message with type in JSON body — handle both
    if (eventName === 'session') {
      try {
        const payload = JSON.parse(joined) as Record<string, unknown>;
        onMessage({ type: 'session', session_id: String(payload.session_id ?? '') });
      } catch { /* ignore */ }
    } else if (eventName === 'sources') {
      try {
        const docs = JSON.parse(joined) as Array<Record<string, unknown>>;
        onMessage({
          type: 'retrieved',
          docs: docs.map((doc, index) => ({
            id: String(doc.doc_id || doc.index || index + 1),
            score: Number(doc.score || 0),
            preview: String(doc.content || doc.snippet || ''),
            doc_name: String(doc.doc_name || doc.filename || `引用 ${index + 1}`),
            clause: String(doc.clause_number || doc.section_title || '法规片段'),
            doc_id: doc.doc_id ? String(doc.doc_id) : undefined,
            download_url: doc.doc_id ? `${AGENT_API_BASE}/documents/download/${String(doc.doc_id)}` : undefined,
          })),
        });
      } catch { /* ignore */ }
    } else if (eventName === 'content') {
      onMessage({ type: 'chunk', text: joined });
    } else if (eventName === 'done') {
      try {
        const payload = JSON.parse(joined) as Record<string, unknown>;
        onMessage({ type: 'done', session_id: payload.session_id ? String(payload.session_id) : undefined });
      } catch {
        onMessage({ type: 'done' });
      }
    } else if (eventName === 'error') {
      onMessage({ type: 'error', text: joined });
    } else if (eventName === 'status') {
      onMessage({ type: 'status', text: joined });
    } else if (eventName === 'message') {
      // /rag/chat format: event:message + JSON body with type field
      try {
        const payload = JSON.parse(joined) as SSEMessage;
        onMessage(payload);
      } catch { /* ignore */ }
    }
  }
}

export async function ragChat(
  query: string,
  topK: number = 5,
  onMessage: (data: SSEMessage) => void,
  onError?: (error: Error) => void,
  onComplete?: () => void,
  filters?: string,
  sessionId?: string,
  signal?: AbortSignal,
): Promise<void> {
  try {
    const response = await fetch(`${AGENT_API_BASE}/agent/chat/stream`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Accept: 'text/event-stream',
        ...(getToken() ? { Authorization: `Bearer ${getToken()}` } : {}),
      },
      body: JSON.stringify({
        query,
        top_k: topK,
        ...(filters ? { filters } : {}),
        ...(sessionId ? { session_id: sessionId } : {}),
      }),
      signal,
    });

    if (!response.ok || !response.body) {
      throw new Error(`HTTP error! status: ${response.status}`);
    }

    const reader = response.body.getReader();
    const decoder = new TextDecoder();
    let buffer = '';

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

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

    if (buffer.trim()) {
      parseSSEChunk(buffer, onMessage);
    }

    if (onComplete) {
      onComplete();
    }
  } catch (error) {
    if (error instanceof DOMException && error.name === 'AbortError') return;
    if (onError) {
      onError(error instanceof Error ? error : new Error(String(error)));
    }
  }
}

export type { QuickQuestionsResponse, SSEMessage };