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

import { useChat } from "@ai-sdk/react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  Banner,
  Button,
  Select,
  Spin,
  Tag,
  TextArea,
  Toast,
} from "@douyinfe/semi-ui-19";
import {
  api,
  formatNumber,
  formatTime,
  jobLabels,
  jobStateLabels,
  post,
  stateLabels,
} from "../api";
import { BacktestToolCard } from "../backtests/BacktestToolCard";
import { PnlChart } from "../components/PnlChart";
import type { Alpha, Job, Pnl, Research } from "../types";
import { chatTransport } from "./transport";
import { runLabels, toolLabels } from "./types";
import type {
  ChatMessage,
  Conversation,
  ConversationDetail,
  ModelSettings,
  PageContext,
  RunSnapshot,
  ToolCard,
  UIAction,
} from "./types";

export function ChatPanel({
  open,
  context,
  onClose,
  onAction,
  onChanged,
  onSettings,
  width,
  onWidth,
  timezone,
  jobs,
}: {
  open: boolean;
  context: PageContext;
  onClose: () => void;
  onAction: (action: UIAction) => void;
  onChanged: () => void;
  onSettings: () => void;
  width: number;
  onWidth: (width: number) => void;
  timezone?: string;
  jobs: Job[];
}) {
  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 seenWrites = 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 observeWrites = useCallback((items: RunSnapshot[]) => {
    let changed = false;
    for (const run of items)
      for (const call of run.tools ?? []) {
        if (
          call.status === "completed" &&
          [
            "update_research",
            "bulk_update_research",
            "create_sync_job",
            "cancel_job",
            "retry_job",
            "start_backtest",
            "control_backtest",
          ].includes(call.name) &&
          !seenWrites.current.has(call.id)
        ) {
          seenWrites.current.add(call.id);
          changed = true;
        }
      }
    if (changed) changedRef.current();
  }, []);
  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;
        observeWrites([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);
      observeWrites(detail.runs);
      setConversations((items) =>
        items.map((item) =>
          item.id === id ? { id, title: detail.title } : item,
        ),
      );
    } catch (e) {
      setFailure((e as Error).message);
    }
  }, [chat.setMessages, observeWrites]);
  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]);
  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();
      changedRef.current();
    } 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" ? (
                <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">
          {context.page === "backtests"
            ? "上下文：回测研究"
            : context.page === "datasets"
              ? "上下文：数据目录（未发送字段与备注）"
              : context.page === "account"
                ? "上下文：个人信息页"
                : `上下文：${context.alpha_id ? `Alpha ${context.alpha_id}` : "Alpha 列表"}${context.selected_ids?.length ? ` · 已选 ${context.selected_ids.length} 条` : ""}`}
        </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>
  );
}

