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

import {
  useCallback,
  useEffect,
  useRef,
  useState,
  type ReactNode,
} from "react";
import { AIChatInput, Button, Tag } from "@douyinfe/semi-ui-19";
import {
  IconAIEditLevel3,
  IconAIEditLevel1,
  IconAIFilledLevel3,
  IconStop,
} from "@douyinfe/semi-icons";

export type ResearchReference = { id: string; name: string };
type Editor = ReturnType<AIChatInput["getEditor"]>;

// JSON keeps expressions and user text literal instead of interpreting them as HTML.
function textDocument(value: string) {
  return {
    type: "doc",
    content: value.split("\n").map((text) => ({
      type: "paragraph",
      ...(text ? { content: [{ type: "text", text }] } : {}),
    })),
  };
}
function editorText(editor?: Editor) {
  return editor?.getText({ blockSeparator: "\n" }).replace(/\u200b/g, "") ?? "";
}

/** Plain-text boundary for Semi's editor; references describe existing business context. */
export function ResearchInput({
  value,
  onChange,
  label,
  placeholder,
  references = [],
  disabled = false,
  canSubmit = true,
  busy = false,
  onSubmit,
  onStop,
  submitLabel = "发送",
  hint,
  prompts = [],
  sendHotKey = "enter",
  maxLength = 20000,
}: {
  value: string;
  onChange: (value: string) => void;
  label: string;
  placeholder: string;
  references?: ResearchReference[];
  disabled?: boolean;
  canSubmit?: boolean;
  busy?: boolean;
  onSubmit: () => void;
  onStop?: () => void;
  submitLabel?: string;
  hint?: ReactNode;
  prompts?: string[];
  sendHotKey?: "enter" | "shift+enter";
  maxLength?: number;
}) {
  const input = useRef<AIChatInput>(null);
  const [editor, setEditor] = useState<Editor>();
  const initialContent = useRef(textDocument(value));
  const currentPlaceholder = useRef(placeholder);
  currentPlaceholder.current = placeholder;
  // Tiptap initializes its placeholder extension once; read current availability lazily.
  const editorPlaceholder = useCallback(() => currentPlaceholder.current, []);
  const tooLong = value.length > maxLength;
  const allowed = !disabled && !busy && canSubmit && !!value.trim() && !tooLong;
  useEffect(() => {
    // Semi exposes its editor after its child mounts, including under StrictMode.
    let frame: number;
    const connect = () => {
      const next = input.current?.getEditor();
      if (next && !next.isDestroyed) setEditor(next);
      else frame = requestAnimationFrame(connect);
    };
    frame = requestAnimationFrame(connect);
    return () => cancelAnimationFrame(frame);
  }, []);
  useEffect(() => {
    if (!editor || editor.isDestroyed) return;
    editor.setOptions({
      editorProps: {
        ...editor.options.editorProps,
        attributes: {
          role: "textbox",
          "aria-label": label,
          "aria-multiline": "true",
          "aria-disabled": String(disabled),
          "aria-invalid": String(tooLong),
        },
      },
    });
    editor.setEditable(!disabled, false);
  }, [editor, label, disabled, tooLong]);
  useEffect(() => {
    if (editor && !editor.isDestroyed && editorText(editor) !== value)
      editor.commands.setContent(textDocument(value));
  }, [editor, value]);
  const submit = () => {
    if (allowed) onSubmit();
  };
  return (
    <div
      className="research-ai-input"
      onKeyDownCapture={(event) => {
        if (event.key === "Enter" && event.nativeEvent.isComposing)
          event.stopPropagation();
      }}
    >
      <AIChatInput
        ref={input}
        className="research-ai-editor"
        defaultContent={initialContent.current}
        placeholder={editorPlaceholder}
        clearContentOnGenerating={false}
        generating={busy}
        canSend={allowed}
        sendHotKey={sendHotKey}
        showUploadButton={false}
        showUploadFile={false}
        references={references.map((reference) => ({
          ...reference,
          type: "text",
        }))}
        renderReference={(reference) => (
          <span
            key={reference.id}
            title={reference.name}
            className="research-ai-reference"
          >
            <Tag
              colorful
              prefixIcon={
                <IconAIFilledLevel3 size="small" aria-hidden="true" />
              }
            >
              {reference.name}
            </Tag>
          </span>
        )}
        onContentChange={() => onChange(editorText(input.current?.getEditor()))}
        onMessageSend={submit}
        onStopGenerate={onStop}
        renderConfigureArea={() => (
          <span className="research-ai-shortcut">
            {sendHotKey === "enter"
              ? "Enter 发送 · Shift + Enter 换行"
              : "Shift + Enter 生成 · Enter 换行"}
          </span>
        )}
        renderActionArea={() =>
          busy && onStop ? (
            <Button
              colorful
              icon={<IconStop aria-hidden="true" />}
              onClick={onStop}
            >
              停止生成
            </Button>
          ) : (
            <Button
              colorful
              theme="solid"
              icon={<IconAIEditLevel1 aria-hidden="true" />}
              disabled={!allowed}
              loading={busy}
              onClick={submit}
            >
              {submitLabel}
            </Button>
          )
        }
      />
      {!value && !busy && prompts.length > 0 && (
        <div className="research-ai-prompts" aria-label="研究提示">
          {prompts.map((prompt) => (
            <Button
              key={prompt}
              colorful
              theme="borderless"
              size="small"
              icon={<IconAIEditLevel3 aria-hidden="true" />}
              disabled={disabled || busy}
              onClick={() => {
                onChange(prompt);
                input.current?.focusEditor("end");
              }}
            >
              {prompt}
            </Button>
          ))}
        </div>
      )}
      {tooLong && (
        <p className="error-text" role="alert">
          内容不能超过 {maxLength.toLocaleString("en-US")}{" "}
          字符，请精简后再发送。
        </p>
      )}
      {hint && <div className="research-ai-hint">{hint}</div>}
    </div>
  );
}