fix

2026-06-04 15:43:44 +08:00
parent ac490d851a
commit 746513cc54
11 changed files with 955 additions and 131 deletions
@@ -1,4 +1,4 @@
-import { useState, useRef, useEffect } from 'react';
+import { useState, useRef, useEffect, useCallback } from 'react';
 import { Topbar } from '../../components/layout/Topbar';
 import { Send, Download } from 'lucide-react';

@@ -6,79 +6,138 @@ interface Message {
  id: string;
  role: 'user' | 'assistant';
  text: string;
+  // citation indices mentioned in this assistant message (1-based, matching citations array)
+  citationRefs?: number[];
 }

 interface Citation {
-  score: number;
-  name: string;
-  clause: string;
-  snippet: string;
+  index: number;   // 1-based, matches [N] markers in text
+  score: number;   // 0–100 display percentage
+  name: string;    // doc_name
+  clause: string;  // section_title or clause
+  snippet: string; // preview text
+  docId?: string;
 }

-const HISTORY = [
-  { id: 'h1', title: 'EU AI Act Article 9 scope', date: '2025-11-18' },
-  { id: 'h2', title: 'MIIT training data requirements', date: '2025-11-15' },
-  { id: 'h3', title: 'ISO 21434 CSMS audit scope', date: '2025-11-10' },
-];
+// Map a raw source doc from the backend "retrieved" event to our Citation shape.
+// Backend fields: { id, score(0-1), preview, doc_name, clause, doc_id }
+function mapSource(s: Record<string, unknown>, idx: number): Citation {
+  const rawScore = typeof s.score === 'number' ? s.score : 0;
+  const displayScore = rawScore <= 1 ? Math.round(rawScore * 100) : Math.round(rawScore);
+  return {
+    index: idx,
+    score: displayScore,
+    name: String(s.doc_name ?? ''),
+    clause: String(s.clause ?? s.section_title ?? ''),
+    snippet: String(s.preview ?? s.text ?? ''),
+    docId: s.doc_id ? String(s.doc_id) : undefined,
+  };
+}

