@stainless-api/docs-ai-chat 0.1.0-beta.5 → 0.1.0-beta.50

package/src/hook.ts

@@ -1,31 +1,14 @@
+import type {
+  AssistantTextMessage,
+  AssistantToolCallMessage,
+  ChatMessage,
+  UserMessage,
+} from '@stainless-api/ai-chat/src/types';
+import type { DocsLanguage } from '@stainless-api/docs-ui/routing';
 import { useCallback, useEffect, useReducer, useRef } from 'react';
 import { getChatResponse } from './api';
-import type { DocsLanguage } from '@stainless-api/docs-ui/routing';
 import type { ResponseChunk } from './api/schemas';
 
-// Base
-type BaseMessage = { id: string };
-type BaseTextMessage = BaseMessage & { content: string };
-// User
-type UserMessage = BaseTextMessage & { role: 'user' };
-// Assistant
-type BaseAssistantMessage = BaseMessage & {
-  role: 'assistant';
-  respondingTo: UserMessage['id'];
-  messageType: ResponseChunk['type'];
-};
-type AssistantTextMessage = (BaseAssistantMessage & BaseTextMessage) & {
-  messageType: Extract<BaseAssistantMessage['messageType'], 'text'>;
-  isComplete: boolean;
-};
-type AssistantToolCallMessage = BaseAssistantMessage & {
-  messageType: Extract<BaseAssistantMessage['messageType'], 'tool_use'>;
-  toolName: string;
-  input: Record<string, unknown> | undefined;
-};
-// All
-export type ChatMessage = UserMessage | AssistantTextMessage | AssistantToolCallMessage;
-
 //
 // Reducer
 //
@@ -53,7 +36,10 @@ type ChatReducerAction =
   | { type: 'beginAssistantMessage'; message: Omit<AssistantTextMessage, 'role' | 'isComplete'> }
   | { type: 'streamMessage'; id: string; newContent: string }
   | { type: 'completeMessage'; id: string }
-  | { type: 'addAssistantToolCall'; message: Omit<AssistantToolCallMessage, 'role' | 'id'> };
+  | { type: 'addAssistantToolCall'; message: Omit<AssistantToolCallMessage, 'role' | 'id'> }
+  // a response potentially contains multiple messages / tool calls
+  | { type: 'completeResponse'; respondingTo: UserMessage['id']; spanId: string }
+  | { type: 'addError'; respondingTo: UserMessage['id']; errorMessage: string };
 
 function chatReducer(state: ChatMessage[], action: ChatReducerAction) {
   if (action.type === 'addUserMessage') {
@@ -98,6 +84,26 @@ function chatReducer(state: ChatMessage[], action: ChatReducerAction) {
     });
   }
 
+  if (action.type === 'completeResponse') {
+    return spliceNewMessage(state, {
+      role: 'assistant',
+      id: crypto.randomUUID(),
+      messageType: 'done',
+      respondingTo: action.respondingTo,
+      spanId: action.spanId,
+    } satisfies Extract<ChatMessage, { role: 'assistant' }>);
+  }
+
+  if (action.type === 'addError') {
+    return spliceNewMessage(state, {
+      role: 'assistant',
+      id: crypto.randomUUID(),
+      messageType: 'error',
+      respondingTo: action.respondingTo,
+      errorMessage: action.errorMessage,
+    });
+  }
+
   return state;
 }
 
