feat: Update AG-UI protocol with Generative UI support

sdks/python/ag_ui/core/__init__.py

@@ -53,6 +53,8 @@
     TextInputContent,
     BinaryInputContent,
     InputContent,
+    GENERATE_UI_TOOL_NAME,
+    GenerateUserInterfaceToolArguments,
 )
 
 __all__ = [
@@ -105,4 +107,6 @@
     "TextInputContent",
     "BinaryInputContent",
     "InputContent",
+    "GENERATE_UI_TOOL_NAME",
+    "GenerateUserInterfaceToolArguments",
 ]

sdks/python/ag_ui/core/types.py

@@ -180,3 +180,15 @@ class RunAgentInput(ConfiguredBaseModel):
 
 # State can be any type
 State = Any
+
+
+GENERATE_UI_TOOL_NAME = "generateUserInterface"
+
+
+class GenerateUserInterfaceToolArguments(ConfiguredBaseModel):
+    """
+    Arguments for the generateUserInterface tool.
+    """
+    description: str
+    data: Optional[Any] = None
+    output: Optional[Any] = None  # JSON Schema for the expected output

sdks/python/tests/test_generative_ui.py

@@ -0,0 +1,23 @@
+from ag_ui.core.types import GENERATE_UI_TOOL_NAME, GenerateUserInterfaceToolArguments
+
+def test_constants():
+    assert GENERATE_UI_TOOL_NAME == "generateUserInterface"
+
+def test_arguments_validation():
+    # Valid arguments
+    args = GenerateUserInterfaceToolArguments(
+        description="A test form",
+        data={"foo": "bar"},
+        output={"type": "object"}
+    )
+    assert args.description == "A test form"
+    assert args.data == {"foo": "bar"}
+    assert args.output == {"type": "object"}
+
+    # Minimal arguments
+    args_minimal = GenerateUserInterfaceToolArguments(
+        description="Minimal form"
+    )
+    assert args_minimal.description == "Minimal form"
+    assert args_minimal.data is None
+    assert args_minimal.output is None

sdks/typescript/packages/client/src/agent/__tests__/generative-ui.test.ts

@@ -0,0 +1,79 @@
+import { AbstractAgent } from "../agent";
+import { UiGenerator } from "../generative-ui";
+import { GenerateUserInterfaceToolArguments, GENERATE_UI_TOOL_NAME, Message, RunAgentInput, BaseEvent } from "@ag-ui/core";
+import { Observable, of } from "rxjs";
+
+class TestAgent extends AbstractAgent {
+    run(input: RunAgentInput): Observable<BaseEvent> {
+        return of();
+    }
+}
+
+describe("Generative UI Integration", () => {
+    let agent: TestAgent;
+    let mockGenerator: UiGenerator;
+
+    beforeEach(() => {
+        agent = new TestAgent();
+        mockGenerator = {
+            generate: jest.fn().mockResolvedValue({ status: "success" }),
+        };
+        agent.registerUiGenerator(mockGenerator);
+    });
+
+    it("should intercept generateUserInterface tool calls and delegate to the generator", async () => {
+        const toolCallArgs: GenerateUserInterfaceToolArguments = {
+            description: "A test form",
+            data: { foo: "bar" },
+            output: { type: "object" },
+        };
+
+        const message: Message = {
+            id: "msg-1",
+            role: "assistant",
+            content: "",
+            toolCalls: [
+                {
+                    id: "call-1",
+                    type: "function",
+                    function: {
+                        name: GENERATE_UI_TOOL_NAME,
+                        arguments: JSON.stringify(toolCallArgs),
+                    },
+                },
+            ],
+        };
+
+        // Simulate adding a message with the tool call
+        agent.addMessage(message);
+
+        // Wait a bit for the async processing (addMessage triggers async subscribers/processing)
+        await new Promise((resolve) => setTimeout(resolve, 100));
+
+        expect(mockGenerator.generate).toHaveBeenCalledWith(toolCallArgs);
+    });
+
+    it("should not intercept other tool calls", async () => {
+        const message: Message = {
+            id: "msg-2",
+            role: "assistant",
+            content: "",
+            toolCalls: [
+                {
+                    id: "call-2",
+                    type: "function",
+                    function: {
+                        name: "someOtherTool",
+                        arguments: "{}",
+                    },
+                },
+            ],
+        };
+
+        agent.addMessage(message);
+
+        await new Promise((resolve) => setTimeout(resolve, 100));
+
+        expect(mockGenerator.generate).not.toHaveBeenCalled();
+    });
+});

sdks/typescript/packages/client/src/agent/agent.ts

@@ -23,6 +23,8 @@ import {
   BackwardCompatibility_0_0_39,
 } from "@/middleware";
 import packageJson from "../../package.json";
