@stainless-api/docs 0.1.0-beta.13 → 0.1.0-beta.130

package/stl-docs/chat/hook.ts

@@ -0,0 +1,215 @@
+import type { AssistantTextMessage, AssistantToolCallMessage, ChatMessage, UserMessage } from './ui/types';
+import { useCallback, useEffect, useReducer, useRef } from 'react';
+import type { ResponseChunk } from './schemas';
+import { DocsChatHandler } from './docs-chat-handler';
+
+//
+// Reducer
+//
+
+/**
+ * Splice a new assistant message into a ChatMessage[] stream based on its respondingTo property.
+ *
+ * if two responses are streaming at once, we want to put all the responses to user Message A
+ * before user message B
+ */
+function spliceNewMessage(messages: ChatMessage[], newMessage: Extract<ChatMessage, { role: 'assistant' }>) {
+  // find the most recent assistant message that's responding to the same user message as we are
+  let insertAfterIdx = messages.findLastIndex(
+    (msg) =>
+      (msg.role === 'assistant' && msg.respondingTo === newMessage.respondingTo) ||
+      // if this is the first assistant message responding to this user message
+      (msg.role === 'user' && msg.id === newMessage.respondingTo),
+  );
+  insertAfterIdx = insertAfterIdx === -1 ? messages.length - 1 : insertAfterIdx;
+  return messages.toSpliced(insertAfterIdx + 1, 0, newMessage);
+}
+
+type ChatReducerAction =
+  | { type: 'addUserMessage'; content: UserMessage['content']; id: string }
+  | { type: 'beginAssistantMessage'; message: Omit<AssistantTextMessage, 'role' | 'isComplete'> }
+  | { type: 'streamMessage'; id: string; newContent: string }
+  | { type: 'completeMessage'; id: string }
+  | { 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') {
+    return [...state, { role: 'user', content: action.content, id: action.id } satisfies ChatMessage];
+  }
+
+  if (action.type === 'beginAssistantMessage') {
+    return spliceNewMessage(state, {
+      role: 'assistant',
+      id: action.message.id,
+      respondingTo: action.message.respondingTo,
+      messageType: action.message.messageType satisfies 'text',
+      content: action.message.content,
+      isComplete: false,
+    });
+  }
+
+  if (action.type === 'streamMessage') {
+    return state.map((msg) =>
+      msg.id === action.id && 'content' in msg
+        ? ({ ...msg, content: `${msg.content}${action.newContent}` } satisfies ChatMessage)
+        : msg,
+    );
+  }
+
+  if (action.type === 'completeMessage') {
+    return state.map((msg) =>
+      msg.id === action.id && msg.role === 'assistant' && msg.messageType === 'text'
+        ? ({ ...msg, isComplete: true } satisfies ChatMessage)
+        : msg,
+    );
+  }
+
+  if (action.type === 'addAssistantToolCall') {
+    return spliceNewMessage(state, {
+      role: 'assistant',
+      id: crypto.randomUUID(),
+      messageType: action.message.messageType satisfies 'tool_use',
+      respondingTo: action.message.respondingTo,
+      toolName: action.message.toolName,
+      input: action.message.input,
+    });
+  }
+
+  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;
+}
+
+//
+// Consumable hook
+//
+
+export function useChat({ handler }: { handler: DocsChatHandler }) {
+  // Used to clean up stray streaming requests on unmount (prevent setState on unmounted component)
+  const abortControllerRef = useRef(new AbortController());
+  useEffect(() => {
+    abortControllerRef.current = abortControllerRef.current.signal.aborted
+      ? new AbortController()
+      : abortControllerRef.current;
+    const ac = abortControllerRef.current;
+    return () => ac.abort('Component unmounted');
+  }, []);
+
+  const [chatMessages, dispatch] = useReducer(chatReducer, []);
+
+  /** Send a message and stream back the response in chat */
+  const sendMessage = useCallback(
+    async (question: string) => {
+      const userMessageId = crypto.randomUUID();
+      dispatch({ type: 'addUserMessage', content: question, id: userMessageId });
+
+      let currentResponseId = crypto.randomUUID(); // for streaming text messages
+      let lastChunkType: ResponseChunk['type'] | undefined = undefined;
+
+      try {
+        let chunk: ResponseChunk | undefined = undefined;
+        for await (chunk of handler.generateResponse(
+          {
+            query: question,
+            priorMessages: chatMessages.filter(
+              (msg) => msg.role === 'user' || (msg.role === 'assistant' && msg.messageType === 'text'),
+            ),
+          },
+          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 });
+          }
+
+          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 === 'tool_use') {
+            dispatch({
+              type: 'addAssistantToolCall',
+              message: {
+                respondingTo: userMessageId,
+                messageType: chunk.type,
+                toolName: chunk.name,
+                input: chunk.input,
+              },
+            });
+          }
+
+          lastChunkType = chunk.type;
+        }
+        if (!chunk || lastChunkType === 'start_session') {
+          dispatch({
+            type: 'addError',
+            respondingTo: userMessageId,
+            errorMessage: 'No response received. Please try again.',
+          });
+        }
+      } catch {
+        dispatch({
+          type: 'addError',
+          respondingTo: userMessageId,
+          errorMessage: 'Something went wrong. Please try again.',
+        });
+      }
+    },
+    [chatMessages, handler],
+  );
+
+  async function rateMessage(spanId: string, rating: 'up' | 'down'): Promise<boolean> {
+    try {
+      const { success } = await handler.onRate(spanId, { up: 1 as const, down: 0 as const }[rating]);
+      return success;
+    } catch {
+      return false;
+    }
+  }
+
+  return { chatMessages, sendMessage, rateMessage, setMetadata: handler.onAssignMetadata.bind(handler) };
+}

