rebuild to use conversation helper

packages/playwright-core/src/utils/isomorphic/llm.ts

@@ -14,13 +14,16 @@
  * limitations under the License.
  */
 
+import { EventEmitter } from '@testIsomorphic/events';
+
 export type LLMMessage = {
-    role: 'user' | 'assistant' | 'developer';
-    content: string;
+  role: 'user' | 'assistant' | 'developer';
+  content: string;
+  displayContent?: string;
 };
 
 export interface LLM {
-    chatCompletion(messages: LLMMessage[], signal: AbortSignal): AsyncGenerator<string>;
+  chatCompletion(messages: LLMMessage[], signal: AbortSignal): AsyncGenerator<string>;
 }
 
 async function *parseSSE(body: Response['body']): AsyncGenerator<string> {
@@ -58,7 +61,7 @@ export class OpenAI implements LLM {
       },
       body: JSON.stringify({
         model: 'gpt-4o',
-        messages,
+        messages: messages.map(({ role, content }) => ({ role, content })),
         stream: true,
       }),
     });
@@ -92,7 +95,7 @@ export class Anthropic implements LLM {
       },
       body: JSON.stringify({
         model: 'claude-3-5-sonnet-20241022',
-        messages: messages.filter(({ role }) => role !== 'developer'),
+        messages: messages.filter(({ role }) => role !== 'developer').map(({ role, content }) => ({ role, content })),
         system: messages.find(({ role }) => role === 'developer')?.content,
         max_tokens: 1024,
         stream: true,
@@ -111,30 +114,33 @@ export class Anthropic implements LLM {
 }
 
 export class LLMChat {
-  conversations: Conversation[] = [];
+  conversations = new Map<string, Conversation>();
 
   constructor(readonly api: LLM) {}
 
-  startConversation(systemPrompt: string) {
-    const conversation = new Conversation(this, systemPrompt);
-    this.conversations.push(conversation);
-    return conversation;
+  getConversation(id: string, systemPrompt: string) {
+    if (!this.conversations.has(id)) {
+      const conversation = new Conversation(this, systemPrompt);
+      this.conversations.set(id, conversation);
+    }
+    return this.conversations.get(id)!;
   }
 }
 
 export class Conversation {
   history: LLMMessage[];
+  onChange = new EventEmitter<void>();
 
   constructor(private chat: LLMChat, systemPrompt: string) {
     this.history = [{ role: 'developer', content: systemPrompt }];
   }
 
-  async *send(message: string, signal: AbortSignal) {
+  async send(content: string, displayContent: string | undefined, signal: AbortSignal) {
     const response: LLMMessage = { role: 'assistant', content: '' };
-    this.history.push({ role: 'user', content: message }, response);
+    this.history.push({ role: 'user', content, displayContent }, response);
     for await (const chunk of this.chat.api.chatCompletion(this.history, signal)) {
       response.content += chunk;
-      yield response.content;
+      this.onChange.fire();
     }
   }
 }

packages/trace-viewer/src/sw/main.ts

@@ -88,10 +88,8 @@ async function doFetch(event: FetchEvent): Promise<Response> {
   if (event.request.url.startsWith('chrome-extension://'))
     return fetch(event.request);
 
-  if (event.request.headers.get('x-pw-serviceworker') === 'forward') {
-    event.request.headers.delete('x-pw-serviceworker');
+  if (event.request.headers.get('x-pw-serviceworker') === 'forward')
     return fetch(event.request);
-  }
 
   const request = event.request;
   const client = await self.clients.get(event.clientId);

packages/trace-viewer/src/ui/aiTab.tsx

@@ -1,9 +1,8 @@
-import { useCallback, useEffect, useMemo, useState } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 import Markdown from 'react-markdown'
 import './aiTab.css';
-import type { LLMMessage } from 'playwright-core/lib/server/llm';
 import { clsx } from '@web/uiUtils';
-import { useLLMChat } from './llm';
+import { useLLMConversation } from './llm';
 
 export interface AIState {
   prompt?: string;
@@ -12,11 +11,7 @@ export interface AIState {
 
 export function AITab({ state }: { state?: AIState }) {
   const [input, setInput] = useState('');
-
-  const [messages, setMessages] = useState<LLMMessage[]>([]);
-  const chat = useLLMChat();
-  const conversation = useMemo(() => chat?.startConversation('You are a helpful assistant, skilled in programming and software testing with Playwright. Help me write good code. Be bold, creative and assertive when you suggest solutions.'), [chat]);
-
+  const [history, conversation] = useLLMConversation('aitab', 'You are a helpful assistant, skilled in programming and software testing with Playwright. Help me write good code. Be bold, creative and assertive when you suggest solutions.');
   const [abort, setAbort] = useState<AbortController>();
 
   const onSubmit = useCallback(async (event: React.FormEvent<HTMLFormElement>) => {
@@ -28,11 +23,9 @@ export function AITab({ state }: { state?: AIState }) {
     const content = new FormData(event.target as any).get('content') as string;
 
     const controller = new AbortController();
-    setAbort(controller);
-
-    try {
-      for await (const _chunk of conversation?.send(content, controller.signal))
-        setMessages([...conversation?.history])
+    try { 
+      setAbort(controller);
+      await conversation.send(content, undefined, controller.signal);
     } finally {
       setAbort(undefined);
     }
@@ -46,7 +39,7 @@ export function AITab({ state }: { state?: AIState }) {
   return (
     <div className="chat-container">
       <div className="messages-container">
-        {messages.filter(({ role }) => role !== 'developer').map((message, index) => (
+        {history.filter(({ role }) => role !== 'developer').map((message, index) => (
           <div
             key={'' + index}
             className={clsx('message', message.role === 'user' && 'user-message')}
@@ -57,7 +50,7 @@ export function AITab({ state }: { state?: AIState }) {
               </div>
             )}
             <div className="message-content">
-              <Markdown>{message.content}</Markdown>
+              <Markdown>{message.displayContent ?? message.content}</Markdown>
             </div>
           </div>
         ))}

packages/trace-viewer/src/ui/llm.tsx

@@ -19,3 +19,21 @@ export function LLMProvider({ openai, anthropic, children }: React.PropsWithChil
 export function useLLMChat() {
     return React.useContext(llmContext);
 };
+
+export function useLLMConversation(id: string, systemPrompt: string) {
+    const chat = useLLMChat();
+    if (!chat)
+        throw new Error('No LLM chat available, make sure theres a LLMProvider above');
+    const conversation = React.useMemo(() => chat.getConversation(id, systemPrompt), [chat, id]);
+    const [history, setHistory] = React.useState(conversation.history);
+    React.useEffect(() => {
+        function update() {
+            setHistory([...conversation.history]);
+        }
+        update();
+        const subscription = conversation.onChange.event(update);
+        return subscription.dispose;
+    }, [conversation]);
+
+    return [history, conversation] as const;
+};

packages/trace-viewer/src/ui/uiModeTraceView.tsx

@@ -31,8 +31,7 @@ export const TraceView: React.FC<{
   onOpenExternally?: (location: SourceLocation) => void,
   revealSource?: boolean,
   pathSeparator: string,
-  llmAvailable?: boolean,
-}> = ({ item, rootDir, onOpenExternally, revealSource, pathSeparator, llmAvailable }) => {
+}> = ({ item, rootDir, onOpenExternally, revealSource, pathSeparator }) => {
   const [model, setModel] = React.useState<{ model: MultiTraceModel, isLive: boolean } | undefined>();
   const [counter, setCounter] = React.useState(0);
   const pollTimer = React.useRef<NodeJS.Timeout | null>(null);
@@ -98,7 +97,6 @@ export const TraceView: React.FC<{
     annotations={item.testCase?.annotations || []}
     onOpenExternally={onOpenExternally}
     revealSource={revealSource}
-    llmAvailable={llmAvailable}
   />;
 };
 

packages/trace-viewer/src/ui/workbench.tsx

@@ -44,6 +44,7 @@ import type { UITestStatus } from './testUtils';
 import type { AfterActionTraceEventAttachment } from '@trace/trace';
 import type { HighlightedElement } from './snapshotTab';
 import { AIState, AITab } from './aiTab';
+import { useLLMChat } from './llm';
 
 export const Workbench: React.FunctionComponent<{
   model?: modelUtil.MultiTraceModel,
@@ -57,8 +58,7 @@ export const Workbench: React.FunctionComponent<{
   inert?: boolean,
   onOpenExternally?: (location: modelUtil.SourceLocation) => void,
   revealSource?: boolean,
-  llmAvailable?: boolean,
-}> = ({ model, showSourcesFirst, rootDir, fallbackLocation, isLive, hideTimeline, status, annotations, inert, onOpenExternally, revealSource, llmAvailable }) => {
+}> = ({ model, showSourcesFirst, rootDir, fallbackLocation, isLive, hideTimeline, status, annotations, inert, onOpenExternally, revealSource }) => {
   const [selectedCallId, setSelectedCallId] = React.useState<string | undefined>(undefined);
   const [revealedError, setRevealedError] = React.useState<ErrorDescription | undefined>(undefined);
   const [revealedAttachment, setRevealedAttachment] = React.useState<[attachment: AfterActionTraceEventAttachment, renderCounter: number] | undefined>(undefined);
@@ -72,7 +72,8 @@ export const Workbench: React.FunctionComponent<{
   const [aiState, setAIState] = React.useState<AIState>();
   const [selectedTime, setSelectedTime] = React.useState<Boundaries | undefined>();
   const [sidebarLocation, setSidebarLocation] = useSetting<'bottom' | 'right'>('propertiesSidebarLocation', 'bottom');
-  const [showAITab, setShowAITab] = useSetting<'on' | 'off'>('aiTab', 'off')
+  const [showAITab, setShowAITab] = useSetting<'on' | 'off'>('aiTab', 'off');
+  const llmAvailable = !!useLLMChat();
 
   const setSelectedAction = React.useCallback((action: modelUtil.ActionTraceEventInContext | undefined) => {
     setSelectedCallId(action?.callId);