chain-of-thought.tsx 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
  1. 'use client';
  2. import { useControllableState } from '@radix-ui/react-use-controllable-state';
  3. import { Badge } from '@/components/ui/badge';
  4. import { Collapsible, CollapsibleContent, CollapsibleTrigger } from '@/components/ui/collapsible';
  5. import { cn } from '@/lib/utils';
  6. import { BrainIcon, ChevronDownIcon, DotIcon, type LucideIcon } from 'lucide-react';
  7. import type { ComponentProps, ReactNode } from 'react';
  8. import { createContext, memo, useContext, useMemo } from 'react';
  9. type ChainOfThoughtContextValue = {
  10. isOpen: boolean;
  11. setIsOpen: (open: boolean) => void;
  12. };
  13. const ChainOfThoughtContext = createContext<ChainOfThoughtContextValue | null>(null);
  14. const useChainOfThought = () => {
  15. const context = useContext(ChainOfThoughtContext);
  16. if (!context) {
  17. throw new Error('ChainOfThought components must be used within ChainOfThought');
  18. }
  19. return context;
  20. };
  21. export type ChainOfThoughtProps = ComponentProps<'div'> & {
  22. open?: boolean;
  23. defaultOpen?: boolean;
  24. onOpenChange?: (open: boolean) => void;
  25. };
  26. export const ChainOfThought = memo(
  27. ({
  28. className,
  29. open,
  30. defaultOpen = false,
  31. onOpenChange,
  32. children,
  33. ...props
  34. }: ChainOfThoughtProps) => {
  35. const [isOpen, setIsOpen] = useControllableState({
  36. prop: open,
  37. defaultProp: defaultOpen,
  38. onChange: onOpenChange,
  39. });
  40. const chainOfThoughtContext = useMemo(() => ({ isOpen, setIsOpen }), [isOpen, setIsOpen]);
  41. return (
  42. <ChainOfThoughtContext.Provider value={chainOfThoughtContext}>
  43. <div className={cn('not-prose max-w-prose space-y-4', className)} {...props}>
  44. {children}
  45. </div>
  46. </ChainOfThoughtContext.Provider>
  47. );
  48. },
  49. );
  50. export type ChainOfThoughtHeaderProps = ComponentProps<typeof CollapsibleTrigger>;
  51. export const ChainOfThoughtHeader = memo(
  52. ({ className, children, ...props }: ChainOfThoughtHeaderProps) => {
  53. const { isOpen, setIsOpen } = useChainOfThought();
  54. return (
  55. <Collapsible onOpenChange={setIsOpen} open={isOpen}>
  56. <CollapsibleTrigger
  57. className={cn(
  58. 'flex w-full items-center gap-2 text-muted-foreground text-sm transition-colors hover:text-foreground',
  59. className,
  60. )}
  61. {...props}
  62. >
  63. <BrainIcon className="size-4" />
  64. <span className="flex-1 text-left">{children ?? 'Chain of Thought'}</span>
  65. <ChevronDownIcon
  66. className={cn('size-4 transition-transform', isOpen ? 'rotate-180' : 'rotate-0')}
  67. />
  68. </CollapsibleTrigger>
  69. </Collapsible>
  70. );
  71. },
  72. );
  73. export type ChainOfThoughtStepProps = ComponentProps<'div'> & {
  74. icon?: LucideIcon;
  75. label: ReactNode;
  76. description?: ReactNode;
  77. status?: 'complete' | 'active' | 'pending';
  78. };
  79. export const ChainOfThoughtStep = memo(
  80. ({
  81. className,
  82. icon: Icon = DotIcon,
  83. label,
  84. description,
  85. status = 'complete',
  86. children,
  87. ...props
  88. }: ChainOfThoughtStepProps) => {
  89. const statusStyles = {
  90. complete: 'text-muted-foreground',
  91. active: 'text-foreground',
  92. pending: 'text-muted-foreground/50',
  93. };
  94. return (
  95. <div
  96. className={cn(
  97. 'flex gap-2 text-sm',
  98. statusStyles[status],
  99. 'fade-in-0 slide-in-from-top-2 animate-in',
  100. className,
  101. )}
  102. {...props}
  103. >
  104. <div className="relative mt-0.5">
  105. <Icon className="size-4" />
  106. <div className="-mx-px absolute top-7 bottom-0 left-1/2 w-px bg-border" />
  107. </div>
  108. <div className="flex-1 space-y-2 overflow-hidden">
  109. <div>{label}</div>
  110. {description && <div className="text-muted-foreground text-xs">{description}</div>}
  111. {children}
  112. </div>
  113. </div>
  114. );
  115. },
  116. );
  117. export type ChainOfThoughtSearchResultsProps = ComponentProps<'div'>;
  118. export const ChainOfThoughtSearchResults = memo(
  119. ({ className, ...props }: ChainOfThoughtSearchResultsProps) => (
  120. <div className={cn('flex flex-wrap items-center gap-2', className)} {...props} />
  121. ),
  122. );
  123. export type ChainOfThoughtSearchResultProps = ComponentProps<typeof Badge>;
  124. export const ChainOfThoughtSearchResult = memo(
  125. ({ className, children, ...props }: ChainOfThoughtSearchResultProps) => (
  126. <Badge
  127. className={cn('gap-1 px-2 py-0.5 font-normal text-xs', className)}
  128. variant="secondary"
  129. {...props}
  130. >
  131. {children}
  132. </Badge>
  133. ),
  134. );
  135. export type ChainOfThoughtContentProps = ComponentProps<typeof CollapsibleContent>;
  136. export const ChainOfThoughtContent = memo(
  137. ({ className, children, ...props }: ChainOfThoughtContentProps) => {
  138. const { isOpen } = useChainOfThought();
  139. return (
  140. <Collapsible open={isOpen}>
  141. <CollapsibleContent
  142. className={cn(
  143. 'mt-2 space-y-3',
  144. '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',
  145. className,
  146. )}
  147. {...props}
  148. >
  149. {children}
  150. </CollapsibleContent>
  151. </Collapsible>
  152. );
  153. },
  154. );
  155. export type ChainOfThoughtImageProps = ComponentProps<'div'> & {
  156. caption?: string;
  157. };
  158. export const ChainOfThoughtImage = memo(
  159. ({ className, children, caption, ...props }: ChainOfThoughtImageProps) => (
  160. <div className={cn('mt-2 space-y-2', className)} {...props}>
  161. <div className="relative flex max-h-[22rem] items-center justify-center overflow-hidden rounded-lg bg-muted p-3">
  162. {children}
  163. </div>
  164. {caption && <p className="text-muted-foreground text-xs">{caption}</p>}
  165. </div>
  166. ),
  167. );
  168. ChainOfThought.displayName = 'ChainOfThought';
  169. ChainOfThoughtHeader.displayName = 'ChainOfThoughtHeader';
  170. ChainOfThoughtStep.displayName = 'ChainOfThoughtStep';
  171. ChainOfThoughtSearchResults.displayName = 'ChainOfThoughtSearchResults';
  172. ChainOfThoughtSearchResult.displayName = 'ChainOfThoughtSearchResult';
  173. ChainOfThoughtContent.displayName = 'ChainOfThoughtContent';
  174. ChainOfThoughtImage.displayName = 'ChainOfThoughtImage';