"use client";
import { Node, mergeAttributes, nodeInputRule } from "@tiptap/react";
import {
NodeViewWrapper,
ReactNodeViewRenderer,
type NodeViewProps,
} from "@tiptap/react";
import { streamText } from "ai";
import {
Check,
Languages,
Loader2,
Minimize2,
RefreshCw,
Sparkles,
TextQuote,
Trash2,
Wand2,
X,
} from "lucide-react";
import * as React from "react";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { cn } from "@/lib/utils";
import {
createOpenAIClient,
EDITOR_SYSTEM_PROMPT,
expand,
rewrite,
selectOpenAIModel,
summarize,
translate,
} from "@tasks/ai";
type AiStatus = "idle" | "loading" | "done" | "error";
function escapeHtml(s: string) {
return s
.replace(/&/g, "&")
.replace(//g, ">")
.replace(/"/g, """);
}
function plainTextToHtml(text: string) {
const trimmed = text.trim();
if (!trimmed) return "
";
return trimmed
.split(/\n\n+/)
.map((block) => `${escapeHtml(block).replace(/\n/g, "
")}
`)
.join("");
}
async function runEditorGeneration(prompt: string, signal: AbortSignal) {
const client = createOpenAIClient();
if (!client.ok) {
throw new Error(
"OpenAI API key is not configured. Set OPENAI_API_KEY or NEXT_PUBLIC_OPENAI_API_KEY.",
);
}
const model = selectOpenAIModel(client.provider, "gpt-4o-mini");
const result = streamText({
model,
abortSignal: signal,
system: EDITOR_SYSTEM_PROMPT,
prompt: prompt.trim(),
});
let out = "";
for await (const chunk of result.textStream) {
out += chunk;
}
return out;
}
async function runActionPair(
systemPrompt: string,
userPrompt: string,
signal: AbortSignal,
) {
const client = createOpenAIClient();
if (!client.ok) {
throw new Error(
"OpenAI API key is not configured. Set OPENAI_API_KEY or NEXT_PUBLIC_OPENAI_API_KEY.",
);
}
const model = selectOpenAIModel(client.provider, "gpt-4o-mini");
const result = streamText({
model,
abortSignal: signal,
system: systemPrompt,
prompt: userPrompt,
});
let out = "";
for await (const chunk of result.textStream) {
out += chunk;
}
return out;
}
function LoadingDots() {
return (
{[0, 1, 2].map((i) => (
))}
);
}
function AiBlockView(props: NodeViewProps) {
const { node, editor, getPos, updateAttributes, deleteNode } = props;
const status = (node.attrs.status as AiStatus) ?? "idle";
const prompt = (node.attrs.prompt as string) ?? "";
const resultHtml = (node.attrs.resultHtml as string) ?? "";
const errorMessage = (node.attrs.errorMessage as string) ?? "";
const inputRef = React.useRef(null);
const abortRef = React.useRef(null);
React.useEffect(() => {
if (status === "idle" && !resultHtml) {
inputRef.current?.focus();
}
}, [status, resultHtml]);
const [selTick, setSelTick] = React.useState(0);
React.useEffect(() => {
const bump = () => setSelTick((n) => n + 1);
editor.on("selectionUpdate", bump);
editor.on("transaction", bump);
return () => {
editor.off("selectionUpdate", bump);
editor.off("transaction", bump);
};
}, [editor]);
const selectionText = React.useMemo(() => {
const { from, to } = editor.state.selection;
if (from === to) return "";
return editor.state.doc.textBetween(from, to, "\n");
// eslint-disable-next-line react-hooks/exhaustive-deps -- sync selection text when editor updates
}, [editor, selTick]);
const stopGeneration = React.useCallback(() => {
abortRef.current?.abort();
abortRef.current = null;
}, []);
const runPrompt = React.useCallback(
async (text: string) => {
const trimmed = text.trim();
if (!trimmed) return;
stopGeneration();
const ac = new AbortController();
abortRef.current = ac;
updateAttributes({
prompt: trimmed,
status: "loading",
resultHtml: "",
errorMessage: "",
});
try {
const out = await runEditorGeneration(trimmed, ac.signal);
updateAttributes({
status: "done",
resultHtml: plainTextToHtml(out),
errorMessage: "",
});
} catch (e) {
if ((e as Error).name === "AbortError") return;
updateAttributes({
status: "error",
errorMessage: (e as Error).message ?? "Something went wrong.",
});
} finally {
abortRef.current = null;
}
},
[stopGeneration, updateAttributes],
);
const runQuick = React.useCallback(
async (kind: "summarize" | "expand" | "simplify" | "translate") => {
const base = selectionText.trim() || prompt.trim();
if (!base) {
updateAttributes({
status: "error",
errorMessage: "Select text in the editor or enter a prompt first.",
});
return;
}
stopGeneration();
const ac = new AbortController();
abortRef.current = ac;
updateAttributes({
status: "loading",
resultHtml: "",
errorMessage: "",
});
try {
let pair;
if (kind === "summarize") pair = summarize(base);
else if (kind === "expand") pair = expand(base);
else if (kind === "simplify") pair = rewrite(base, "concise");
else pair = translate(base, "Spanish");
const out = await runActionPair(
pair.systemPrompt,
pair.userPrompt,
ac.signal,
);
updateAttributes({
status: "done",
prompt: prompt || `[${kind}]`,
resultHtml: plainTextToHtml(out),
errorMessage: "",
});
} catch (e) {
if ((e as Error).name === "AbortError") return;
updateAttributes({
status: "error",
errorMessage: (e as Error).message ?? "Something went wrong.",
});
} finally {
abortRef.current = null;
}
},
[prompt, selectionText, stopGeneration, updateAttributes],
);
const accept = React.useCallback(() => {
const pos = getPos();
if (typeof pos !== "number") return;
const html = resultHtml || "";
editor
.chain()
.focus()
.deleteRange({ from: pos, to: pos + node.nodeSize })
.insertContentAt(pos, html)
.run();
}, [editor, getPos, node.nodeSize, resultHtml]);
const discard = React.useCallback(() => {
deleteNode();
}, [deleteNode]);
const regenerate = React.useCallback(() => {
if (prompt.trim()) void runPrompt(prompt);
}, [prompt, runPrompt]);
const onKeyDownInput = (e: React.KeyboardEvent) => {
if (e.key === "Enter") {
e.preventDefault();
void runPrompt(inputRef.current?.value ?? "");
}
if (e.key === "Escape") {
e.preventDefault();
discard();
}
};
return (
AI assistant
/ai
Ask in natural language, then accept or regenerate.
{status !== "loading" && !(status === "done" && resultHtml) ? (
) : null}
{status === "loading" ? (
Generating
) : null}
{status === "error" && errorMessage ? (
{errorMessage}
) : null}
{status === "done" && resultHtml ? (
) : null}
{status === "idle" ||
status === "loading" ||
status === "error" ? (
) : null}
);
}
function QuickBtn({
icon: Icon,
label,
onClick,
}: {
icon: React.ComponentType<{ className?: string }>;
label: string;
onClick: () => void;
}) {
return (
);
}
export const AiBlock = Node.create({
name: "aiBlock",
group: "block",
atom: true,
draggable: true,
addAttributes() {
return {
prompt: { default: "" },
status: { default: "idle" },
resultHtml: { default: "" },
errorMessage: { default: "" },
};
},
parseHTML() {
return [{ tag: 'div[data-type="ai-block"]' }];
},
renderHTML({
HTMLAttributes,
}: {
HTMLAttributes: Record;
}) {
return [
"div",
mergeAttributes(HTMLAttributes, { "data-type": "ai-block" }),
];
},
addNodeView() {
return ReactNodeViewRenderer(AiBlockView);
},
addInputRules() {
return [
nodeInputRule({
find: /(^|\s)\/ai$/,
type: this.type,
getAttributes: () => ({
prompt: "",
status: "idle",
resultHtml: "",
errorMessage: "",
}),
}),
];
},
});