-const QUICK = [
+// Parse message text and replace [N] with clickable <button class="cite-ref"> elements.
+function renderWithCitations(
+  text: string,
+  onCiteClick: (n: number) => void,
+): React.ReactNode[] {
+  const parts = text.split(/(\[\d+\])/g);
+  return parts.map((part, i) => {
+    const m = part.match(/^\[(\d+)\]$/);
+    if (m) {
+      const n = parseInt(m[1], 10);
+      return (
+        <button
+          key={i}
+          className="cite-ref"
+          onClick={() => onCiteClick(n)}
+          title={`Jump to source [${n}]`}
+        >
+          {n}
+        </button>
+      );
+    }
+    return part;
+  });
+}
+
+const MOCK_QUICK = [
  'What does EU AI Act Art. 9 require for risk management?',
  'Which documents need CSMS certification?',
  'Summarize MIIT training data rules',
  'What are high-risk AI categories under Annex III?',
 ];

-const MOCK_CITATIONS: Citation[] = [
-  {
-    score: 94, name: 'EU AI Act', clause: 'Art. 9(1)',
-    snippet: 'Providers of high-risk AI systems shall establish a risk management system consisting of a continuous iterative process run throughout the entire lifecycle.'
-  },
-  {
-    score: 87, name: 'Vehicle AI Safety Manual', clause: '§4.2.1',
-    snippet: 'All AI systems classified as high-risk must maintain a documented risk register with quarterly review cadence.'
-  },
-  {
-    score: 72, name: 'ISO/SAE 21434', clause: 'Clause 9.3',
-    snippet: 'The cybersecurity management system shall include AI model update governance procedures and audit log retention policy.'
-  },
-];
-
 export function RagChatPage() {
  const [messages, setMessages] = useState<Message[]>([
    {
      id: 'init', role: 'assistant',
-      text: 'Hello! I can answer questions about your indexed regulations and compliance documents. Try asking about EU AI Act requirements, MIIT rules, or ISO/SAE 21434 scope.'
+      text: 'Hello! I can answer questions about your indexed regulations and compliance documents. Try asking about EU AI Act requirements, MIIT rules, or ISO/SAE 21434 scope.',
    }
  ]);
+  const [quickPrompts, setQuickPrompts] = useState<string[]>(MOCK_QUICK);
  const [input, setInput] = useState('');
  const [streaming, setStreaming] = useState(false);
-  const [citations, setCitations] = useState<Citation[]>(MOCK_CITATIONS);
+  const [citations, setCitations] = useState<Citation[]>([]);
+  const [highlightedCit, setHighlightedCit] = useState<number | null>(null);
+  const [sessionId, setSessionId] = useState<string | null>(null);
+
  const bottomRef = useRef<HTMLDivElement>(null);
+  const citRailRef = useRef<HTMLDivElement>(null);
+  const citItemRefs = useRef<Record<number, HTMLDivElement | null>>({});
  const abortRef = useRef<AbortController | null>(null);

+  // Fetch quick questions from backend on mount
+  useEffect(() => {
+    fetch('/api/v1/rag/quick-questions')
+      .then(r => r.json())
+      .then(d => {
+        if (Array.isArray(d?.questions) && d.questions.length > 0) {
+          setQuickPrompts(d.questions.map((q: { question: string }) => q.question));
+        }
+      })
+      .catch(() => { /* keep mock */ });
+  }, []);
+
+  // Auto-scroll to latest message
  useEffect(() => {
    bottomRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);

+  // Jump to citation N in the rail and highlight it
+  const jumpToCitation = useCallback((n: number) => {
+    setHighlightedCit(n);
+    const el = citItemRefs.current[n];
+    if (el) {
+      el.scrollIntoView({ behavior: 'smooth', block: 'nearest' });
+    }
+    // Clear highlight after 3s
+    setTimeout(() => setHighlightedCit(h => h === n ? null : h), 3000);
+  }, []);
+
  async function send(text?: string) {
    const q = (text ?? input).trim();
    if (!q || streaming) return;
    setInput('');
+
    const userMsg: Message = { id: Date.now().toString(), role: 'user', text: q };
    setMessages(m => [...m, userMsg]);

    const assistantId = (Date.now() + 1).toString();
    setMessages(m => [...m, { id: assistantId, role: 'assistant', text: '' }]);
    setStreaming(true);
+    setCitations([]);
+    setHighlightedCit(null);

    const ctrl = new AbortController();
    abortRef.current = ctrl;

    try {
+      const body: Record<string, unknown> = { query: q, top_k: 5 };
+      if (sessionId) body.session_id = sessionId;
+
      const res = await fetch('/api/v1/rag/chat', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({ question: q }),
+        body: JSON.stringify(body),
        signal: ctrl.signal,
      });

@@ -86,29 +145,65 @@ export function RagChatPage() {
      const reader = res.body.getReader();
      const dec = new TextDecoder();
      let buffer = '';
+      const newCitations: Citation[] = [];

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
-        buffer += dec.decode(value);
-        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]') break;
-            try {
-              const j = JSON.parse(data);
-              if (j.text) setMessages(m => m.map(msg =>
-                msg.id === assistantId ? { ...msg, text: msg.text + j.text } : msg
-              ));
-              if (j.citations) setCitations(j.citations);
-            } catch {
+        buffer += dec.decode(value, { stream: true });
+
+        // SSE blocks separated by double newline
+        const blocks = buffer.split('\n\n');
+        buffer = blocks.pop() ?? '';
+
+        for (const block of blocks) {
+          const dataLine = block.split('\n').find(l => l.startsWith('data: '));
+          if (!dataLine) continue;
+          const raw = dataLine.slice(6).trim();
+          if (!raw) continue;
+          try {
+            const j = JSON.parse(raw);
+
+            if (j.type === 'session') {
+              // Backend assigned a session_id — persist for next request
+              if (j.session_id) setSessionId(j.session_id);
+
+            } else if (j.type === 'retrieved' && Array.isArray(j.docs)) {
+              // Sources arrive before the answer starts
+              const mapped = j.docs.map((d: Record<string, unknown>, i: number) => mapSource(d, i + 1));
+              newCitations.push(...mapped);
+              setCitations([...mapped]);
+
+            } else if (j.type === 'chunk' && j.text) {
              setMessages(m => m.map(msg =>
-                msg.id === assistantId ? { ...msg, text: msg.text + data } : msg
+                msg.id === assistantId
+                  ? { ...msg, text: msg.text + (j.text as string) }
+                  : msg
+              ));
+
+            } else if (j.type === 'status') {
+              // Status message (e.g. "找到N条相关法规…") — could show in UI if desired
+              // For now we ignore it to keep the bubble clean
+
+            } else if (j.type === 'done') {
+              // Extract which citation numbers appear in the final answer
+              setMessages(m => m.map(msg => {
+                if (msg.id !== assistantId) return msg;
+                const refs = [...new Set(
+                  [...msg.text.matchAll(/\[(\d+)\]/g)].map(r => parseInt(r[1], 10))
+                )].filter(n => n >= 1 && n <= newCitations.length);
+                return { ...msg, citationRefs: refs };
+              }));
+              break;
+
+            } else if (j.type === 'error') {
+              setMessages(m => m.map(msg =>
+                msg.id === assistantId
+                  ? { ...msg, text: `Error: ${j.text ?? 'Unknown error'}` }
+                  : msg
              ));
            }
-          }
+          } catch { /* malformed JSON chunk, skip */ }
        }
      }
    } catch (e: unknown) {
@@ -130,30 +225,44 @@ export function RagChatPage() {
    <div className="chat-page">
      <Topbar
        title="Regulation Q&A"
-        actions={<button className="btn sm"><Download size={13} />Export chat</button>}
+        actions={
+          <button
+            className="btn sm"
+            onClick={() => {
+              const text = messages.map(m => `${m.role === 'user' ? 'Q' : 'A'}: ${m.text}`).join('\n\n');
+              const blob = new Blob([text], { type: 'text/plain' });
+              const url = URL.createObjectURL(blob);
+              const a = document.createElement('a'); a.href = url; a.download = 'chat-export.txt'; a.click();
+              URL.revokeObjectURL(url);
+            }}
+          >
+            <Download size={13} />Export chat
+          </button>
+        }
      />
+
      <div className="chat-body">
+        {/* ── History pane ── */}
        <div className="history-pane">
-          <div className="history-header">Chat history</div>
-          {HISTORY.map(h => (
-            <div key={h.id} className="history-item">
-              <div className="history-title">{h.title}</div>
-              <div className="history-date">{h.date}</div>
-            </div>
-          ))}
-          <div className="quick-header">Quick prompts</div>
-          {QUICK.map(q => (
-            <button key={q} className="quick-item" onClick={() => send(q)}>{q}</button>
+          <div className="history-header">Quick prompts</div>
+          {quickPrompts.map(q => (
+            <button key={q} className="quick-item" onClick={() => send(q)}>
+              {q}
+            </button>
          ))}
        </div>

+        {/* ── Chat main ── */}
        <div className="chat-main">
          <div className="messages">
            {messages.map(msg => (
              <div key={msg.id} className={`message msg-${msg.role}`}>
                {msg.role === 'assistant' && <div className="msg-avatar">AI</div>}
                <div className="msg-bubble">
-                  {msg.text}
+                  {msg.role === 'assistant'
+                    ? renderWithCitations(msg.text, jumpToCitation)
+                    : msg.text
+                  }
                  {streaming && msg.id === lastAssistantId && (
                    <span className="blink-cursor">▋</span>
                  )}
@@ -166,7 +275,7 @@ export function RagChatPage() {

          <div className="composer">
            <div className="quick-chips">
-              {QUICK.slice(0, 3).map(q => (
+              {quickPrompts.slice(0, 3).map(q => (
                <button key={q} className="chip" onClick={() => send(q)}>
                  {q.length > 42 ? q.slice(0, 42) + '…' : q}
                </button>
@@ -175,7 +284,7 @@ export function RagChatPage() {
            <div className="composer-row">
              <textarea
                className="composer-input"
-                placeholder="Ask about your regulations..."
+                placeholder="Ask about your regulations…"
                value={input}
                onChange={e => setInput(e.target.value)}
                onKeyDown={e => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send(); } }}
@@ -192,13 +301,29 @@ export function RagChatPage() {
          </div>
        </div>

-        <div className="citation-rail">
-          <div className="citation-header">Sources</div>
+        {/* ── Citation rail ── */}
+        <div className="citation-rail" ref={citRailRef}>
+          <div className="citation-header">
+            Sources {citations.length > 0 && `(${citations.length})`}
+          </div>
+          {citations.length === 0 && (
+            <p style={{ padding: '12px 16px', fontSize: 12, color: 'var(--muted)', lineHeight: 1.5 }}>
+              Citations will appear here after a response is generated.
+            </p>
+          )}
          {citations.map(c => (
-            <div key={`${c.name}-${c.clause}`} className="citation-item">
-              <div className="cit-score">{c.score}%</div>
-              <div>
-                <div className="cit-name">{c.name} <span className="cit-clause">{c.clause}</span></div>
+            <div
+              key={c.index}
+              ref={el => { citItemRefs.current[c.index] = el; }}
+              className={`citation-item${highlightedCit === c.index ? ' highlighted' : ''}`}
+            >
+              <div className="cit-index">[{c.index}]</div>
+              <div style={{ flex: 1, minWidth: 0 }}>
+                <div style={{ display: 'flex', alignItems: 'baseline', gap: 4, marginBottom: 3 }}>
+                  <div className="cit-name">{c.name}</div>
+                  {c.clause && <span className="cit-clause">{c.clause}</span>}
+                  <span className="cit-score" style={{ marginLeft: 'auto' }}>{c.score}%</span>
+                </div>
                <div className="cit-snippet">{c.snippet}</div>
              </div>
            </div>