refactor

packages/playwright-core/src/server/trace/viewer/traceViewer.ts

@@ -152,8 +152,10 @@ export async function installRootRedirect(server: HttpServer, traceUrls: string[
     params.append('project', project);
   for (const reporter of options.reporter || [])
     params.append('reporter', reporter);
-  if (options.llm)
-    params.append('llm', 'true');
+  if (process.env.OPENAI_API_KEY)
+    params.append('openai_api_key', process.env.OPENAI_API_KEY);
+  if (process.env.ANTHROPIC_API_KEY)
+    params.append('anthropic_api_key', process.env.ANTHROPIC_API_KEY);
 
   let baseUrl = '.';
   if (process.env.PW_HMR) {

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

@@ -0,0 +1,138 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export type LLMMessage = {
+    role: 'user' | 'assistant' | 'developer';
+    content: string;
+};
+
+export interface LLM {
+    chatCompletion(messages: LLMMessage[], signal: AbortSignal): AsyncGenerator<string>;
+}
+
+async function *parseSSE(body: Response['body']): AsyncGenerator<string> {
+  const reader = body!.getReader();
+  const decoder = new TextDecoder();
+  let buffer = '';
+  while (true) {
+    const { value, done } = await reader.read();
+    if (done)
+      break;
+    buffer += decoder.decode(value, { stream: true });
+    const events = buffer.split('\n\n');
+    buffer = events.pop()!;
+    for (const event of events) {
+      const contentStart = event.indexOf('data: ');
+      if (contentStart === -1)
+        continue;
+      yield event.substring(contentStart + 'data: '.length);
+    }
+  }
+}
+
+export class OpenAI implements LLM {
+
+  constructor(private apiKey: string, private baseURL = 'https://api.openai.com') {}
+
+  async *chatCompletion(messages: LLMMessage[])  {
+    const url = new URL('./v1/chat/completions', this.baseURL);
+    const response = await fetch(url, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'Authorization': `Bearer ${this.apiKey}`,
+      },
+      body: JSON.stringify({
+        model: 'gpt-4o',
+        messages,
+        stream: true,
+      }),
+    });
+
+    if (response.status !== 200 || !response.body)
+      throw new Error('Failed to chat with OpenAI, unexpected status: ' + response.status + await response.text());
+
+    for await (const eventString of parseSSE(response.body)) {
+      if (eventString === '[DONE]')
+        break;
+      const event = JSON.parse(eventString);
+      if (event.object === 'chat.completion.chunk') {
+        if (event.choices[0].finish_reason)
+          break;
+        yield event.choices[0].delta.content;
+      }
+    }
+  }
+}
+
+export class Anthropic implements LLM {
+  constructor(private apiKey: string, private baseURL = 'https://api.anthropic.com') {}
+  async *chatCompletion(messages: LLMMessage[]): AsyncGenerator<string> {
+    const response = await fetch(new URL('./v1/messages', this.baseURL), {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'x-api-key': this.apiKey,
+        'anthropic-version': '2023-06-01'
+      },
+      body: JSON.stringify({
+        model: 'claude-3-5-sonnet-20241022',
+        messages: messages.filter(({ role }) => role !== 'developer'),
+        system: messages.find(({ role }) => role === 'developer')?.content,
+        max_tokens: 1024,
+        stream: true,
+      })
+    });
+
+    if (response.status !== 200 || !response.body)
+      throw new Error('Failed to chat with Anthropic, unexpected status: ' + response.status + await response.text());
+
+    for await (const eventString of parseSSE(response.body)) {
+      const event = JSON.parse(eventString);
+      if (event.type  === 'content_block_delta')
+        yield event.delta.text;
+    }
+  }
+}
+
+export class LLMChat {
+  conversations: Conversation[] = [];
+
+  constructor(readonly api: LLM) {}
+
+  startConversation(systemPrompt: string) {
+    const conversation = new Conversation(this, systemPrompt);
+    this.conversations.push(conversation);
+    return conversation;
+  }
+}
+
+export class Conversation {
+  history: LLMMessage[];
+
+  constructor(private chat: LLMChat, systemPrompt: string) {
+    this.history = [{ role: 'developer', content: systemPrompt }];
+  }
+
+  async *send(message: string, signal: AbortSignal) {
+    const response: LLMMessage = { role: 'assistant', content: '' };
+    this.history.push({ role: 'user', content: message }, response);
+    for await (const chunk of this.chat.api.chatCompletion(this.history, signal)) {
+      response.content += chunk;
+      yield response.content;
+    }
+  }
+}

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