+import { UiGenerator, DefaultUiGenerator } from "./generative-ui";
+import { GENERATE_UI_TOOL_NAME, GenerateUserInterfaceToolArgumentsSchema } from "@ag-ui/core";
 
 export interface RunAgentResult {
   result: any;
@@ -39,6 +41,7 @@ export abstract class AbstractAgent {
   public subscribers: AgentSubscriber[] = [];
   public isRunning: boolean = false;
   private middlewares: Middleware[] = [];
+  private uiGenerator: UiGenerator = new DefaultUiGenerator();
   // Emits to immediately detach from the active run (stop processing its stream)
   private activeRunDetach$?: Subject<void>;
   private activeRunCompletionPromise?: Promise<void>;
@@ -86,6 +89,10 @@ export abstract class AbstractAgent {
     return this;
   }
 
+  public registerUiGenerator(generator: UiGenerator) {
+    this.uiGenerator = generator;
+  }
+
   public async runAgent(
     parameters?: RunAgentParameters,
     subscriber?: AgentSubscriber,
@@ -231,7 +238,7 @@ export abstract class AbstractAgent {
     }
   }
 
-  public abortRun() {}
+  public abortRun() { }
 
   public async detachActiveRun(): Promise<void> {
     if (!this.activeRunDetach$) {
@@ -459,6 +466,27 @@ export abstract class AbstractAgent {
       // Fire onNewToolCall if the message is from assistant and contains tool calls
       if (message.role === "assistant" && message.toolCalls) {
         for (const toolCall of message.toolCalls) {
+          // Intercept Generate UI tool calls
+          if (toolCall.function.name === GENERATE_UI_TOOL_NAME) {
+            try {
+              const args = JSON.parse(toolCall.function.arguments);
+              const parsedArgs = GenerateUserInterfaceToolArgumentsSchema.parse(args);
+
+              // Execute UI generation in background
+              this.uiGenerator.generate(parsedArgs).then((result) => {
+                // Here we would ideally send a tool output back to the agent
+                // For now, we just log or could emit a custom event if needed
+                if (this.debug) {
+                  console.debug("[GenerativeUI] Generated UI:", result);
+                }
+              }).catch(err => {
+                console.error("[GenerativeUI] Error generating UI:", err);
+              });
+            } catch (e) {
+              console.error("[GenerativeUI] Invalid arguments for UI generation:", e);
+            }
+          }
+
           for (const subscriber of this.subscribers) {
             await subscriber.onNewToolCall?.({
               toolCall,
@@ -502,6 +530,25 @@ export abstract class AbstractAgent {
         // Fire onNewToolCall if the message is from assistant and contains tool calls
         if (message.role === "assistant" && message.toolCalls) {
           for (const toolCall of message.toolCalls) {
+            // Intercept Generate UI tool calls
+            if (toolCall.function.name === GENERATE_UI_TOOL_NAME) {
+              try {
+                const args = JSON.parse(toolCall.function.arguments);
+                const parsedArgs = GenerateUserInterfaceToolArgumentsSchema.parse(args);
+
+                // Execute UI generation in background
+                this.uiGenerator.generate(parsedArgs).then((result) => {
+                  if (this.debug) {
+                    console.debug("[GenerativeUI] Generated UI:", result);
+                  }
+                }).catch(err => {
+                  console.error("[GenerativeUI] Error generating UI:", err);
+                });
+              } catch (e) {
+                console.error("[GenerativeUI] Invalid arguments for UI generation:", e);
+              }
+            }
+
             for (const subscriber of this.subscribers) {
               await subscriber.onNewToolCall?.({
                 toolCall,

sdks/typescript/packages/client/src/agent/generative-ui.ts

@@ -0,0 +1,21 @@
+import { GenerateUserInterfaceToolArguments } from "@ag-ui/core";
+
+export interface UiGenerator {
+    generate(args: GenerateUserInterfaceToolArguments): Promise<any>;
+}
+
+export class DefaultUiGenerator implements UiGenerator {
+    async generate(args: GenerateUserInterfaceToolArguments): Promise<any> {
+        // In a real implementation, this might call an external service or LLM
+        // For now, we just return the arguments as a placeholder
+        return {
+            status: "generated",
+            originalArgs: args,
+            generatedUi: {
+                // This would be the actual generated UI schema/code
+                type: "placeholder",
+                message: "UI generation not yet implemented in default generator",
+            },
+        };
+    }
+}

sdks/typescript/packages/core/src/types.ts

@@ -167,3 +167,15 @@ export class AGUIConnectNotImplementedError extends AGUIError {
     super("Connect not implemented. This method is not supported by the current agent.");
   }
 }
+
+export const GENERATE_UI_TOOL_NAME = "generateUserInterface";
+
+export const GenerateUserInterfaceToolArgumentsSchema = z.object({
+  description: z.string(),
+  data: z.any().optional(),
+  output: z.any().optional(), // JSON Schema for the expected output
+});
+
+export type GenerateUserInterfaceToolArguments = z.infer<
+  typeof GenerateUserInterfaceToolArgumentsSchema
+>;

sdks/typescript/packages/proto/package.json

@@ -15,7 +15,7 @@
     "lint": "eslint \"src/**/*.ts*\"",
     "clean": "rm -rf dist .turbo node_modules",
     "test": "jest",
-    "generate": "mkdir -p ./src/generated && npx protoc --plugin=./node_modules/.bin/protoc-gen-ts_proto --ts_proto_out=./src/generated --ts_proto_opt=esModuleInterop=true,outputJsonMethods=false,outputClientImpl=false -I ./src/proto ./src/proto/*.proto",
+    "generate": "node -e \"require('fs').mkdirSync('./src/generated', {recursive:true})\" && npx protoc --plugin=./node_modules/.bin/protoc-gen-ts_proto.cmd --ts_proto_out=./src/generated --ts_proto_opt=esModuleInterop=true,outputJsonMethods=false,outputClientImpl=false -I ./src/proto ./src/proto/*.proto",
     "link:global": "pnpm link --global",
     "unlink:global": "pnpm unlink --global"
   },
@@ -33,4 +33,4 @@
     "tsup": "^8.0.2",
     "typescript": "^5.8.2"
   }
-}
+}