function BusinessCard({
  call,
  onDecision,
  disabled,
  onAction,
  timezone,
  jobs,
}: {
  call: ToolCard;
  onDecision: (id: string, approved: boolean) => Promise<void>;
  disabled: boolean;
  onAction: (action: UIAction) => void;
  timezone?: string;
  jobs: Job[];
}) {
  const result = call.result ?? {};
  const jobId = result.job_id ?? result.id;
  const job = jobs.find((item) => item.id === jobId);
  const jobOperation = [
    "get_job_status",
    "create_sync_job",
    "retry_job",
    "cancel_job",
  ].includes(call.name);
  const pending = call.status === "pending";
  const items = Array.isArray(result.items) ? result.items : [];
  const isSearch = call.name === "search_alphas";
  const alphas = (
    isSearch ? items : call.name === "get_alpha" && result.id ? [result] : []
  ) as Alpha[];
  const labels: Record<string, string> = {
    pending: "待确认",
    completed: "已完成",
    failed: "失败",
    denied: "已拒绝",
    cancelled: "已取消",
  };
  const openAlpha = (id: string) =>
    onAction({ type: "open_alpha", alpha_id: id, nonce: Date.now() });
  return (
    <section className="ai-tool-card" data-status={call.status}>
      <div className="ai-card-title">
        <strong>{toolLabels[call.name] ?? "业务操作"}</strong>
        <Tag
          color={pending ? "orange" : call.status === "failed" ? "red" : "grey"}
        >
          {labels[call.status] ?? call.status}
        </Tag>
      </div>
      {call.name.includes("backtest") && (
        <BacktestToolCard call={call} onAction={onAction} />
      )}
      {call.preview.targets?.map((target) => (
        <details
          key={target.alpha_id}
          open={call.preview.targets!.length === 1}
        >
          <summary>{target.alpha_id} 的修改</summary>
          {(["note", "tags", "favorite", "state"] as const)
            .filter(
              (key) =>
                JSON.stringify(target.before[key]) !==
                JSON.stringify(target.after[key]),
            )
            .map((key) => (
              <div className="ai-diff" key={key}>
                <strong>
                  {
                    {
                      note: "备注",
                      tags: "标签",
                      favorite: "收藏",
                      state: "研究状态",
                    }[key]
                  }
                </strong>
                <del>
                  <span>修改前</span>
                  {researchValue(key, target.before[key])}
                </del>
                <ins>
                  <span>修改后</span>
                  {researchValue(key, target.after[key])}
                </ins>
              </div>
            ))}
        </details>
      ))}
      {call.preview.operation && (
        <p>
          {jobLabels[String(call.preview.operation.kind)] ?? "创建任务"} ·{" "}
          {Array.isArray(call.preview.operation.alpha_ids) &&
          call.preview.operation.alpha_ids.length
            ? call.preview.operation.alpha_ids.join("、")
            : call.preview.operation.submission === "UNSUBMITTED"
              ? "待提交 Alpha"
              : call.preview.operation.submission === "SUBMITTED"
                ? "已提交 Alpha"
                : "全部 Alpha"}
          {call.preview.operation.kind === "daily_sync" && (
            <>
              {" · "}
              {call.preview.operation.submission === "UNSUBMITTED"
                ? "创建日期"
                : "提交日期"}{" "}
              {String(call.preview.operation.date_from)} 至{" "}
              {String(call.preview.operation.date_to)}（UTC）
            </>
          )}
        </p>
      )}
      {call.preview.job && (
        <p>
          {jobLabels[String(call.preview.job.kind)]} ·{" "}
          {String(call.preview.job.id)} ·{" "}
          {jobStateLabels[String(call.preview.job.status)]}
        </p>
      )}
      {pending && (
        <div className="ai-approval">
          <p className="muted">
            {call.preview.targets?.length
              ? `将修改 ${call.preview.targets.length} 条研究记录。`
              : "将执行以上任务操作。"}
            确认后执行。
          </p>
          <div className="inline-actions">
            <Button
              theme="solid"
              disabled={disabled}
              onClick={() => void onDecision(call.id, true)}
            >
              确认执行
            </Button>
            <Button
              type="tertiary"
              disabled={disabled}
              onClick={() => void onDecision(call.id, false)}
            >
              拒绝
            </Button>
          </div>
        </div>
      )}
      {typeof result.error === "string" && (
        <p className="error-text">{result.error}</p>
      )}
      {result.denied === true && <p>已拒绝，未执行此操作。</p>}
      {alphas.length > 0 && (
        <div className="ai-alpha-results">
          {alphas.map((alpha) => (
            <div key={alpha.id}>
              <button className="text-link" onClick={() => openAlpha(alpha.id)}>
                {alpha.name || alpha.id}
              </button>
              <small>
                {alpha.id} · {alpha.region ?? "地区未提供"}
              </small>
              <table>
                <thead>
                  <tr>
                    <th>Sharpe</th>
                    <th>Fitness</th>
                    <th>Turnover</th>
                  </tr>
                </thead>
                <tbody>
                  <tr>
                    <td>{formatNumber(alpha.sharpe)}</td>
                    <td>{formatNumber(alpha.fitness)}</td>
                    <td>
                      {alpha.turnover == null
                        ? "未提供"
                        : `${formatNumber(alpha.turnover * 100)}%`}
                    </td>
                  </tr>
                </tbody>
              </table>
              <small>本地快照 · {formatTime(alpha.synced_at, timezone)}</small>
            </div>
          ))}
        </div>
      )}
      {isSearch && (
        <p>
          共 {String(result.total ?? 0)} 条，当前返回 {items.length} 条。
          {!!result.filters && (
            <button
              className="text-link"
              onClick={() =>
                onAction({
                  type: "apply_filters",
                  filters: result.filters as Record<string, unknown>,
                  nonce: Date.now(),
                })
              }
            >
              应用到列表
            </button>
          )}
        </p>
      )}
      {call.name === "get_alpha_facets" && call.status === "completed" && (
        <p>
          已同步 {String(result.total ?? 0)} 条 Alpha · 收藏{" "}
          {String(result.favorites ?? 0)} 条<br />
          地区：
          {Array.isArray(result.region) ? result.region.join("、") : "未提供"}
        </p>
      )}
      {call.name === "get_alpha_pnl" && call.status === "completed" && (
        <ChatPnl result={result} timezone={timezone} />
      )}
      {call.name === "list_jobs" &&
        items.map((item) => {
          const job = item as Record<string, unknown>;
          return (
            <p key={String(job.id)}>
              {jobLabels[String(job.kind)]} ·{" "}
              {jobStateLabels[String(job.status)]}
              <small>{String(job.id)}</small>
            </p>
          );
        })}
      {jobOperation && !!jobId && (
        <p>
          {jobLabels[job?.kind ?? String(result.kind)] ?? "任务"} ·{" "}
          {jobStateLabels[job?.status ?? String(result.status)] ??
            String(result.status ?? "已处理")}
          {job && (
            <span>
              {" "}
              · {job.processed}/{job.total ?? "未知"}，失败 {job.failed}
              {job.error ? ` · ${job.error}` : ""}
            </span>
          )}
          <small>{String(jobId)}</small>
        </p>
      )}
      {result.ok === true && (
        <p>
          操作已保存。
          {typeof result.alpha_id === "string" && (
            <button
              className="text-link"
              onClick={() => openAlpha(result.alpha_id as string)}
            >
              打开 Alpha
            </button>
          )}
        </p>
      )}
      {typeof result.updated === "number" && (
        <p>已更新 {result.updated} 条研究记录。</p>
      )}
    </section>
  );
}

