diff --git a/backend/openui/litellm.py b/backend/openui/litellm.py index 1dd73b8..618d171 100644 --- a/backend/openui/litellm.py +++ b/backend/openui/litellm.py @@ -31,25 +31,25 @@ def generate_config(): { "model_name": "claude-sonnet-4-0", "litellm_params": { - "model": "claude-sonnet-4-0", + "model": "anthropic/claude-sonnet-4-0", }, }, { "model_name": "claude-opus-4-0", "litellm_params": { - "model": "claude-opus-4-0", + "model": "anthropic/claude-opus-4-0", }, }, { "model_name": "claude-3-7-sonnet", "litellm_params": { - "model": "claude-3-7-sonnet-latest", + "model": "anthropic/claude-3-7-sonnet-latest", }, }, { "model_name": "claude-3-5-haiku", "litellm_params": { - "model": "claude-3-5-haiku-latest", + "model": "anthropic/claude-3-5-haiku-latest", }, }, ] diff --git a/backend/openui/server.py b/backend/openui/server.py index 8733d43..1c28c32 100644 --- a/backend/openui/server.py +++ b/backend/openui/server.py @@ -233,6 +233,11 @@ async def chat_completions( data["max_tokens"] = 4096 - input_tokens - 20 logger.info("Starting trace %s", request.headers.get("X-Wandb-Trace-Id")) with Trace(request.headers.get("X-Wandb-Trace-Id"), user_id): + # TODO: make the frontend remove this? + if "iframeId" in data: + del data["iframeId"] + if "sessionId" in data: + del data["sessionId"] return await model_router.stream_chat_completion(data, input_tokens, user_id) except (ResponseError, APIStatusError) as e: traceback.print_exc() diff --git a/frontend/src/api/openai.ts b/frontend/src/api/openai.ts index 94a15b5..b8f4779 100644 --- a/frontend/src/api/openai.ts +++ b/frontend/src/api/openai.ts @@ -1,3 +1,4 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ import { OpenAI } from 'openai' import type { ToolFinishEvent } from '../state' @@ -8,7 +9,7 @@ function host() { } /* I patched OpenAI here so that users can use basic auth behind a proxy if they want */ class MyOpenAI extends OpenAI { - // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unused-vars + // eslint-disable-next-line @typescript-eslint/no-unused-vars protected override authHeaders(_opts: any) { return {} } @@ -223,6 +224,57 @@ function processToolCalls( return toolCallAccumulator } +function postToolCallsToIframe( + iframeId: string, + toolCalls: Record< + number, + OpenAI.Chat.Completions.ChatCompletionMessageToolCall + > +) { + const iframe = document.getElementById( + `iframe-${iframeId}` + ) as HTMLIFrameElement + const iframeWindow = iframe?.contentWindow + if (!iframeWindow) { + console.error('No iframe found', iframeId) + return + } + for (const toolCall of Object.values(toolCalls)) { + if (toolCall.function.name === 'exec-script') { + const { javascript, description } = JSON.parse( + toolCall.function.arguments + ) + iframeWindow.postMessage( + { + action: 'exec-script', + id: iframeId, + toolCallId: toolCall.id, + javascript, + description + }, + '*' + ) + } else if (toolCall.function.name === 'edit') { + const { mode, selector, html, description, multiple } = JSON.parse( + toolCall.function.arguments + ) + iframeWindow.postMessage( + { + action: 'edit', + id: iframeId, + toolCallId: toolCall.id, + mode, + selector, + html, + description, + multiple + }, + '*' + ) + } + } +} + type Response = { body: string toolCalls: Record< @@ -231,17 +283,150 @@ type Response = { > } +// Helper functions to identify Anthropic tool_use and tool_result messages +function isAnthropicToolUse(msg: any) { + return ( + msg.role === 'assistant' && + Array.isArray(msg.content) && + msg.content.some((c: any) => c.type === 'tool_use') + ) +} +function isAnthropicToolResult(msg: any) { + return ( + msg.role === 'user' && + Array.isArray(msg.content) && + msg.content.some((c: any) => c.type === 'tool_result') + ) +} + export async function respondToToolCalls( - ctx: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming, + ctx: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming & { + iframeId?: string + }, toolCalls: ToolFinishEvent[], - sessionId: string -) { - //ctx.messages.push({}) - await openai.chat.completions.create(ctx, { + sessionId: string, + callback: (response: string) => void +): Promise { + const isAnthropic = + ctx.model && + (ctx.model.includes('claude') || ctx.model.includes('anthropic')) + + // Remove only the most recent contiguous block of assistant/tool_call and tool messages + while (ctx.messages.length > 0) { + const last = ctx.messages[ctx.messages.length - 1] + if ( + last.role === 'assistant' && + ('tool_calls' in last || isAnthropicToolUse(last)) + ) { + ctx.messages.pop() + // Also pop any immediately following tool messages + while ( + ctx.messages.length > 0 && + ctx.messages[ctx.messages.length - 1].role === 'tool' + ) { + ctx.messages.pop() + } + break // Only remove the most recent block + } else if (last.role === 'tool' || isAnthropicToolResult(last)) { + ctx.messages.pop() + } else { + break + } + } + + if (isAnthropic) { + // Anthropic: use tool_use and tool_result blocks + const toolUseBlocks = toolCalls + .filter(tc => tc.call) + .map(tc => ({ + type: 'tool_use', + id: tc.call!.id, + name: tc.call!.function.name, + input: JSON.parse(tc.call!.function.arguments) + })) + if (toolUseBlocks.length > 0) { + ctx.messages.push({ + role: 'assistant', + content: toolUseBlocks as any + }) + const toolResultBlocks = toolCalls + .filter(tc => tc.call && tc.result) + .map(tc => ({ + type: 'tool_result', + tool_use_id: tc.call!.id, + content: + typeof tc.result === 'string' + ? tc.result + : JSON.stringify(tc.result) + })) + if (toolResultBlocks.length > 0) { + ctx.messages.push({ + role: 'user', + content: toolResultBlocks as any + }) + } + } + } else { + // OpenAI: always follow tool_calls with tool messages for each call + const assistantToolCalls = toolCalls + .map(tc => tc.call) + .filter( + (c): c is OpenAI.Chat.Completions.ChatCompletionMessageToolCall => !!c + ) + if (assistantToolCalls.length > 0) { + ctx.messages.push({ + role: 'assistant', + content: '', + tool_calls: assistantToolCalls + }) + const totalToolCalls = assistantToolCalls.length + let calledTools = 0 + // For each tool_call, add a tool message with the correct tool_call_id + for (const tc of toolCalls) { + if (tc.call && tc.result !== undefined) { + ctx.messages.push({ + role: 'tool', + tool_call_id: tc.call.id, + content: + typeof tc.result === 'string' + ? tc.result + : JSON.stringify(tc.result) + }) + calledTools++ + } + } + if (calledTools !== totalToolCalls) { + console.error('Called tools mismatch', calledTools, totalToolCalls) + return { body: '', toolCalls: {} } + } + } + } + + // DEBUG: Output the full message array before making the API call + console.log( + 'DEBUG: ctx.messages before OpenAI API call:', + JSON.stringify(ctx.messages, null, 2) + ) + + ctx.stream = true + const response = await openai.chat.completions.create(ctx, { headers: { 'X-Wandb-Trace-Id': sessionId } }) + let markdown = '' + const finalToolCalls = processToolCalls(undefined, {}) + for await (const chunk of response) { + const part = chunk.choices[0]?.delta?.content ?? '' + markdown += part + callback(part) + processToolCalls(chunk.choices[0]?.delta?.tool_calls, finalToolCalls) + } + const iframeId = ctx.iframeId + if (iframeId) { + postToolCallsToIframe(iframeId, finalToolCalls) + } + return { body: markdown, toolCalls: finalToolCalls } } export async function createOrRefine( @@ -360,20 +545,25 @@ emoji: 🎉 // TODO: use sessionId instead, modify the DOM of jotai dev tools // jotai-devtools-root to include a link to weave console.log('Session ID:', sessionId) - const context: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + const context: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming & { + iframeId: string + sessionId: string + } = { model, messages, temperature, stream: true, max_tokens: GPT4_MAX_TOKENS, - tools + tools, + iframeId, + sessionId: sessionId ?? '' } if (storeContext) { storeContext(context) } const response = await openai.chat.completions.create(context, { headers: { - 'X-Wandb-Trace-Id': iframeId + 'X-Wandb-Trace-Id': sessionId } }) let markdown = '' @@ -393,66 +583,9 @@ emoji: 🎉 }) } console.table(toolTable, ['id', 'name', 'args']) - const iframe = document.getElementById( - `iframe-${iframeId}` - ) as HTMLIFrameElement - let iframeWindow - if (iframe) { - iframeWindow = iframe.contentWindow - } - if (!iframeWindow) { - console.error('No iframe found', iframeId) - return { body: markdown, toolCalls: finalToolCalls } - } - // TODO: move these into UI context - for (const toolCall of Object.values(finalToolCalls)) { - if (toolCall.function.name === 'exec-script') { - const { javascript, description } = JSON.parse( - toolCall.function.arguments - ) - iframeWindow.postMessage( - { - action: 'exec-script', - id: iframeId, - toolCallId: toolCall.id, - javascript, - description - }, - '*' - ) - console.log( - 'Sent exec-script to iframe:', - javascript, - description, - iframeId - ) - } else if (toolCall.function.name === 'edit') { - const { mode, selector, html, description, multiple } = JSON.parse( - toolCall.function.arguments - ) - iframeWindow.postMessage( - { - action: 'edit', - id: iframeId, - toolCallId: toolCall.id, - mode, - selector, - html, - description, - multiple - }, - '*' - ) - console.log( - 'Sent edit DOM to iframe:', - mode, - selector, - html, - description, - multiple, - iframeId - ) - } + console.log('frame', iframeId, finalToolCalls) + if (iframeId) { + postToolCallsToIframe(iframeId, finalToolCalls) } return { body: markdown, diff --git a/frontend/src/components/Chat.tsx b/frontend/src/components/Chat.tsx index 90b154d..e1f3294 100644 --- a/frontend/src/components/Chat.tsx +++ b/frontend/src/components/Chat.tsx @@ -35,7 +35,7 @@ import { uiStateAtom, uiThemeAtom } from 'state' -import { CurrentUIProvider } from './CurrentUiContext' +import CurrentUIProvider from './CurrentUiContext' import ShareDialog from './ShareDialog' import { Button } from './ui/button' diff --git a/frontend/src/components/CodeEditor.tsx b/frontend/src/components/CodeEditor.tsx index 7a5882e..6a80c20 100644 --- a/frontend/src/components/CodeEditor.tsx +++ b/frontend/src/components/CodeEditor.tsx @@ -17,7 +17,7 @@ import { useSaveHistory, type Framework } from 'state' -import CurrentUiContext from './CurrentUiContext' +import { CurrentUIContext } from './CurrentUiContext' import 'monaco-editor/esm/vs/basic-languages/css/css.contribution' import 'monaco-editor/esm/vs/basic-languages/html/html.contribution' @@ -112,7 +112,7 @@ export default function CodeEditor({ const printWidth = 200 const params = useParams() const id = params.id ?? 'new' - const uiContext = useContext(CurrentUiContext) + const uiContext = useContext(CurrentUIContext) const [readOnly, setReadOnly] = useState(framework !== 'html') const editor = useRef() diff --git a/frontend/src/components/CurrentUiContext.tsx b/frontend/src/components/CurrentUiContext.tsx index 4ba9b8a..d44c93d 100644 --- a/frontend/src/components/CurrentUiContext.tsx +++ b/frontend/src/components/CurrentUiContext.tsx @@ -16,13 +16,9 @@ import { import type { ToolEvent, ToolFinishEvent } from 'state' export type { IFrameEvent } from 'state' -const CurrentUIContext = createContext(eventEmitter) +export const CurrentUIContext = createContext(eventEmitter) -export const CurrentUIProvider = ({ - children -}: { - children: React.ReactNode -}) => { +const CurrentUIProvider = ({ children }: { children: React.ReactNode }) => { const { id } = useParams() const [finishedToolCalls, setFinishedToolCalls] = useAtom( finishedToolCallsAtom @@ -50,6 +46,7 @@ export const CurrentUIProvider = ({ prompt: item.prompt(versionIdx) ?? '' } if (update.pureHTML === '') { + console.log('No pureHTML', item.markdown) update.error = `No HTML in LLM response, received: \n${item.markdown}` } setUiState({ ...cleanUiState, ...update }) @@ -105,16 +102,22 @@ export const CurrentUIProvider = ({ }, [id, uiState.toolCalls, setFinishedToolCalls]) useEffect(() => { - const toolCallIds = Object.values(uiState.toolCalls).map( - toolCall => toolCall.id + const toolCallIds = Object.values(uiState.toolCalls).map(toolCall => + String(toolCall.id) ) const finishedToolCallIds = Object.keys(finishedToolCalls) const missingToolCallIds = toolCallIds.filter( id => !finishedToolCallIds.includes(id) ) + console.log('toolCallIds:', toolCallIds) + console.log('finishedToolCallIds:', finishedToolCallIds) + console.log('missingToolCallIds:', missingToolCallIds) if (missingToolCallIds.length > 0) { console.warn('Missing tool call ids', missingToolCallIds) - } else if (Object.keys(finishedToolCalls).length > 0) { + } else if ( + toolCallIds.length > 0 && + finishedToolCallIds.length === toolCallIds.length + ) { eventEmitter.emit(`tool-calls-finished`, finishedToolCalls) } }, [finishedToolCalls, uiState.toolCalls]) @@ -165,4 +168,4 @@ export const CurrentUIProvider = ({ ) } -export default CurrentUIContext +export default CurrentUIProvider diff --git a/frontend/src/components/HtmlAnnotator.tsx b/frontend/src/components/HtmlAnnotator.tsx index 026fd11..5d92b77 100644 --- a/frontend/src/components/HtmlAnnotator.tsx +++ b/frontend/src/components/HtmlAnnotator.tsx @@ -48,7 +48,8 @@ import { } from 'lucide-react' import { useNavigate } from 'react-router-dom' import CodeViewer from './CodeViewer' -import CurrentUIContext, { type IFrameEvent } from './CurrentUiContext' +import type { IFrameEvent } from './CurrentUiContext' +import { CurrentUIContext } from './CurrentUiContext' import { Checkbox } from './ui/checkbox' import { Label } from './ui/label' import { Popover, PopoverContent, PopoverTrigger } from './ui/popover' diff --git a/frontend/src/components/Prompt.tsx b/frontend/src/components/Prompt.tsx index 13c1dcf..2d570f9 100644 --- a/frontend/src/components/Prompt.tsx +++ b/frontend/src/components/Prompt.tsx @@ -1,4 +1,10 @@ -import { convert, createOrRefine, systemPrompt, type Action } from 'api/openai' +import { + convert, + createOrRefine, + respondToToolCalls, + systemPrompt, + type Action +} from 'api/openai' import { Tooltip, TooltipContent, TooltipTrigger } from 'components/ui/tooltip' import { useThrottle, useVersion } from 'hooks' import { useAtom, useAtomValue, useSetAtom } from 'jotai' @@ -34,9 +40,10 @@ import { temperatureAtom, openAIContextAtom, uiStateAtom, - useSaveHistory + useSaveHistory, + type ToolFinishEvent } from 'state' -import CurrentUIContext from './CurrentUiContext' +import { CurrentUIContext } from './CurrentUiContext' import { Button } from './ui/button' import { Textarea } from './ui/textarea' @@ -57,6 +64,7 @@ export default function Prompt({ imageUploadRef: React.RefObject }) { const currentUI = useContext(CurrentUIContext) + const openAIContext = useAtomValue(openAIContextAtom) const setOpenAIContext = useSetAtom(openAIContextAtom) const params = useParams() const [searchParams, setSearchParams] = useSearchParams() @@ -110,6 +118,8 @@ export default function Prompt({ const [animate, setAnimate] = useState(false) const [textareaHeight, setTextareaHeight] = useState() + const [sessionUuid, setSessionUuid] = useState('') + const action: Action = isEditing ? 'refine' : 'create' // Save our streamed markdown const saveMarkdown = useCallback( @@ -129,17 +139,42 @@ export default function Prompt({ // Continue tool calls when they finish useEffect(() => { - function continueToolCalls(finishedToolCalls: unknown) { - console.log('Tool calls finished in Prompt.tsx', finishedToolCalls) + async function continueToolCalls(finishedToolCalls: unknown) { + if (!openAIContext) return + const calls = Object.values( + finishedToolCalls as Record + ) + console.log('Continuing tool calls', calls) + currentUI.emit('ui-state', { rendering: true }) + try { + const response = await respondToToolCalls( + openAIContext, + calls, + sessionUuid, + md => setLiveMarkdown(prev => (prev || '') + md) + ) + setOpenAIContext(openAIContext) + setLiveMarkdown(response.body) + currentUI.emit('ui-state', { + rendering: false, + toolCalls: response.toolCalls + }) + saveMarkdown(response.body) + } catch (error) { + console.error(error) + currentUI.emit('ui-state', { + rendering: false, + error: (error as Error).message + }) + } } currentUI.on('tool-calls-finished', continueToolCalls) return () => { currentUI.off('tool-calls-finished', continueToolCalls) } - }, [currentUI]) + }, [currentUI, openAIContext, sessionUuid, setOpenAIContext, saveMarkdown]) // UUID state and session logic - const [sessionUuid, setSessionUuid] = useState('') const inactivityTimeout = useRef(null) // Helper to get/set uuid in sessionStorage @@ -361,6 +396,7 @@ export default function Prompt({ error: undefined }) } else if (!rendering) { + console.log('No HTML state emitted') currentUI.emit('ui-state', { rendering: false, error: `No HTML in LLM response, received: \n${liveMarkdown}` @@ -600,7 +636,7 @@ export default function Prompt({ className={ /* TODO: make this width calculation dynamic */ cn( - 'my-auto max-h-[130px] flex-1 resize-none items-center justify-center overflow-y-hidden rounded-none align-middle text-lg placeholder:text-lg', + 'my-auto max-h-[130px] flex-1 resize-none items-center justify-center overflow-y-hidden rounded-none align-middle !text-lg placeholder:text-lg', 'bg-muted dark:focus-visible:bg-muted border-none ring-0 outline-hidden transition-all focus-visible:bg-white focus-visible:ring-0 focus-visible:ring-offset-0' ) } diff --git a/frontend/src/components/Settings.tsx b/frontend/src/components/Settings.tsx index c94614b..def1255 100644 --- a/frontend/src/components/Settings.tsx +++ b/frontend/src/components/Settings.tsx @@ -265,7 +265,7 @@ export default function Settings({ trigger }: { trigger: JSX.Element }) { }} onCheckedChange={checked => setModelSupportsImages(checked)} /> -
+
We attempt to detect if the model has vision capabilities. You can override this if you're sure it does. {model === 'gpt-3.5-turbo' && ( @@ -282,7 +282,7 @@ export default function Settings({ trigger }: { trigger: JSX.Element }) { System Prompt