@@ -1,8 +1,9 @@
-import { useCallback, useEffect, useState } from 'react';
+import { useCallback, useEffect, useMemo, 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';
 
 export interface AIState {
   prompt?: string;
@@ -11,62 +12,27 @@ export interface AIState {
 
 export function AITab({ state }: { state?: AIState }) {
   const [input, setInput] = useState('');
-  const [messages, setMessages] = useState<LLMMessage[]>([
-    {
-      role: "developer",
-      content: '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 [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 [abort, setAbort] = useState<AbortController>();
 
   const onSubmit = useCallback(async (event: React.FormEvent<HTMLFormElement>) => {
+    if (!conversation)
+      return;
+
     event.preventDefault();
     setInput('');
     const content = new FormData(event.target as any).get('content') as string;
-    const messages = await new Promise<LLMMessage[]>(resolve => {
-      setMessages(messages => {
-          const newMessages = [...messages, { role: 'user', content } as LLMMessage];
-          resolve(newMessages);
-          return newMessages;
-      })
-    });
+
     const controller = new AbortController();
     setAbort(controller);
 
-    const hydratedMessages = messages.map(message => {
-      let content = message.content;
-      for (const [variable, value] of Object.entries(state?.variables || {})) {
-        content = content.replaceAll(variable, value);
-      }
-      return { ...message, content };
-    })
-
     try {
-      const response = await fetch('./llm/chat-completion', {
-        method: 'POST',
-        headers: {
-          'Content-Type': 'application/json'
-        },
-        body: JSON.stringify(hydratedMessages),
-        signal: controller.signal,
-      });
-      const decoder = new TextDecoder();
-      let reply = '';
-      function update() {
-        setMessages(messages => {
-          return messages.slice(0, -1).concat([{ role: 'assistant', content: reply }]);
-        });
-      }
-      await response.body?.pipeTo(new WritableStream({
-        write(chunk) {
-          reply += decoder.decode(chunk, { stream: true });
-          update();
-        },
-        close() {
-          reply += decoder.decode();
-          update();
-        }
-      }));
+      for await (const _chunk of conversation?.send(content, controller.signal))
+        setMessages([...conversation?.history])
     } finally {
       setAbort(undefined);
     }

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

@@ -0,0 +1,21 @@
+import * as React from 'react';
+import { LLMChat, LLM, OpenAI, Anthropic } from '@isomorphic/llm';
+
+const llmContext = React.createContext<LLMChat | undefined>(undefined);
+
+export function LLMProvider({ openai, anthropic, children }: React.PropsWithChildren<{ openai?: string, anthropic?: string }>) {
+    const chat = React.useMemo(() => {
+        let llm: LLM | undefined;
+        if (openai)
+            llm = new OpenAI(openai);
+        if (anthropic)
+            llm = new Anthropic(anthropic);
+        if (llm)
+            return new LLMChat(llm);
+    }, [openai, anthropic]);
+  return <llmContext.Provider value={chat}>{children}</llmContext.Provider>;
+};
+
+export function useLLMChat() {
+    return React.useContext(llmContext);
+};

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

@@ -38,6 +38,7 @@ import { TraceView } from './uiModeTraceView';
 import { SettingsView } from './settingsView';
 import { DefaultSettingsView } from './defaultSettingsView';
 import { GitCommitInfoProvider } from './errorsTab';
+import { LLMProvider } from './llm';
 
 let xtermSize = { cols: 80, rows: 24 };
 const xtermDataSource: XtermDataSource = {
@@ -61,7 +62,8 @@ const queryParams = {
   updateSnapshots: (searchParams.get('updateSnapshots') as 'all' | 'none' | 'missing' | undefined) || undefined,
   reporters: searchParams.has('reporter') ? searchParams.getAll('reporter') : undefined,
   pathSeparator: searchParams.get('pathSeparator') || '/',
-  llm: searchParams.has('llm'),
+  openai_api_key: searchParams.get('openai_api_key') || undefined,
+  anthropic_api_key: searchParams.get('anthropic_api_key') || undefined,
 };
 if (queryParams.updateSnapshots && !['all', 'none', 'missing'].includes(queryParams.updateSnapshots))
   queryParams.updateSnapshots = undefined;
@@ -399,7 +401,7 @@ export const UIModeView: React.FC<{}> = ({
     });
   }, [closeInstallDialog, testServerConnection]);
 
-  return <div className='vbox ui-mode'>
+  return <LLMProvider openai={queryParams.openai_api_key} anthropic={queryParams.anthropic_api_key}><div className='vbox ui-mode'>
     {!hasBrowsers && <dialog ref={dialogRef}>
       <div className='title'><span className='codicon codicon-lightbulb'></span>Install browsers</div>
       <div className='body'>
@@ -439,7 +441,6 @@ export const UIModeView: React.FC<{}> = ({
               rootDir={testModel?.config?.rootDir}
               revealSource={revealSource}
               onOpenExternally={location => testServerConnection?.openNoReply({ location: { file: location.file, line: location.line, column: location.column } })}
-              llmAvailable={queryParams.llm}
             />
           </GitCommitInfoProvider>
         </div>
@@ -529,5 +530,5 @@ export const UIModeView: React.FC<{}> = ({
       </div>
       }
     />
-  </div>;
+  </div></LLMProvider>;
 };