ai-design-system 0.1.0

package/components/ai-elements/branch.tsx

@@ -0,0 +1,212 @@
+"use client";
+
+import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
+import type { UIMessage } from "ai";
+import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
+import type { ComponentProps, HTMLAttributes, ReactElement } from "react";
+import { createContext, useContext, useEffect, useState } from "react";
+
+type BranchContextType = {
+  currentBranch: number;
+  totalBranches: number;
+  goToPrevious: () => void;
+  goToNext: () => void;
+  branches: ReactElement[];
+  setBranches: (branches: ReactElement[]) => void;
+};
+
+const BranchContext = createContext<BranchContextType | null>(null);
+
+const useBranch = () => {
+  const context = useContext(BranchContext);
+
+  if (!context) {
+    throw new Error("Branch components must be used within Branch");
+  }
+
+  return context;
+};
+
+export type BranchProps = HTMLAttributes<HTMLDivElement> & {
+  defaultBranch?: number;
+  onBranchChange?: (branchIndex: number) => void;
+};
+
+export const Branch = ({
+  defaultBranch = 0,
+  onBranchChange,
+  className,
+  ...props
+}: BranchProps) => {
+  const [currentBranch, setCurrentBranch] = useState(defaultBranch);
+  const [branches, setBranches] = useState<ReactElement[]>([]);
+
+  const handleBranchChange = (newBranch: number) => {
+    setCurrentBranch(newBranch);
+    onBranchChange?.(newBranch);
+  };
+
+  const goToPrevious = () => {
+    const newBranch =
+      currentBranch > 0 ? currentBranch - 1 : branches.length - 1;
+    handleBranchChange(newBranch);
+  };
+
+  const goToNext = () => {
+    const newBranch =
+      currentBranch < branches.length - 1 ? currentBranch + 1 : 0;
+    handleBranchChange(newBranch);
+  };
+
+  const contextValue: BranchContextType = {
+    currentBranch,
+    totalBranches: branches.length,
+    goToPrevious,
+    goToNext,
+    branches,
+    setBranches,
+  };
+
+  return (
+    <BranchContext.Provider value={contextValue}>
+      <div
+        className={cn("grid w-full gap-2 [&>div]:pb-0", className)}
+        {...props}
+      />
+    </BranchContext.Provider>
+  );
+};
+
+export type BranchMessagesProps = HTMLAttributes<HTMLDivElement>;
+
+export const BranchMessages = ({ children, ...props }: BranchMessagesProps) => {
+  const { currentBranch, setBranches, branches } = useBranch();
+  const childrenArray = Array.isArray(children) ? children : [children];
+
+  // Use useEffect to update branches when they change
+  useEffect(() => {
+    if (branches.length !== childrenArray.length) {
+      setBranches(childrenArray);
+    }
+  }, [childrenArray, branches, setBranches]);
+
+  return childrenArray.map((branch, index) => (
+    <div
+      className={cn(
+        "grid gap-2 overflow-hidden [&>div]:pb-0",
+        index === currentBranch ? "block" : "hidden"
+      )}
+      key={branch.key}
+      {...props}
+    >
+      {branch}
+    </div>
+  ));
+};
+
+export type BranchSelectorProps = HTMLAttributes<HTMLDivElement> & {
+  from: UIMessage["role"];
+};
+
+export const BranchSelector = ({
+  className,
+  from,
+  ...props
+}: BranchSelectorProps) => {
+  const { totalBranches } = useBranch();
+
+  // Don't render if there's only one branch
+  if (totalBranches <= 1) {
+    return null;
+  }
+
+  return (
+    <div
+      className={cn(
+        "flex items-center gap-2 self-end px-10",
+        from === "assistant" ? "justify-start" : "justify-end",
+        className
+      )}
+      {...props}
+    />
+  );
+};
+
+export type BranchPreviousProps = ComponentProps<typeof Button>;
+
+export const BranchPrevious = ({
+  className,
+  children,
+  ...props
+}: BranchPreviousProps) => {
+  const { goToPrevious, totalBranches } = useBranch();
+
+  return (
+    <Button
+      aria-label="Previous branch"
+      className={cn(
+        "size-7 shrink-0 rounded-full text-muted-foreground transition-colors",
+        "hover:bg-accent hover:text-foreground",
+        "disabled:pointer-events-none disabled:opacity-50",
+        className
+      )}
+      disabled={totalBranches <= 1}
+      onClick={goToPrevious}
+      size="icon"
+      type="button"
+      variant="ghost"
+      {...props}
+    >
+      {children ?? <ChevronLeftIcon size={14} />}
+    </Button>
+  );
+};
+
+export type BranchNextProps = ComponentProps<typeof Button>;
+
+export const BranchNext = ({
+  className,
+  children,
+  ...props
+}: BranchNextProps) => {
+  const { goToNext, totalBranches } = useBranch();
+
+  return (
+    <Button
+      aria-label="Next branch"
+      className={cn(
+        "size-7 shrink-0 rounded-full text-muted-foreground transition-colors",
+        "hover:bg-accent hover:text-foreground",
+        "disabled:pointer-events-none disabled:opacity-50",
+        className
+      )}
+      disabled={totalBranches <= 1}
+      onClick={goToNext}
+      size="icon"
+      type="button"
+      variant="ghost"
+      {...props}
+    >
+      {children ?? <ChevronRightIcon size={14} />}
+    </Button>
+  );
+};
+
+export type BranchPageProps = HTMLAttributes<HTMLSpanElement>;
+
+export const BranchPage = ({ className, ...props }: BranchPageProps) => {
+  const { currentBranch, totalBranches } = useBranch();
+
+  return (
+    <span
+      className={cn(
+        "font-medium text-muted-foreground text-xs tabular-nums",
+        className
+      )}
+      {...props}
+    >
+      {currentBranch + 1} of {totalBranches}
+    </span>
+  );
+};