function researchValue(key: keyof Research, value: unknown) {
  if (key === "favorite") return value ? "已收藏" : "未收藏";
  if (key === "state")
    return (
      stateLabels[String(value) as keyof typeof stateLabels] ?? String(value)
    );
  if (Array.isArray(value)) return value.join("、") || "无";
  return String(value || "空");
}

function ChatPnl({
  result,
  timezone,
}: {
  result: Record<string, unknown>;
  timezone?: string;
}) {
  const [pnl, setPnl] = useState<Pnl | null>(null);
  useEffect(() => {
    let active = true;
    if (result.cached && typeof result.alpha_id === "string")
      api<Pnl>(`/alphas/${encodeURIComponent(result.alpha_id)}/pnl`)
        .then((value) => {
          if (active) setPnl(value);
        })
        .catch((error) => {
          if (active) Toast.error(error.message);
        });
    return () => {
      active = false;
    };
  }, [result.alpha_id, result.fetched_at, result.cached]);
  return result.cached ? (
    <div>
      {pnl && <PnlChart points={pnl.points} />}
      <small>
        {String(result.count)} 条记录 · 缓存于{" "}
        {formatTime(result.fetched_at as string, timezone)}
      </small>
    </div>
  ) : (
    <p>尚无 PnL 缓存，可以请求创建 PnL 刷新任务。</p>
  );
}