ai-branch

1'use client';
2 
3import { ChevronLeft, ChevronRight } from 'lucide-react';
4import type { HTMLAttributes, ReactElement, ReactNode } from 'react';
5import { createContext, useContext, useEffect, useState } from 'react';
6import { Button } from '@repo/shadcn-ui/components/ui/button';
7import { cn } from '@repo/shadcn-ui/lib/utils';
8 
9type AIBranchContextType = {
10  currentBranch: number;
11  totalBranches: number;
12  goToPrevious: () => void;
13  goToNext: () => void;
14  branches: ReactElement[];
15  setBranches: (branches: ReactElement[]) => void;
16};
17 
18const AIBranchContext = createContext<AIBranchContextType | null>(null);
19 
20const useAIBranch = () => {
21  const context = useContext(AIBranchContext);
22 
23  if (!context) {
24    throw new Error('AIBranch components must be used within AIBranch');
25  }
26 
27  return context;
28};
29 
30export type AIBranchProps = HTMLAttributes<HTMLDivElement> & {
31  defaultBranch?: number;
32  onBranchChange?: (branchIndex: number) => void;
33};
34 
35export const AIBranch = ({
36  defaultBranch = 0,
37  onBranchChange,
38  className,
39  ...props
40}: AIBranchProps) => {
41  const [currentBranch, setCurrentBranch] = useState(defaultBranch);
42  const [branches, setBranches] = useState<ReactElement[]>([]);
43 
44  const handleBranchChange = (newBranch: number) => {
45    setCurrentBranch(newBranch);
46    onBranchChange?.(newBranch);
47  };
48 
49  const goToPrevious = () => {
50    const newBranch =
51      currentBranch > 0 ? currentBranch - 1 : branches.length - 1;
52    handleBranchChange(newBranch);
53  };
54 
55  const goToNext = () => {
56    const newBranch =
57      currentBranch < branches.length - 1 ? currentBranch + 1 : 0;
58    handleBranchChange(newBranch);
59  };
60 
61  const contextValue: AIBranchContextType = {
62    currentBranch,
63    totalBranches: branches.length,
64    goToPrevious,
65    goToNext,
66    branches,
67    setBranches,
68  };
69 
70  return (
71    <AIBranchContext.Provider value={contextValue}>
72      <div
73        className={cn('grid w-full gap-2 [&>div]:pb-0', className)}
74        {...props}
75      />
76    </AIBranchContext.Provider>
77  );
78};
79 
80export type AIBranchMessagesProps = {
81  children: ReactElement | ReactElement[];
82};
83 
84export const AIBranchMessages = ({ children }: AIBranchMessagesProps) => {
85  const { currentBranch, setBranches, branches } = useAIBranch();
86  const childrenArray = Array.isArray(children) ? children : [children];
87 
88  // Use useEffect to update branches when they change
89  useEffect(() => {
90    if (branches.length !== childrenArray.length) {
91      setBranches(childrenArray);
92    }
93// … truncated