package/components/ai-elements/canvas.tsx

@@ -0,0 +1,24 @@
+import { Background, ReactFlow, type ReactFlowProps } from "@xyflow/react";
+import type { ReactNode } from "react";
+import "@xyflow/react/dist/style.css";
+import { Controls } from "./controls";
+
+type CanvasProps = ReactFlowProps & {
+  children?: ReactNode;
+};
+
+export const Canvas = ({ children, ...props }: CanvasProps) => (
+  <ReactFlow
+    deleteKeyCode={["Backspace", "Delete"]}
+    fitView
+    panOnDrag={false}
+    panOnScroll
+    selectionOnDrag={true}
+    zoomOnDoubleClick={false}
+    {...props}
+  >
+    <Background bgColor="var(--sidebar)" />
+    <Controls />
+    {children}
+  </ReactFlow>
+);

package/components/ai-elements/chain-of-thought.tsx

@@ -0,0 +1,228 @@
+"use client";
+
+import { useControllableState } from "@radix-ui/react-use-controllable-state";
+import { Badge } from "@/components/ui/badge";
+import {
+  Collapsible,
+  CollapsibleContent,
+  CollapsibleTrigger,
+} from "@/components/ui/collapsible";
+import { cn } from "@/lib/utils";
+import {
+  BrainIcon,
+  ChevronDownIcon,
+  DotIcon,
+  type LucideIcon,
+} from "lucide-react";
+import type { ComponentProps } from "react";
+import { createContext, memo, useContext, useMemo } from "react";
+
+type ChainOfThoughtContextValue = {
+  isOpen: boolean;
+  setIsOpen: (open: boolean) => void;
+};
+
+const ChainOfThoughtContext = createContext<ChainOfThoughtContextValue | null>(
+  null
+);
+
+const useChainOfThought = () => {
+  const context = useContext(ChainOfThoughtContext);
+  if (!context) {
+    throw new Error(
+      "ChainOfThought components must be used within ChainOfThought"
+    );
+  }
+  return context;
+};
+
+export type ChainOfThoughtProps = ComponentProps<"div"> & {
+  open?: boolean;
+  defaultOpen?: boolean;
+  onOpenChange?: (open: boolean) => void;
+};
+
+export const ChainOfThought = memo(
+  ({
+    className,
+    open,
+    defaultOpen = false,
+    onOpenChange,
+    children,
+    ...props
+  }: ChainOfThoughtProps) => {
+    const [isOpen, setIsOpen] = useControllableState({
+      prop: open,
+      defaultProp: defaultOpen,
+      onChange: onOpenChange,
+    });
+
+    const chainOfThoughtContext = useMemo(
+      () => ({ isOpen, setIsOpen }),
+      [isOpen, setIsOpen]
+    );
+
+    return (
+      <ChainOfThoughtContext.Provider value={chainOfThoughtContext}>
+        <div
+          className={cn("not-prose max-w-prose space-y-4", className)}
+          {...props}
+        >
+          {children}
+        </div>
+      </ChainOfThoughtContext.Provider>
+    );
+  }
+);
+
+export type ChainOfThoughtHeaderProps = ComponentProps<
+  typeof CollapsibleTrigger
+>;
+
+export const ChainOfThoughtHeader = memo(
+  ({ className, children, ...props }: ChainOfThoughtHeaderProps) => {
+    const { isOpen, setIsOpen } = useChainOfThought();
+
+    return (
+      <Collapsible onOpenChange={setIsOpen} open={isOpen}>
+        <CollapsibleTrigger
+          className={cn(
+            "flex w-full items-center gap-2 text-muted-foreground text-sm transition-colors hover:text-foreground",
+            className
+          )}
+          {...props}
+        >
+          <BrainIcon className="size-4" />
+          <span className="flex-1 text-left">
+            {children ?? "Chain of Thought"}
+          </span>
+          <ChevronDownIcon
+            className={cn(
+              "size-4 transition-transform",
+              isOpen ? "rotate-180" : "rotate-0"
+            )}
+          />
+        </CollapsibleTrigger>
+      </Collapsible>
+    );
+  }
+);
+
+export type ChainOfThoughtStepProps = ComponentProps<"div"> & {
+  icon?: LucideIcon;
+  label: string;
+  description?: string;
+  status?: "complete" | "active" | "pending";
+};
+
+export const ChainOfThoughtStep = memo(
+  ({
+    className,
+    icon: Icon = DotIcon,
+    label,
+    description,
+    status = "complete",
+    children,
+    ...props
+  }: ChainOfThoughtStepProps) => {
+    const statusStyles = {
+      complete: "text-muted-foreground",
+      active: "text-foreground",
+      pending: "text-muted-foreground/50",
+    };
+
+    return (
+      <div
+        className={cn(
+          "flex gap-2 text-sm",
+          statusStyles[status],
+          "fade-in-0 slide-in-from-top-2 animate-in",
+          className
+        )}
+        {...props}
+      >
+        <div className="relative mt-0.5">
+          <Icon className="size-4" />
+          <div className="-mx-px absolute top-7 bottom-0 left-1/2 w-px bg-border" />
+        </div>
+        <div className="flex-1 space-y-2">
+          <div>{label}</div>
+          {description && (
+            <div className="text-muted-foreground text-xs">{description}</div>
+          )}
+          {children}
+        </div>
+      </div>
+    );
+  }
+);
+
+export type ChainOfThoughtSearchResultsProps = ComponentProps<"div">;
+
+export const ChainOfThoughtSearchResults = memo(
+  ({ className, ...props }: ChainOfThoughtSearchResultsProps) => (
+    <div className={cn("flex items-center gap-2", className)} {...props} />
+  )
+);
+
+export type ChainOfThoughtSearchResultProps = ComponentProps<typeof Badge>;
+
+export const ChainOfThoughtSearchResult = memo(
+  ({ className, children, ...props }: ChainOfThoughtSearchResultProps) => (
+    <Badge
+      className={cn("gap-1 px-2 py-0.5 font-normal text-xs", className)}
+      variant="secondary"
+      {...props}
+    >
+      {children}
+    </Badge>
+  )
+);
+
+export type ChainOfThoughtContentProps = ComponentProps<
+  typeof CollapsibleContent
+>;
+
+export const ChainOfThoughtContent = memo(
+  ({ className, children, ...props }: ChainOfThoughtContentProps) => {
+    const { isOpen } = useChainOfThought();
+
+    return (
+      <Collapsible open={isOpen}>
+        <CollapsibleContent
+          className={cn(
+            "mt-2 space-y-3",
+            "data-[state=closed]:fade-out-0 data-[state=closed]:slide-out-to-top-2 data-[state=open]:slide-in-from-top-2 text-popover-foreground outline-none data-[state=closed]:animate-out data-[state=open]:animate-in",
+            className
+          )}
+          {...props}
+        >
+          {children}
+        </CollapsibleContent>
+      </Collapsible>
+    );
+  }
+);
+
+export type ChainOfThoughtImageProps = ComponentProps<"div"> & {
+  caption?: string;
+};
+
+export const ChainOfThoughtImage = memo(
+  ({ className, children, caption, ...props }: ChainOfThoughtImageProps) => (
+    <div className={cn("mt-2 space-y-2", className)} {...props}>
+      <div className="relative flex max-h-88 items-center justify-center overflow-hidden rounded-lg bg-muted p-3">
+        {children}
+      </div>
+      {caption && <p className="text-muted-foreground text-xs">{caption}</p>}
+    </div>
+  )
+);
+
+ChainOfThought.displayName = "ChainOfThought";
+ChainOfThoughtHeader.displayName = "ChainOfThoughtHeader";
+ChainOfThoughtStep.displayName = "ChainOfThoughtStep";
+ChainOfThoughtSearchResults.displayName = "ChainOfThoughtSearchResults";
+ChainOfThoughtSearchResult.displayName = "ChainOfThoughtSearchResult";
+ChainOfThoughtContent.displayName = "ChainOfThoughtContent";
+ChainOfThoughtImage.displayName = "ChainOfThoughtImage";

