diff --git a/packages/playwright-core/src/utils/isomorphic/llm.ts b/packages/playwright-core/src/utils/isomorphic/llm.ts index 401bfe7748..cd5651cf18 100644 --- a/packages/playwright-core/src/utils/isomorphic/llm.ts +++ b/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; + chatCompletion(messages: LLMMessage[], signal: AbortSignal): AsyncGenerator; } async function *parseSSE(body: Response['body']): AsyncGenerator { @@ -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(); 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(); 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(); } } } diff --git a/packages/trace-viewer/src/sw/main.ts b/packages/trace-viewer/src/sw/main.ts index d965cabc9d..cf6f19f087 100644 --- a/packages/trace-viewer/src/sw/main.ts +++ b/packages/trace-viewer/src/sw/main.ts @@ -88,10 +88,8 @@ async function doFetch(event: FetchEvent): Promise { 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); diff --git a/packages/trace-viewer/src/ui/aiTab.tsx b/packages/trace-viewer/src/ui/aiTab.tsx index 5501b6b02e..d53541d98e 100644 --- a/packages/trace-viewer/src/ui/aiTab.tsx +++ b/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([]); - 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(); const onSubmit = useCallback(async (event: React.FormEvent) => { @@ -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 (
- {messages.filter(({ role }) => role !== 'developer').map((message, index) => ( + {history.filter(({ role }) => role !== 'developer').map((message, index) => (
)}
- {message.content} + {message.displayContent ?? message.content}
))} diff --git a/packages/trace-viewer/src/ui/llm.tsx b/packages/trace-viewer/src/ui/llm.tsx index 98a8801844..9bfced6008 100644 --- a/packages/trace-viewer/src/ui/llm.tsx +++ b/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; +}; diff --git a/packages/trace-viewer/src/ui/uiModeTraceView.tsx b/packages/trace-viewer/src/ui/uiModeTraceView.tsx index 41aec33621..cf35d89007 100644 --- a/packages/trace-viewer/src/ui/uiModeTraceView.tsx +++ b/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(null); @@ -98,7 +97,6 @@ export const TraceView: React.FC<{ annotations={item.testCase?.annotations || []} onOpenExternally={onOpenExternally} revealSource={revealSource} - llmAvailable={llmAvailable} />; }; diff --git a/packages/trace-viewer/src/ui/workbench.tsx b/packages/trace-viewer/src/ui/workbench.tsx index a92c417750..281e4e0060 100644 --- a/packages/trace-viewer/src/ui/workbench.tsx +++ b/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(undefined); const [revealedError, setRevealedError] = React.useState(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(); const [selectedTime, setSelectedTime] = React.useState(); 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);