@@ -105,14 +111,22 @@ function chatReducer(state: ChatMessage[], action: ChatReducerAction) {
 // Consumable hook
 //
 
-export function useChat({ projectId, language }: { projectId: string; language: DocsLanguage }) {
+export function useChat({
+  projectId,
+  language,
+  siteTitle,
+}: {
+  projectId: string;
+  language: DocsLanguage;
+  siteTitle: string | undefined;
+}) {
   // Used to clean up stray streaming requests on unmount (prevent setState on unmounted component)
-  const abortController = useRef(new AbortController());
+  const abortControllerRef = useRef(new AbortController());
   useEffect(() => {
-    abortController.current = abortController.current.signal.aborted
+    abortControllerRef.current = abortControllerRef.current.signal.aborted
       ? new AbortController()
-      : abortController.current;
-    const ac = abortController.current;
+      : abortControllerRef.current;
+    const ac = abortControllerRef.current;
     return () => ac.abort('Component unmounted');
   }, []);
 
@@ -127,66 +141,83 @@ export function useChat({ projectId, language }: { projectId: string; language:
       let currentResponseId = crypto.randomUUID(); // for streaming text messages
       let lastChunkType: ResponseChunk['type'] | undefined = undefined;
 
-      for await (const chunk of getChatResponse(
-        {
-          query: question,
-          project: projectId,
-          language,
-          priorMessages: chatMessages.filter(
-            (msg) => msg.role === 'user' || (msg.role === 'assistant' && msg.messageType === 'text'),
-          ),
-        },
-        abortController.current.signal,
-      )) {
-        if (abortController.current.signal.aborted) break;
-
-        // mark complete when text messages finish streaming
-        if (lastChunkType === 'text' && chunk.type !== 'text') {
-          dispatch({ type: 'completeMessage', id: currentResponseId });
-        }
+      try {
+        let chunk: ResponseChunk | undefined = undefined;
+        for await (chunk of getChatResponse(
+          {
+            query: question,
+            project: projectId,
+            language,
+            priorMessages: chatMessages.filter(
+              (msg) => msg.role === 'user' || (msg.role === 'assistant' && msg.messageType === 'text'),
+            ),
+            siteTitle,
+          },
+          abortControllerRef.current.signal,
+        )) {
+          if (abortControllerRef.current.signal.aborted) break;
+
+          // mark complete when text messages finish streaming
+          if (lastChunkType === 'text' && chunk.type !== 'text') {
+            dispatch({ type: 'completeMessage', id: currentResponseId });
+          }
 
-        // stop reading from the stream on done
-        // TODO: record 'done' in the state so that we can render feedback buttons
-        if (chunk.type === 'done') break;
+          if (chunk.type === 'done') {
+            dispatch({ type: 'completeResponse', respondingTo: userMessageId, spanId: chunk.span_id });
+            // stop reading from the stream on done
+            break;
+          }
+
+          if (chunk.type === 'text') {
+            if (lastChunkType !== 'text') {
+              // start a new text message
+              currentResponseId = crypto.randomUUID();
+              dispatch({
+                type: 'beginAssistantMessage',
+                message: {
+                  content: chunk.text,
+                  id: currentResponseId,
+                  messageType: chunk.type,
+                  respondingTo: userMessageId,
+                },
+              });
+            } else {
+              // continue the current message with the new content
+              dispatch({ type: 'streamMessage', id: currentResponseId, newContent: chunk.text });
+            }
+          }
 
-        if (chunk.type === 'text') {
-          if (lastChunkType !== 'text') {
-            // start a new text message
-            currentResponseId = crypto.randomUUID();
+          if (chunk.type === 'tool_use') {
             dispatch({
-              type: 'beginAssistantMessage',
+              type: 'addAssistantToolCall',
               message: {
-                content: chunk.text,
-                id: currentResponseId,
-                messageType: chunk.type,
                 respondingTo: userMessageId,
+                messageType: chunk.type,
+                toolName: chunk.name,
+                input: chunk.input,
               },
             });
-          } else {
-            // continue the current message with the new content
-            dispatch({ type: 'streamMessage', id: currentResponseId, newContent: chunk.text });
           }
-        }
 
-        if (chunk.type === 'tool_use') {
+          lastChunkType = chunk.type;
+        }
+        if (!chunk || lastChunkType === 'start_session') {
           dispatch({
-            type: 'addAssistantToolCall',
-            message: {
-              respondingTo: userMessageId,
-              messageType: chunk.type,
-              toolName: chunk.name,
-              input: chunk.input,
-            },
+            type: 'addError',
+            respondingTo: userMessageId,
+            errorMessage: 'No response received. Please try again.',
           });
         }
-
-        lastChunkType = chunk.type;
+      } catch {
+        dispatch({
+          type: 'addError',
+          respondingTo: userMessageId,
+          errorMessage: 'Something went wrong. Please try again.',
+        });
       }
     },
-    [language, projectId, chatMessages],
+    [language, projectId, siteTitle, chatMessages],
   );
 
-  // TODO: error handling
-
   return { chatMessages, sendMessage };
 }

package/tsconfig.json

@@ -3,5 +3,5 @@
   "compilerOptions": {
     "jsx": "react-jsx"
   },
-  "include": ["src", "*.d.ts", "plugin.tsx"]
+  "include": ["src", "*.d.ts", "*.config.*", "plugin.tsx"]
 }

package/src/AiChat.module.css

@@ -1,221 +0,0 @@
-.outer-wrapper {
-  display: contents;
-  font-size: var(--stl-typography-scale-sm);
-
-  --shadow-color: light-dark(rgb(0 0 0 / 0.15), var(--stl-color-background));
-  --trigger-size: 3rem;
-  --fixed-inset-bottom: calc(env(safe-area-inset-bottom, 0px) + 1rem);
-  --fixed-inset-right: calc(env(safe-area-inset-right, 0px) + 1rem);
-}
-
-.trigger-border {
-  position: fixed;
-  bottom: var(--fixed-inset-bottom);
-  right: var(--fixed-inset-right);
-  z-index: 51;
-
-  padding: 1px;
-  background-color: var(--stl-color-border);
-  display: block;
-  transition: background-color 0.1s ease;
-
-  &:hover {
-    background-color: var(--stl-color-border-strong);
-  }
-
-  &:has(textarea:focus-visible) {
-    background-color: var(--stl-color-blue-border-strong);
-    transition: background-color 0.25s ease;
-  }
-
-  &:not(:has(.trigger.expanded)) {
-    cursor: pointer;
-  }
-}
-
-.trigger {
-  display: flex;
-  align-items: stretch;
-  min-height: var(--trigger-size);
-  min-width: var(--trigger-size);
-  position: relative;
-
-  background-color: var(--stl-color-background);
-  color: var(--stl-color-foreground);
-  --padding: 0.5rem;
-  padding: var(--padding);
-
-  overflow: clip;
-
-  .bot-icon {
-    position: absolute;
-    --icon-size: 1.5rem;
-    top: calc(var(--trigger-size) / 2 - var(--icon-size) / 2);
-    left: calc(var(--trigger-size) / 2 - var(--icon-size) / 2);
-    width: var(--icon-size);
-    height: var(--icon-size);
-    pointer-events: none;
-  }
-
-  .focused-contents {
-    display: flex;
-    gap: 0.2rem;
-    align-items: center;
-    user-select: none;
-  }
-  &:not(.expanded) .focused-contents {
-    position: absolute;
-    pointer-events: none;
-    textarea {
-      user-select: none;
-    }
-  }
-
-  textarea {
-    display: block;
-    resize: none;
-    width: 24ch;
-    height: auto;
-
-    font-size: var(--stl-typography-scale-base);
-    line-height: 1.35;
-    padding: 0.25em 0.25em 0.25em 0.5em;
-
-    color: inherit;
-    background: none;
-
-    border: none;
-    user-select: initial;
-    &:focus-visible {
-      outline: none;
-    }
-  }
-
-  button[type='submit'] {
-    align-self: flex-end;
-
-    width: calc(var(--trigger-size) - var(--padding) * 2);
-    height: calc(var(--trigger-size) - var(--padding) * 2);
-    display: flex;
-    align-items: center;
-    justify-content: center;
-
-    border: none;
-    border-radius: calc(var(--border-radius) - var(--padding));
-
-    background-color: var(--stl-color-accent-inverse-background);
-    color: var(--stl-color-accent-inverse-foreground);
-
-    cursor: pointer;
-
-    svg {
-      width: 1.15rem;
-      height: 1.15rem;
-      * {
-        stroke-width: 2.25px;
-      }
-    }
-
-    &:disabled {
-      background-color: var(--stl-color-muted-background);
-      color: var(--stl-color-foreground);
-      opacity: 0.6;
-      cursor: default;
-    }
-  }
-}
-
-.chat-area-border {
-  position: fixed;
-  --chat-area-inset-bottom: calc(var(--fixed-inset-bottom) + var(--trigger-size) + 0.5rem);
-  bottom: var(--chat-area-inset-bottom);
-  right: var(--fixed-inset-right);
-  z-index: 50;
-}
-
-.chat-area {
-  display: flex;
-  flex-direction: column-reverse;
-  align-items: center;
-
-  max-width: 48ch;
-  max-height: calc(75svh - var(--chat-area-inset-bottom) - 1rem);
-  overflow-y: auto;
-  overscroll-behavior: contain;
-
-  padding: 1em;
-
-  background-color: var(--stl-color-background);
-  background-image: linear-gradient(
-    to bottom,
-    var(--stl-color-ui-background),
-    var(--stl-color-ui-background)
-  );
-}
-
-.message-log {
-  padding: 0;
-  list-style-type: none;
-
-  flex: 0;
-  position: relative;
-  z-index: 1;
-
-  display: flex;
-  flex-direction: column;
-  gap: 0.5rem;
-
-  width: calc(100vw - 2rem);
-  max-width: 100%;
-
-  &:not(:has(.chat-message)) {
-    display: none;
-  }
-}
-
-.chat-message {
-  list-style-type: none;
-  text-wrap: pretty;
-  font-size: 15px;
-  line-height: 1.4;
-
-  &[data-message-role='user'] {
-    max-width: 40ch;
-    align-self: flex-end;
-    padding: 0.5em 0.75em;
-    background-color: var(--stl-color-accent-inverse-background);
-    color: var(--stl-color-accent-inverse-foreground);
-    overflow: clip;
-
-    ::selection {
-      background-color: rgb(from var(--stl-color-accent-inverse-foreground) r g b / 0.25);
-    }
-  }
-
-  &[data-message-role='assistant'] {
-    align-self: stretch;
-    padding: 0.25em 0;
-    color: var(--stl-color-foreground-reduced);
-  }
-
-  pre {
-    padding: 0.25em 0.65em;
-    font-size: var(--stl-typography-scale-sm);
-    line-height: 1.5;
-  }
-
-  &.tool-use {
-    color: var(--stl-color-foreground-muted);
-    opacity: 0.8;
-
-    strong,
-    em {
-      font-weight: 500;
-      color: var(--stl-color-foreground-reduced);
-    }
-
-    & + .tool-use {
-      margin-top: -0.75em;
-    }
-  }
-}

package/src/AiChat.tsx

@@ -1,84 +0,0 @@
-import { useEffect, useRef, useState } from 'react';
-import { useChat } from './hook';
-import { motion } from 'motion/react';
-
-import type { DocsLanguage } from '@stainless-api/docs-ui/routing';
-
-import styles from './AiChat.module.css';
-import ChatLog from './components/ChatLog';
-import AiChatTrigger from './Trigger';
-
-const borderRadius = 16;
-
-export default function DocsChat({
-  projectId,
-  language,
-}: {
-  projectId: string;
-  language: DocsLanguage | undefined;
-}) {
-  const { chatMessages, sendMessage } = useChat({ projectId, language: language ?? 'typescript' });
-
-  const baseRef = useRef<HTMLDivElement>(null);
-  const inputRef = useRef<HTMLTextAreaElement>(null);
-  const [focused, setFocused] = useState(false);
-  const chatExpanded = focused && chatMessages.length > 0;
-
-  // Manage “focus” state
-  // prettier-ignore
-  useEffect(() => {
-    const ac = new AbortController();
-    // “focus” in/out with click
-    window.addEventListener('click', (e) => {
-      if (!(e.target instanceof Element)) return;
-      if (baseRef.current?.contains(e.target)) { setFocused(true); }
-      else { setFocused(false); }
-    }, { signal: ac.signal });
-    // leave with escape
-    document.addEventListener('keydown', (e) => {
-      if (e.key === 'Escape') {
-        setFocused(false);
-        inputRef.current?.blur(); // this is the one case where the input won’t have already lost focus
-      }
-    }, { signal: ac.signal });
-
-    // record focus state when our chat elements receive focus
-    // unfocus when another element outside of our component gets focus (incl. by keyboard)
-    document.addEventListener('focusin', (e) => {
-      if (!(e.target instanceof HTMLElement) || !baseRef.current) return;
-      setFocused(baseRef.current.contains(e.target) ?? false);
-    }, { signal: ac.signal });
-    return () => ac.abort();
-  }, []);
-
-  return (
-    <div className={styles['outer-wrapper']} ref={baseRef}>
-      <AiChatTrigger
-        focused={focused}
-        updateFocused={setFocused}
-        sendMessage={sendMessage}
-        inputRef={inputRef}
-        borderRadius={borderRadius}
-      />
-
-      {/* TODO: enter and leave animations */}
-      <motion.div
-        layout
-        className={styles['chat-area-border']}
-        style={{
-          display: chatExpanded ? 'block' : 'none', // Activity doesn’t play nice with framer-motion layout animations
-          borderRadius,
-          boxShadow: '0 8px 20px -6px var(--shadow-color)',
-        }}
-      >
-        <motion.div
-          layout
-          className={styles['chat-area']}
-          style={{ borderRadius, boxShadow: 'inset 0 0 0 1px var(--stl-color-border)' }}
-        >
-          <ChatLog messages={chatMessages} />
-        </motion.div>
-      </motion.div>
-    </div>
-  );
-}

package/src/Trigger.tsx

@@ -1,135 +0,0 @@
-import React, { useState } from 'react';
-import { ArrowUpIcon, BotMessageSquareIcon } from 'lucide-react';
-import clsx from 'clsx';
-import { Transition } from 'motion';
-import styles from './AiChat.module.css';
-import { motion } from 'motion/react';
-
-const MotionBotIcon = motion.create(BotMessageSquareIcon);
-
-export default function AiChatTrigger({
-  focused,
-  updateFocused,
-  sendMessage,
-  inputRef,
-  borderRadius,
-}: {
-  focused: boolean;
-  updateFocused: (focused: boolean) => void;
-  sendMessage: (question: string) => void;
-  inputRef: React.RefObject<HTMLTextAreaElement | null>;
-  borderRadius: number;
-}) {
-  const [empty, setEmpty] = useState(true);
-  const [resetKey, setResetKey] = useState(0);
-
-  const layoutTransition = {
-    type: 'spring',
-    mass: 0.7,
-    stiffness: 275,
-    damping: 20,
-  } satisfies Transition;
-
-  const crossBlurTransition = {
-    delay: focused ? 0.07 : 0,
-    duration: 0.1,
-    ease: 'easeInOut',
-  } satisfies Transition;
-
-  return (
-    <form
-      style={{ display: 'contents' }}
-      action={(formData) => {
-        const question = formData.get('question');
-        if (typeof question === 'string' && question.trim().length) {
-          sendMessage(question);
-          setResetKey((k) => k + 1);
-          setEmpty(true);
-        }
-      }}
-    >
-      <motion.label
-        layout
-        transition={layoutTransition}
-        className={styles['trigger-border']}
-        style={{
-          borderRadius: borderRadius + 1,
-          boxShadow: '0 4px 12px -4px var(--shadow-color)',
-        }}
-      >
-        <motion.div
-          layout
-          transition={layoutTransition}
-          className={clsx(styles.trigger, focused && styles.expanded)}
-          style={{ borderRadius: borderRadius, '--border-radius': `${borderRadius}px` }}
-        >
-          {/* Bot icon is visible while closed */}
-          <MotionBotIcon
-            layout
-            className={styles['bot-icon']}
-            animate={{
-              opacity: focused ? 0 : 1,
-              scale: focused ? 0.75 : 1,
-              filter: focused ? 'blur(4px)' : 'blur(0px)',
-            }}
-            style={{ willChange: 'filter, transform' }}
-            transition={crossBlurTransition}
-            aria-label="AI chat"
-          />
-
-          {/* Input & send button are visible while open */}
-          <motion.div
-            layout
-            className={styles['focused-contents']}
-            initial={{ opacity: 0 }}
-            animate={{
-              opacity: focused ? 1 : 0,
-              filter: focused ? 'blur(0px)' : 'blur(4px)',
-            }}
-            style={{ willChange: 'filter, transform' }}
-            transition={crossBlurTransition}
-          >
-            <motion.textarea
-              layout
-              transition={layoutTransition}
-              name="question"
-              rows={1}
-              placeholder="Ask a question"
-              // Keep track of whether the question is submittable
-              onChange={(e: React.ChangeEvent<HTMLTextAreaElement>) => {
-                setEmpty(e.target.value.trim().length === 0);
-              }}
-              // Submit on Cmd+Enter
-              onKeyDown={(e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-                if (e.key === 'Enter' && !e.shiftKey) {
-                  e.preventDefault();
-                  e.currentTarget.form?.requestSubmit();
-                }
-              }}
-              // Update textarea height to fit content as user types
-              ref={(el: HTMLTextAreaElement | null) => {
-                inputRef.current = el;
-                if (!el) return;
-                const updateSize = () => {
-                  el.style.height = 'auto';
-                  el.style.height = `${el.scrollHeight}px`;
-                };
-                const ac = new AbortController();
-                el.addEventListener('input', updateSize, { signal: ac.signal });
-                updateSize();
-                // in case the user focused it before we mounted
-                if (document.activeElement === el) updateFocused(true);
-                return () => ac.abort();
-              }}
-              // make the ref re-mount so we get a fresh height measurement after reset
-              key={resetKey}
-            />
-            <motion.button layout type="submit" disabled={empty} transition={layoutTransition}>
-              <ArrowUpIcon aria-label="Send" />
-            </motion.button>
-          </motion.div>
-        </motion.div>
-      </motion.label>
-    </form>
-  );
-}