From 28b75d947939a369999741a4eb7f9e7371c7ce31 Mon Sep 17 00:00:00 2001 From: Simon Knott Date: Wed, 12 Feb 2025 09:46:31 +0100 Subject: [PATCH] refactor --- .../src/server/trace/viewer/traceViewer.ts | 6 +- .../src/utils/isomorphic/llm.ts | 138 ++++++++++++++++++ packages/trace-viewer/src/ui/aiTab.tsx | 60 ++------ packages/trace-viewer/src/ui/llm.tsx | 21 +++ packages/trace-viewer/src/ui/uiModeView.tsx | 9 +- 5 files changed, 181 insertions(+), 53 deletions(-) create mode 100644 packages/playwright-core/src/utils/isomorphic/llm.ts create mode 100644 packages/trace-viewer/src/ui/llm.tsx diff --git a/packages/playwright-core/src/server/trace/viewer/traceViewer.ts b/packages/playwright-core/src/server/trace/viewer/traceViewer.ts index a80d16d113..5f06dd9196 100644 --- a/packages/playwright-core/src/server/trace/viewer/traceViewer.ts +++ b/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) { diff --git a/packages/playwright-core/src/utils/isomorphic/llm.ts b/packages/playwright-core/src/utils/isomorphic/llm.ts new file mode 100644 index 0000000000..157541b6f2 --- /dev/null +++ b/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; +} + +async function *parseSSE(body: Response['body']): AsyncGenerator { + 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 { + 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; + } + } +} diff --git a/packages/trace-viewer/src/ui/aiTab.tsx b/packages/trace-viewer/src/ui/aiTab.tsx index 57cd55be5b..5501b6b02e 100644 --- a/packages/trace-viewer/src/ui/aiTab.tsx +++ b/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([ - { - 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([]); + 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(); const onSubmit = useCallback(async (event: React.FormEvent) => { + if (!conversation) + return; + event.preventDefault(); setInput(''); const content = new FormData(event.target as any).get('content') as string; - const messages = await new Promise(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); } diff --git a/packages/trace-viewer/src/ui/llm.tsx b/packages/trace-viewer/src/ui/llm.tsx new file mode 100644 index 0000000000..98a8801844 --- /dev/null +++ b/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(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 {children}; +}; + +export function useLLMChat() { + return React.useContext(llmContext); +}; diff --git a/packages/trace-viewer/src/ui/uiModeView.tsx b/packages/trace-viewer/src/ui/uiModeView.tsx index 78f8083419..5fb981bd4a 100644 --- a/packages/trace-viewer/src/ui/uiModeView.tsx +++ b/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
+ return
{!hasBrowsers &&
Install browsers
@@ -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} />
@@ -529,5 +530,5 @@ export const UIModeView: React.FC<{}> = ({
} /> -
; + ; };