worldquant-alpha-system/frontend/src/ai/ChatPanel.tsx

import { useChat } from "@ai-sdk/react";
import {
  lazy,
  Suspense,
  useCallback,
  useEffect,
  useMemo,
  useRef,
  useState,
} from "react";
import {
  Banner,
  Button,
  Select,
  Spin,
  Tag,
  TextArea,
} from "@douyinfe/semi-ui-19";
import { api, post } from "../api";
import type { Job } from "../types";
import { BusinessCard } from "./ToolCard";
import { contextLabel } from "./workspace";
import { chatTransport } from "./transport";
import { runLabels } from "./types";
import type {
  ChatMessage,
  Conversation,
  ConversationDetail,
  ModelSettings,
  PageContext,
  RunSnapshot,
  Resource,
  UIAction,
} from "./types";

const MessageMarkdown = lazy(() =>
  import("./MessageMarkdown").then((module) => ({
    default: module.MessageMarkdown,
  })),
);

export function ChatPanel({
  open,
  context,
  onClose,
  onAction,
  onChanged,
  onSettings,
  width,
  onWidth,
  timezone,
  jobs,
  action,
}: {
  open: boolean;
  context: PageContext;
  onClose: () => void;
  onAction: (action: UIAction) => void;
  onChanged: (resources: Resource[]) => void;
  onSettings: () => void;
  width: number;
  onWidth: (width: number) => void;
  timezone?: string;
  jobs: Job[];
  action: UIAction | null;
}) {
  const [settings, setSettings] = useState<ModelSettings | null>(null);
  const [conversations, setConversations] = useState<Conversation[]>([]);
  const [conversationId, setConversationId] = useState("");
  const [runs, setRuns] = useState<RunSnapshot[]>([]);
  const [text, setText] = useState("");
  const [failure, setFailure] = useState("");
  const [loading, setLoading] = useState(false);
  const root = useRef<HTMLElement>(null);
  const input = useRef<HTMLDivElement>(null);
  const bottom = useRef<HTMLDivElement>(null);
  const previousFocus = useRef<HTMLElement | null>(null);
  const seenChanges = useRef(new Set<string>());
  const current = useRef({ conversationId, context });
  current.current = { conversationId, context };
  const refreshRef = useRef<() => Promise<void>>(async () => {});
  const changedRef = useRef(onChanged);
  changedRef.current = onChanged;
  const observeChanges = useCallback((items: RunSnapshot[]) => {
    const changed = new Set<Resource>();
    for (const run of items)
      for (const call of run.tools ?? []) {
        if (
          call.status === "completed" &&
          call.presentation?.refresh.length &&
          !seenChanges.current.has(call.id)
        ) {
          seenChanges.current.add(call.id);
          call.presentation.refresh.forEach((resource) =>
            changed.add(resource),
          );
        }
      }
    if (changed.size) changedRef.current([...changed]);
  }, []);
  const transport = useMemo(() => chatTransport(() => current.current), []);
  const chat = useChat<ChatMessage>({
    id: conversationId || "unselected",
    transport,
    experimental_throttle: 50,
    onData: (part) => {
      if (part.type === "data-run") {
        const next = part.data;
        if (next.conversation_id !== current.current.conversationId) return;
        observeChanges([next]);
        setRuns((items) => {
          const old = items.find((run) => run.id === next.id);
          const merged = {
            ...old,
            ...next,
            tools: next.tools ?? old?.tools ?? [],
          };
          return [...items.filter((run) => run.id !== next.id), merged];
        });
      }
    },
    onFinish: async () => {
      await refreshRef.current();
    },
    onError: (error) => {
      setFailure(error.message);
      void refreshRef.current();
    },
  });
  const refreshConversation = useCallback(async () => {
    const id = current.current.conversationId;
    if (!id) return;
    try {
      const detail = await api<ConversationDetail>(`/ai/conversations/${id}`);
      if (current.current.conversationId !== id) return;
      chat.setMessages(detail.messages);
      setRuns(detail.runs);
      observeChanges(detail.runs);
      setConversations((items) =>
        items.map((item) =>
          item.id === id ? { id, title: detail.title } : item,
        ),
      );
    } catch (e) {
      setFailure((e as Error).message);
    }
  }, [chat.setMessages, observeChanges]);
  refreshRef.current = refreshConversation;

  useEffect(() => {
    let active = true;
    const loadSettings = () =>
      api<ModelSettings>("/ai/settings")
        .then((value) => {
          if (active) setSettings(value);
        })
        .catch((e) => {
          if (active) setFailure(e.message);
        });
    void loadSettings();
    api<Conversation[]>("/ai/conversations")
      .then((items) => {
        if (active) {
          setConversations(items);
          setConversationId(items[0]?.id ?? "");
        }
      })
      .catch((e) => {
        if (active) setFailure(e.message);
      });
    window.addEventListener("ai-settings-changed", loadSettings);
    return () => {
      active = false;
      window.removeEventListener("ai-settings-changed", loadSettings);
    };
  }, []);
  useEffect(() => {
    setRuns([]);
    setFailure("");
    setLoading(true);
    void refreshConversation().finally(() => setLoading(false));
  }, [conversationId, refreshConversation]);
  useEffect(() => {
    if (action?.type !== "open_conversation") return;
    let live = true;
    void api<ConversationDetail>(
      `/ai/conversations/${encodeURIComponent(action.conversation_id)}`,
    )
      .then(async (detail) => {
        if (!live) return;
        await chat.stop();
        if (!live) return;
        setConversations((items) =>
          items.some((item) => item.id === detail.id)
            ? items
            : [{ id: detail.id, title: detail.title }, ...items],
        );
        setConversationId(detail.id);
        setText("");
        setFailure("");
      })
      .catch((e) => {
        if (live) setFailure(e.message);
      });
    return () => {
      live = false;
    };
  }, [action]);
  const streaming = chat.status === "streaming" || chat.status === "submitted";
  const activeRun = runs.find((run) =>
    ["running", "waiting_approval"].includes(run.status),
  );
  useEffect(() => {
    if (!conversationId || streaming || activeRun?.status !== "running") return;
    const timer = window.setInterval(() => void refreshRef.current(), 3000);
    return () => clearInterval(timer);
  }, [conversationId, streaming, activeRun?.status]);
  useEffect(() => {
    if (!open) return;
    previousFocus.current = document.activeElement as HTMLElement;
    const textarea = input.current?.querySelector("textarea");
    if (textarea && !textarea.disabled) textarea.focus();
    else
      root.current
        ?.querySelector<HTMLButtonElement>('button[aria-label="收起研究助手"]')
        ?.focus();
    return () => {
      requestAnimationFrame(() => {
        if (
          previousFocus.current?.isConnected &&
          previousFocus.current !== document.body
        )
          previousFocus.current.focus();
        else
          document
            .querySelector<HTMLElement>('[aria-label="打开研究助手"]')
            ?.focus();
      });
    };
  }, [open]);
  useEffect(() => {
    if (open) bottom.current?.scrollIntoView({ block: "nearest" });
  }, [chat.messages, open]);

  async function createConversation() {
    setLoading(true);
    try {
      await chat.stop();
      const value = await post<Conversation>("/ai/conversations");
      setConversations((items) => [value, ...items]);
      setConversationId(value.id);
      setText("");
    } catch (e) {
      setFailure((e as Error).message);
    } finally {
      setLoading(false);
    }
  }
  async function send() {
    if (!text.trim() || !conversationId || activeRun || streaming) return;
    const message = text;
    setText("");
    setFailure("");
    chat.clearError();
    try {
      await chat.sendMessage({ text: message });
    } catch (e) {
      setFailure((e as Error).message);
    }
  }
  async function decide(id: string, approved: boolean) {
    setFailure("");
    chat.clearError();
    try {
      await chat.sendMessage(undefined, {
        body: { decision: { id, approved } },
      });
    } catch (e) {
      setFailure((e as Error).message);
    }
  }
  async function stop() {
    if (!activeRun) return;
    try {
      await post(`/ai/runs/${activeRun.id}/cancel`);
      await chat.stop();
      await refreshConversation();
      // Completed tools are reconciled by refreshConversation; cancelling has no new business effect.
    } catch (e) {
      setFailure((e as Error).message);
    }
  }
  const cards = new Map(
    runs.flatMap((run) => run.tools ?? []).map((call) => [call.id, call]),
  );
  const lastRun = runs.at(-1);
  const ready = !!settings?.enabled && settings.ready;
  return (
    <aside
      ref={root}
      hidden={!open}
      className="ai-chat"
      id="research-assistant"
      aria-label="AI 研究助手"
      onKeyDown={(event) => {
        if (event.key === "Escape") {
          event.stopPropagation();
          onClose();
        }
        if (event.key === "Tab" && window.innerWidth < 1440) {
          const elements = [
            ...(root.current?.querySelectorAll<HTMLElement>(
              'button:not(:disabled), textarea:not(:disabled), input:not(:disabled), [tabindex="0"]',
            ) ?? []),
          ].filter((el) => el.offsetParent !== null);
          if (event.shiftKey && document.activeElement === elements[0]) {
            event.preventDefault();
            elements.at(-1)?.focus();
          } else if (
            !event.shiftKey &&
            document.activeElement === elements.at(-1)
          ) {
            event.preventDefault();
            elements[0]?.focus();
          }
        }
      }}
    >
      <div
        className="ai-resize"
        role="separator"
        aria-label="调整助手宽度"
        aria-orientation="vertical"
        aria-valuemin={360}
        aria-valuemax={640}
        aria-valuenow={width}
        tabIndex={0}
        onKeyDown={(event) => {
          if (event.key === "ArrowLeft") onWidth(Math.min(640, width + 20));
          if (event.key === "ArrowRight") onWidth(Math.max(360, width - 20));
        }}
        onPointerDown={(event) => {
          event.currentTarget.setPointerCapture(event.pointerId);
        }}
        onPointerMove={(event) => {
          if (event.currentTarget.hasPointerCapture(event.pointerId))
            onWidth(
              Math.max(360, Math.min(640, window.innerWidth - event.clientX)),
            );
        }}
      />
      <header className="ai-header">
        <h2>AI 研究助手</h2>
        <Button
          type="tertiary"
          aria-label="收起研究助手"
          theme="borderless"
          onClick={onClose}
        >
          收起
        </Button>
      </header>
      <div className="ai-conversations">
        <Select
          aria-label="选择聊天会话"
          value={conversationId || undefined}
          placeholder="选择会话"
          optionList={conversations.map((item) => ({
            value: item.id,
            label: item.title,
          }))}
          onChange={async (value) => {
            await chat.stop();
            setConversationId(String(value));
          }}
          disabled={loading}
        />
        <Button
          type="tertiary"
          onClick={() => void createConversation()}
          disabled={loading}
        >
          新会话
        </Button>
      </div>
      {!ready && (
        <Banner
          type="info"
          description={
            <span>
              配置并测试模型服务后即可聊天。
              <button className="text-link" onClick={onSettings}>
                前往配置
              </button>
            </span>
          }
        />
      )}
      {failure && <Banner type="danger" description={failure} />}
      <div className="ai-messages" aria-live="polite" aria-busy={streaming}>
        {loading && <Spin />}
        {!chat.messages.length && (
          <div className="ai-empty">
            <h3>从当前研究出发</h3>
            <p>
              可以让我选择数据字段、构建候选并预览回测，也可以查询 Alpha
              和已有结果。
            </p>
            {!conversationId && (
              <Button
                type="tertiary"
                onClick={() => void createConversation()}
                disabled={loading}
              >
                开始新会话
              </Button>
            )}
            <span>修改和任务操作会先展示预览，确认后执行。</span>
          </div>
        )}
        {chat.messages.map((message) => (
          <article className={`ai-message ${message.role}`} key={message.id}>
            <span className="ai-role">
              {message.role === "user" ? "你" : "研究助手"}
            </span>
            {message.parts.map((part, index) =>
              part.type === "text" ? (
                message.role === "assistant" ? (
                  <Suspense
                    key={index}
                    fallback={<div className="ai-text">{part.text}</div>}
                  >
                    <MessageMarkdown
                      content={part.text}
                      final={
                        !streaming ||
                        message.id !== chat.messages.at(-1)?.id ||
                        part.state === "done"
                      }
                    />
                  </Suspense>
                ) : (
                  <div className="ai-text" key={index}>
                    {part.text}
                  </div>
                )
              ) : part.type === "data-tool" ? (
                <BusinessCard
                  key={part.id ?? index}
                  call={cards.get(part.data.id) ?? part.data}
                  onDecision={decide}
                  disabled={
                    streaming || activeRun?.status !== "waiting_approval"
                  }
                  onAction={onAction}
                  timezone={timezone}
                  jobs={jobs}
                />
              ) : null,
            )}
          </article>
        ))}
        {lastRun && (
          <div className="ai-run-state">
            <Tag color={lastRun.status === "failed" ? "red" : "grey"}>
              {runLabels[lastRun.status]}
            </Tag>
            {lastRun.error && <p className="error-text">{lastRun.error}</p>}
            {lastRun.usage && (
              <small>
                {lastRun.model} · {(lastRun.elapsed_ms / 1000).toFixed(1)} 秒 ·{" "}
                {lastRun.usage.token_usage_known &&
                (lastRun.usage.input_tokens ?? 0) +
                  (lastRun.usage.output_tokens ?? 0) >
                  0
                  ? `${lastRun.usage.input_tokens} 输入 / ${lastRun.usage.output_tokens} 输出 tokens`
                  : "用量未提供"}
              </small>
            )}
          </div>
        )}
        <div ref={bottom} />
      </div>
      <footer className="ai-composer" ref={input}>
        <div className="ai-context">{contextLabel(context)}</div>
        <TextArea
          aria-label="发送给研究助手"
          placeholder={
            conversationId
              ? "询问当前 Alpha，或描述要执行的操作…"
              : "请先新建会话"
          }
          value={text}
          onChange={setText}
          rows={3}
          maxCount={20000}
          disabled={!ready || !conversationId || loading}
          onKeyDown={(event) => {
            if (
              event.key === "Enter" &&
              !event.shiftKey &&
              !event.nativeEvent.isComposing
            ) {
              event.preventDefault();
              void send();
            }
          }}
        />
        <div className="ai-send">
          <small>Enter 发送 · Shift + Enter 换行</small>
          {activeRun ? (
            <Button type="tertiary" onClick={() => void stop()}>
              停止生成
            </Button>
          ) : (
            <Button
              theme="solid"
              disabled={
                !ready ||
                !text.trim() ||
                !conversationId ||
                loading ||
                streaming
              }
              onClick={() => void send()}
            >
              发送
            </Button>
          )}
        </div>
      </footer>
    </aside>
  );
}