Add classes for BaseAgent and RootAgent + add Agent Streamer (#2922)

* feat: Add classes for BaseAgent and RootAgent

Author: samschooler

apps/web/client/src/app/api/chat/route.ts

@@ -1,14 +1,11 @@
 import { api } from '@/trpc/server';
 import { trackEvent } from '@/utils/analytics/server';
-import { convertToStreamMessages, getToolSetFromType } from '@onlook/ai';
+import { AgentStreamer, RootAgent } from '@onlook/ai';
 import { toDbMessage } from '@onlook/db';
-import { ChatType, type ChatMessage, type ChatMetadata } from '@onlook/models';
-import { stepCountIs, streamText } from 'ai';
+import { ChatType, type ChatMessage } from '@onlook/models';
 import { type NextRequest } from 'next/server';
 import { v4 as uuidv4 } from 'uuid';
-import { checkMessageLimit, decrementUsage, errorHandler, getModelFromType, getSupabaseUser, getSystemPromptFromType, incrementUsage, repairToolCall } from './helpers';
-
-const MAX_STEPS = 20;
+import { checkMessageLimit, decrementUsage, errorHandler, getSupabaseUser, incrementUsage, repairToolCall } from './helpers';
 
 export async function POST(req: NextRequest) {
     try {
@@ -77,65 +74,41 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
         if (chatType === ChatType.EDIT) {
             usageRecord = await incrementUsage(req, traceId);
         }
-        const modelConfig = await getModelFromType(chatType);
-        const { model, providerOptions, headers } = modelConfig;
-        const systemPrompt = getSystemPromptFromType(chatType);
-        const tools = getToolSetFromType(chatType);
-        const result = streamText({
-            model,
-            headers,
-            tools,
-            stopWhen: stepCountIs(MAX_STEPS),
-            messages: [
-                {
-                    role: 'system',
-                    content: systemPrompt,
-                    providerOptions,
-                },
-                ...convertToStreamMessages(messages),
-            ],
-            experimental_telemetry: {
-                isEnabled: true,
-                metadata: {
-                    conversationId,
-                    projectId,
-                    userId,
-                    chatType: chatType,
-                    tags: ['chat'],
-                    langfuseTraceId: traceId,
-                    sessionId: conversationId,
-                },
-            },
-            experimental_repairToolCall: repairToolCall,
-            onError: async (error) => {
-                console.error('Error in chat stream call', error);
-                // if there was an error with the API, do not penalize the user
-                await decrementUsage(req, usageRecord);
 
-                // Ensure the stream stops on error by re-throwing
-                if (error instanceof Error) {
-                    throw error;
-                } else {
-                    const errorMessage = typeof error === 'string' ? error : JSON.stringify(error);
-                    throw new Error(errorMessage);
-                }
-            }
-        })
+        // Create RootAgent instance
+        const agent = await RootAgent.create(chatType);
+        const streamer = new AgentStreamer(agent, conversationId);
 
-        return result.toUIMessageStreamResponse<ChatMessage>(
-            {
-                originalMessages: messages,
-                generateMessageId: () => uuidv4(),
-                messageMetadata: ({ part }) => {
-                    return {
-                        createdAt: new Date(),
+        return streamer.streamText(messages, {
+            streamTextConfig: {
+                experimental_telemetry: {
+                    isEnabled: true,
+                    metadata: {
                         conversationId,
-                        context: [],
-                        checkpoints: [],
-                        finishReason: part.type === 'finish-step' ? part.finishReason : undefined,
-                        usage: part.type === 'finish-step' ? part.usage : undefined,
-                    } satisfies ChatMetadata;
+                        projectId,
+                        userId,
+                        chatType: chatType,
+                        tags: ['chat'],
+                        langfuseTraceId: traceId,
+                        sessionId: conversationId,
+                    },
                 },
+                experimental_repairToolCall: repairToolCall,
+                onError: async (error) => {
+                    console.error('Error in chat stream call', error);
+                    // if there was an error with the API, do not penalize the user
+                    await decrementUsage(req, usageRecord);
+
+                    // Ensure the stream stops on error by re-throwing
+                    if (error instanceof Error) {
+                        throw error;
+                    } else {
+                        const errorMessage = typeof error === 'string' ? error : JSON.stringify(error);
+                        throw new Error(errorMessage);
+                    }
+                },
+            },
+            toUIMessageStreamResponseConfig: {
                 onFinish: async ({ messages: finalMessages }) => {
                     const messagesToStore = finalMessages
                         .filter(msg =>
@@ -149,8 +122,8 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
                     });
                 },
                 onError: errorHandler,
-            }
-        );
+            },
+        });
     } catch (error) {
         console.error('Error in streamResponse setup', error);
         // If there was an error setting up the stream and we incremented usage, revert it

packages/ai/package.json

@@ -48,6 +48,7 @@
         "marked": "^15.0.7",
         "openai": "^4.103.0",
         "zod": "^4.1.3",
-        "@onlook/ui": "*"
+        "@onlook/ui": "*",
+        "uuid": "^11.1.0"
     }
 }

packages/ai/src/agents/classes/index.ts

@@ -0,0 +1 @@
+export { RootAgent } from './root';

packages/ai/src/agents/classes/root.ts

@@ -0,0 +1,57 @@
+import { ChatType, LLMProvider, OPENROUTER_MODELS, type ModelConfig } from '@onlook/models';
+import { initModel } from '../../chat/providers';
+import { getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt } from '../../prompt';
+import { getToolSetFromType } from '../../tools/toolset';
+import { BaseAgent } from '../models/base';
+
+export class RootAgent extends BaseAgent {
+    readonly id = 'root-agent';
+    private readonly chatType: ChatType;
+    readonly modelConfig: ModelConfig;
+
+    constructor(chatType: ChatType, modelConfig: ModelConfig) {
+        super(getToolSetFromType(chatType));
+
+        this.chatType = chatType;
+        this.modelConfig = modelConfig;
+    }
+
+    get systemPrompt(): string {
+        return this.getSystemPromptFromType(this.chatType);
+    }
+
+    private getSystemPromptFromType(chatType: ChatType): string {
+        switch (chatType) {
+            case ChatType.CREATE:
+                return getCreatePageSystemPrompt();
+            case ChatType.ASK:
+                return getAskModeSystemPrompt();
+            case ChatType.EDIT:
+            default:
+                return getSystemPrompt();
+        }
+    }
+
+    static async create(chatType: ChatType): Promise<RootAgent> {
+        const modelConfig = await RootAgent.getModelFromType(chatType);
+        return new RootAgent(chatType, modelConfig);
+    }
+
+    private static async getModelFromType(chatType: ChatType): Promise<ModelConfig> {
+        switch (chatType) {
+            case ChatType.CREATE:
+            case ChatType.FIX:
+                return await initModel({
+                    provider: LLMProvider.OPENROUTER,
+                    model: OPENROUTER_MODELS.OPEN_AI_GPT_5,
+                });
+            case ChatType.ASK:
+            case ChatType.EDIT:
+            default:
+                return await initModel({
+                    provider: LLMProvider.OPENROUTER,
+                    model: OPENROUTER_MODELS.CLAUDE_4_SONNET,
+                });
+        }
+    }
+}

packages/ai/src/agents/index.ts

@@ -0,0 +1,3 @@
+export * from './models';
+export * from './classes';
+export { AgentStreamer } from './streamer';

packages/ai/src/agents/models/base.ts

@@ -0,0 +1,52 @@
+import type { ModelConfig } from '@onlook/models';
+import { stepCountIs, streamText, type ModelMessage, type ToolSet } from 'ai';
+
+export abstract class BaseAgent {
+    abstract readonly id: string;
+    abstract readonly modelConfig: ModelConfig;
+    abstract readonly systemPrompt: string;
+    protected readonly toolSet: ToolSet
+
+    constructor(toolSet: ToolSet = {}) {
+        this.toolSet = toolSet;
+    }
+
+    getSystemPrompt() {
+        return this.systemPrompt;
+    }
+
+    getToolSet() {
+        return this.toolSet;
+    }
+
+    // Default streamText configuration - can be overridden by subclasses
+    protected getStreamTextConfig() {
+        return {
+            maxSteps: 20,
+        };
+    }
+
+    // Main streamText method that uses the agent's configuration
+    streamText(
+        messages: ModelMessage[] = [],
+        additionalConfig: Omit<Partial<Parameters<typeof streamText>[0]>, 'model' | 'headers' | 'tools' | 'stopWhen' | 'messages' | 'prompt'> = {},
+    ) {
+        const config = this.getStreamTextConfig();
+
+        return streamText({
+            model: this.modelConfig.model,
+            headers: this.modelConfig.headers,
+            tools: this.getToolSet(),
+            stopWhen: stepCountIs(config.maxSteps),
+            messages: [
+                {
+                    role: 'system',
+                    content: this.getSystemPrompt(),
+                    providerOptions: this.modelConfig.providerOptions,
+                },
+                ...messages
+            ],
+            ...additionalConfig,
+        })
+    }
+}

packages/ai/src/agents/models/index.ts

@@ -0,0 +1 @@
+export { BaseAgent } from './base';

packages/ai/src/agents/streamer.ts

@@ -0,0 +1,37 @@
+import type { ChatMessage, ChatMetadata } from "@onlook/models";
+import type { BaseAgent } from "./models";
+import { convertToStreamMessages } from "../stream";
+import { v4 as uuidv4 } from 'uuid';
+import type { streamText, UIMessageStreamOnFinishCallback, UIMessageStreamOptions } from "ai";
+
+export class AgentStreamer {
+    private readonly agent: BaseAgent;
+    private readonly conversationId: string;
+
+    constructor(agent: BaseAgent, conversationId: string) {
+        this.agent = agent;
+        this.conversationId = conversationId;
+    }
+
+    async streamText(messages: ChatMessage[], { streamTextConfig, toUIMessageStreamResponseConfig }: { streamTextConfig: Omit<Partial<Parameters<typeof streamText>[0]>, 'model' | 'headers' | 'tools' | 'stopWhen' | 'messages' | 'prompt'>, toUIMessageStreamResponseConfig: UIMessageStreamOptions<ChatMessage> }) {
+        const conversationId = this.conversationId;
+
+        const result = await this.agent.streamText(convertToStreamMessages(messages), streamTextConfig);
+
+        return result.toUIMessageStreamResponse<ChatMessage>({
+            originalMessages: messages,
+            generateMessageId: () => uuidv4(),
+            messageMetadata: ({ part }) => {
+                return {
+                    createdAt: new Date(),
+                    conversationId,
+                    context: [],
+                    checkpoints: [],
+                    finishReason: part.type === 'finish-step' ? part.finishReason : undefined,
+                    usage: part.type === 'finish-step' ? part.usage : undefined,
+                } satisfies ChatMetadata;
+            },
+            ...toUIMessageStreamResponseConfig,
+        });
+    }
+}

packages/ai/src/index.ts

@@ -1,3 +1,4 @@
+export * from './agents';
 export * from './apply';
 export * from './chat';
 export * from './prompt';