package/components/ai-elements/code-block.tsx

@@ -0,0 +1,179 @@
+"use client";
+
+import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
+import type { Element } from "hast";
+import { CheckIcon, CopyIcon } from "lucide-react";
+import {
+  type ComponentProps,
+  createContext,
+  type HTMLAttributes,
+  useContext,
+  useEffect,
+  useRef,
+  useState,
+} from "react";
+import { type BundledLanguage, codeToHtml, type ShikiTransformer } from "shiki";
+
+type CodeBlockProps = HTMLAttributes<HTMLDivElement> & {
+  code: string;
+  language: BundledLanguage;
+  showLineNumbers?: boolean;
+};
+
+type CodeBlockContextType = {
+  code: string;
+};
+
+const CodeBlockContext = createContext<CodeBlockContextType>({
+  code: "",
+});
+
+const lineNumberTransformer: ShikiTransformer = {
+  name: "line-numbers",
+  line(node: Element, line: number) {
+    node.children.unshift({
+      type: "element",
+      tagName: "span",
+      properties: {
+        className: [
+          "inline-block",
+          "min-w-10",
+          "mr-4",
+          "text-right",
+          "select-none",
+          "text-muted-foreground",
+        ],
+      },
+      children: [{ type: "text", value: String(line) }],
+    });
+  },
+};
+
+export async function highlightCode(
+  code: string,
+  language: BundledLanguage,
+  showLineNumbers = false
+) {
+  const transformers: ShikiTransformer[] = showLineNumbers
+    ? [lineNumberTransformer]
+    : [];
+
+  return await Promise.all([
+    codeToHtml(code, {
+      lang: language,
+      theme: "one-light",
+      transformers,
+    }),
+    codeToHtml(code, {
+      lang: language,
+      theme: "one-dark-pro",
+      transformers,
+    }),
+  ]);
+}
+
+export const CodeBlock = ({
+  code,
+  language,
+  showLineNumbers = false,
+  className,
+  children,
+  ...props
+}: CodeBlockProps) => {
+  const [html, setHtml] = useState<string>("");
+  const [darkHtml, setDarkHtml] = useState<string>("");
+  const mounted = useRef(false);
+
+  useEffect(() => {
+    highlightCode(code, language, showLineNumbers).then(([light, dark]) => {
+      if (!mounted.current) {
+        setHtml(light);
+        setDarkHtml(dark);
+        mounted.current = true;
+      }
+    });
+
+    return () => {
+      mounted.current = false;
+    };
+  }, [code, language, showLineNumbers]);
+
+  return (
+    <CodeBlockContext.Provider value={{ code }}>
+      <div
+        className={cn(
+          "group relative w-full overflow-hidden rounded-md border bg-background text-foreground",
+          className
+        )}
+        {...props}
+      >
+        <div className="relative">
+          <div
+            className="overflow-hidden dark:hidden [&>pre]:m-0 [&>pre]:bg-background! [&>pre]:p-4 [&>pre]:text-foreground! [&>pre]:text-sm [&_code]:font-mono [&_code]:text-sm"
+            // biome-ignore lint/security/noDangerouslySetInnerHtml: "this is needed."
+            dangerouslySetInnerHTML={{ __html: html }}
+          />
+          <div
+            className="hidden overflow-hidden dark:block [&>pre]:m-0 [&>pre]:bg-background! [&>pre]:p-4 [&>pre]:text-foreground! [&>pre]:text-sm [&_code]:font-mono [&_code]:text-sm"
+            // biome-ignore lint/security/noDangerouslySetInnerHtml: "this is needed."
+            dangerouslySetInnerHTML={{ __html: darkHtml }}
+          />
+          {children && (
+            <div className="absolute top-2 right-2 flex items-center gap-2">
+              {children}
+            </div>
+          )}
+        </div>
+      </div>
+    </CodeBlockContext.Provider>
+  );
+};
+
+export type CodeBlockCopyButtonProps = ComponentProps<typeof Button> & {
+  onCopy?: () => void;
+  onError?: (error: Error) => void;
+  timeout?: number;
+};
+
+export const CodeBlockCopyButton = ({
+  onCopy,
+  onError,
+  timeout = 2000,
+  children,
+  className,
+  ...props
+}: CodeBlockCopyButtonProps) => {
+  const [isCopied, setIsCopied] = useState(false);
+  const { code } = useContext(CodeBlockContext);
+
+  const copyToClipboard = async () => {
+    if (typeof window === "undefined" || !navigator?.clipboard?.writeText) {
+      onError?.(new Error("Clipboard API not available"));
+      return;
+    }
+
+    try {
+      await navigator.clipboard.writeText(code);
+      setIsCopied(true);
+      onCopy?.();
+      setTimeout(() => setIsCopied(false), timeout);
+    } catch (error) {
+      onError?.(error as Error);
+    }
+  };
+
+  const Icon = isCopied ? CheckIcon : CopyIcon;
+
+  return (
+    <Button
+      className={cn("shrink-0", className)}
+      onClick={copyToClipboard}
+      size="icon"
+      variant="ghost"
+      {...props}
+    >
+      {children ?? <Icon size={14} />}
+    </Button>
+  );
+};