package/stl-docs/chat/schemas.ts

@@ -0,0 +1,70 @@
+import z from 'zod';
+
+// TODO: replace with generated SDK types instead of copy/pasting from other repo
+export const requestBody = z.object({
+  query: z.string(),
+  sdk: z.object({
+    project: z.string(),
+    language: z.string(),
+    version: z.string().optional(),
+  }),
+  stream: z.boolean().optional(),
+  budget: z
+    .object({
+      maxTokens: z.number().optional(),
+    })
+    .optional(),
+  session_id: z.string().optional(),
+  additionalContext: z
+    .object({
+      prior_messages: z.array(
+        z.object({
+          role: z.enum(['user', 'assistant']),
+          content: z.string(),
+        }),
+      ),
+      code: z.string().optional(),
+      intent: z.string().optional(),
+      lsp: z.string().optional(),
+      errors: z.string().optional(),
+    })
+    .optional(),
+  browser_id: z.string().optional(),
+});
+export type RequestBody = z.input<typeof requestBody>;
+
+export const responseChunk = z.discriminatedUnion('type', [
+  z.object({
+    type: z.literal('text'),
+    text: z.string(),
+  }),
+  z.object({
+    type: z.literal('tool_use'),
+    name: z.string(),
+    input: z.record(z.string(), z.unknown()).optional(),
+  }),
+  z.object({
+    type: z.literal('tool_result'),
+    tool_use_id: z.string(),
+    content: z.string(),
+  }),
+  z.object({
+    type: z.literal('done'),
+    span_id: z.string(),
+  }),
+  z.object({
+    type: z.literal('start_session'),
+    session_id: z.string(),
+  }),
+]);
+export type ResponseChunk = z.infer<typeof responseChunk>;
+
+export const feedbackRequestBody = z.object({ score: z.number().min(0).max(1) });
+export type FeedbackRequestBody = z.infer<typeof feedbackRequestBody>;
+export const feedbackResponseBody = z.object({ success: z.boolean() });
+export type FeedbackResponseBody = z.infer<typeof feedbackResponseBody>;
+
+export const metadataRequestBody = z.object({ metadata: z.record(z.string(), z.string()) });
+export type MetadataRequestBody = z.infer<typeof metadataRequestBody>;
+export const metadataResponseBody = z.object({ success: z.boolean() });
+export type MetadataResponseBody = z.infer<typeof metadataResponseBody>;

package/stl-docs/chat/stainless-handler/index.ts

@@ -0,0 +1,126 @@
+import type { DocsLanguage } from '@stainless-api/docs-ui/routing';
+import { JSONParser } from '@streamparser/json-whatwg';
+import {
+  type RequestBody,
+  responseChunk,
+  type FeedbackRequestBody,
+  feedbackResponseBody,
+  type MetadataRequestBody,
+  metadataResponseBody,
+} from '../schemas';
+import { streamAsyncIterator } from '../stream-util';
+import { DocsChatHandler } from '../docs-chat-handler';
+
+const API_URL = new URL('https://app.stainless.com/api/');
+const CHAT_ENDPOINT = new URL('ai/get-agentic-help', API_URL);
+
+const FEEDBACK_ENDPOINT = (spanId: string) => new URL(`ai/agentic-help/${spanId}/score`, API_URL);
+const METADATA_ENDPOINT = (spanId: string) => new URL(`ai/agentic-help/${spanId}/metadata`, API_URL);
+
+/**
+ * Identifier for tracking unique users in braintrust
+ */
+function getClientId() {
+  let clientId = localStorage.getItem('stainless-client-id');
+  if (!clientId) {
+    clientId = crypto.randomUUID();
+    localStorage.setItem('stainless-client-id', clientId);
+  }
+  return clientId;
+}
+
+/** Context on what the user is currently viewing to pass to the agent */
+function getPageContext({ siteTitle }: { siteTitle: string | undefined }) {
+  const { href } = window.location;
+  const markdownUrl = `${href.replace(/\/$/, '')}/index.md`;
+  const pageTitle = document.querySelector('h1')?.textContent;
+  return [
+    `The user is viewing a documentation page${siteTitle ? ` for ${siteTitle}` : ''}.`,
+    `- Content URL: ${markdownUrl}`,
+    pageTitle && `- Page title: "${pageTitle}"`,
+    // TODO: include stainless path here? does the agent know how to use it?
+    // TODO: pass more of the page content into context without the agent having to retrieve it
+  ]
+    .filter(Boolean)
+    .join('\n');
+}
+
+export class StainlessHandler implements DocsChatHandler {
+  constructor(
+    private language: DocsLanguage,
+    private siteTitle: string | undefined,
+    private project: string,
+  ) {}
+  /**
+   * Stream chat response from the server
+   */
+  async *generateResponse(
+    {
+      query,
+      priorMessages,
+    }: {
+      query: string;
+      priorMessages: NonNullable<RequestBody['additionalContext']>['prior_messages'];
+    },
+    abortSignal: AbortSignal,
+  ) {
+    const res = await fetch(CHAT_ENDPOINT, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({
+        query,
+        sdk: { project: this.project, language: this.language },
+        stream: true,
+        additionalContext: {
+          prior_messages: priorMessages,
+          intent: getPageContext({ siteTitle: this.siteTitle }),
+        },
+        browser_id: getClientId(),
+      } satisfies RequestBody),
+
+      signal: abortSignal,
+    });
+
+    if (!res.ok || !res.body) throw new Error(`Chat request failed with status ${res.status}`);
+
+    const parser = new JSONParser({ separator: '\n', paths: ['$'] });
+    for await (const chunk of streamAsyncIterator(res.body.pipeThrough(parser))) {
+      const chunkParsed = responseChunk.safeParse(chunk.value);
+      if (chunkParsed.success) yield chunkParsed.data;
+    }
+  }
+
+  /**
+   * Attach a score to a response
+   */
+  async onRate(spanId: string, score: 0 | 1) {
+    const res = await fetch(FEEDBACK_ENDPOINT(spanId), {
+      method: 'PUT',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({ score } satisfies FeedbackRequestBody),
+    });
+
+    if (!res.ok) throw new Error(`Feedback request failed with status ${res.status}`);
+    return feedbackResponseBody.parse(await res.json());
+  }
+
+  /**
+   * Attach metadata to a response
+   */
+  async onAssignMetadata(spanId: string, metadata: Record<string, string>) {
+    const res = await fetch(METADATA_ENDPOINT(spanId), {
+      method: 'PUT',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({ metadata } satisfies MetadataRequestBody),
+    });
+
+    if (!res.ok) throw new Error(`Metadata request failed with status ${res.status}`);
+    return metadataResponseBody.parse(await res.json());
+  }
+}

package/stl-docs/chat/stream-util.ts

@@ -0,0 +1,16 @@
+/*
+ * https://jakearchibald.com/2017/async-iterators-and-generators/#making-streams-iterate
+ * safari does not yet support consuming ReadableStream as AsyncIterable
+ */
+export async function* streamAsyncIterator<T>(stream: ReadableStream<T>) {
+  const reader = stream.getReader();
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) return;
+      yield value;
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}