feat: Pydantic AI support (#133)

* feat: Pydantic AI support

Add support for the Pydantic AI integration.

Add missing background change confirmation.

Increase agentic generation ui page size to allow it to display results
from the Pydantic AI integration better.

Add a description for the human in the loop generate_task_steps tool, so
the Pydantic AI example works as expected.

Merge write_document and confirm_changes tool into one, which allows it
work correctly with the Pydantic AI example, without the need for the
agent to perform two separate tool calls.

Sort the features, so they are easier to find when reading the code and
implementing new integrations.

Fixes: #5

* chore: revert sorting of features

Revert the sorting of features, as per review feedback.

* fix: predictive state examples

Fix predictive state examples by renaming the `write_document` tool to
`write_document_local` and updating the metadata accordingly. This
ensures that the examples correctly reflect the intended functionality
of the predictive state updates, while still allowing the Pydantic AI
example to leverage the AG-UI defined `write_document` tool.

Restore the `confirm_changes` tool to enable the confirmation of
changes for existing examples.

Author: stevenh

typescript-sdk/.gitignore

@@ -41,4 +41,7 @@ yarn-error.log*
 packages/proto/src/generated
 
 # LangGraph API
-**/**/.langgraph_api
+**/**/.langgraph_api
+
+# Python
+__pycache__/

typescript-sdk/apps/dojo/package.json

@@ -16,6 +16,7 @@
     "@ag-ui/llamaindex": "workspace:*",
     "@ag-ui/mastra": "workspace:*",
     "@ag-ui/middleware-starter": "workspace:*",
+    "@ag-ui/pydantic-ai": "workspace:*",
     "@ag-ui/server-starter": "workspace:*",
     "@ag-ui/server-starter-all-features": "workspace:*",
     "@ag-ui/vercel-ai-sdk": "workspace:*",

typescript-sdk/apps/dojo/src/agents.ts

@@ -13,6 +13,7 @@ import { AgnoAgent } from "@ag-ui/agno";
 import { LlamaIndexAgent } from "@ag-ui/llamaindex";
 import { CrewAIAgent } from "@ag-ui/crewai";
 import { mastra } from "./mastra";
+import { PydanticAIAgent } from "@ag-ui/pydantic-ai";
 
 export const agentsIntegrations: AgentIntegrationConfig[] = [
   {
@@ -23,6 +24,31 @@ export const agentsIntegrations: AgentIntegrationConfig[] = [
       };
     },
   },
+  {
+    id: "pydantic-ai",
+    agents: async () => {
+      return {
+        agentic_chat: new PydanticAIAgent({
+          url: "http://localhost:9000/agentic_chat/",
+        }),
+        agentic_generative_ui: new PydanticAIAgent({
+          url: "http://localhost:9000/agentic_generative_ui/",
+        }),
+        human_in_the_loop: new PydanticAIAgent({
+          url: "http://localhost:9000/human_in_the_loop/",
+        }),
+        predictive_state_updates: new PydanticAIAgent({
+          url: "http://localhost:9000/predictive_state_updates/",
+        }),
+        shared_state: new PydanticAIAgent({
+          url: "http://localhost:9000/shared_state/",
+        }),
+        tool_based_generative_ui: new PydanticAIAgent({
+          url: "http://localhost:9000/tool_based_generative_ui/",
+        }),
+      };
+    },
+  },
   {
     id: "server-starter",
     agents: async () => {

typescript-sdk/apps/dojo/src/app/[integrationId]/feature/agentic_chat/page.tsx

@@ -42,6 +42,10 @@ const Chat = () => {
     ],
     handler: ({ background }) => {
       setBackground(background);
+      return {
+        status: "success",
+        message: `Background changed to ${background}`,
+      };
     },
   });
 

typescript-sdk/apps/dojo/src/app/[integrationId]/feature/agentic_generative_ui/page.tsx

@@ -42,7 +42,7 @@ const Chat = () => {
 
       return (
         <div className="flex">
-          <div className="bg-gray-100 rounded-lg w-[500px] p-4 text-black space-y-2">
+          <div className="bg-gray-100 rounded-lg w-[800px] p-4 text-black space-y-2">
             {state.steps.map((step, index) => {
               if (step.status === "completed") {
                 return (
@@ -55,7 +55,7 @@ const Chat = () => {
                 index === state.steps.findIndex((s) => s.status === "pending")
               ) {
                 return (
-                  <div key={index} className="text-3xl font-bold text-slate-700">
+                  <div key={index} className="text-2xl font-bold text-slate-700">
                     <Spinner />
                     {step.description}
                   </div>

typescript-sdk/apps/dojo/src/app/[integrationId]/feature/human_in_the_loop/page.tsx

@@ -99,6 +99,7 @@ const Chat = () => {
   });
   useCopilotAction({
     name: "generate_task_steps",
+    description: "Generates a list of steps for the user to perform",
     parameters: [
       {
         name: "steps",

typescript-sdk/apps/dojo/src/app/[integrationId]/feature/predictive_state_updates/page.tsx

@@ -127,6 +127,7 @@ const DocumentEditor = () => {
     }
   }, [text]);
 
+  // TODO(steve): Remove this when all agents have been updated to use write_document tool.
   useCopilotAction({
     name: "confirm_changes",
     renderAndWaitForResponse: ({ args, respond, status }) => (
@@ -147,6 +148,40 @@ const DocumentEditor = () => {
     ),
   });
 
+  // Action to write the document.
+  useCopilotAction({
+    name: "write_document",
+    description: `Present the proposed changes to the user for review`,
+    parameters: [
+      {
+        name: "document",
+        type: "string",
+        description: "The full updated document in markdown format",
+      },
+    ],
+    renderAndWaitForResponse({ args, status, respond }) {
+      if (status === "executing") {
+        return (
+          <ConfirmChanges
+            args={args}
+            respond={respond}
+            status={status}
+            onReject={() => {
+              editor?.commands.setContent(fromMarkdown(currentDocument));
+              setAgentState({ document: currentDocument });
+            }}
+            onConfirm={() => {
+              editor?.commands.setContent(fromMarkdown(agentState?.document || ""));
+              setCurrentDocument(agentState?.document || "");
+              setAgentState({ document: agentState?.document || "" });
+            }}
+          />
+        );
+      }
+      return <></>;
+    },
+  });
+
   return (
     <div className="relative min-h-screen w-full">
       {placeholderVisible && (

typescript-sdk/apps/dojo/src/menu.ts

@@ -6,6 +6,18 @@ export const menuIntegrations: MenuIntegrationConfig[] = [
     name: "Middleware Starter",
     features: ["agentic_chat"],
   },
+  {
+    id: "pydantic-ai",
+    name: "Pydantic AI",
+    features: [
+      "agentic_chat",
+      "human_in_the_loop",
+      "agentic_generative_ui",
+      "tool_based_generative_ui",
+      "shared_state",
+      "predictive_state_updates",
+    ],
+  },
   {
     id: "server-starter",
     name: "Server Starter",

typescript-sdk/integrations/crewai/python/ag_ui_crewai/examples/predictive_state_updates.py

@@ -16,7 +16,7 @@
 WRITE_DOCUMENT_TOOL = {
     "type": "function",
     "function": {
-        "name": "write_document",
+        "name": "write_document_local",
         "description": " ".join("""
             Write a document. Use markdown formatting to format the document.
             It's good to format the document extensively so it's easy to read.
@@ -63,19 +63,19 @@ async def chat(self):
         Standard chat node.
         """
         system_prompt = f"""
-        You are a helpful assistant for writing documents. 
-        To write the document, you MUST use the write_document tool.
+        You are a helpful assistant for writing documents.
+        To write the document, you MUST use the write_document_local tool.
         You MUST write the full document, even when changing only a few words.
         When you wrote the document, DO NOT repeat it as a message. 
         Just briefly summarize the changes you made. 2 sentences max.
         This is the current state of the document: ----\n {self.state.document}\n-----
         """
 
-        # 1. Here we specify that we want to stream the tool call to write_document
+        # 1. Here we specify that we want to stream the tool call to write_document_local
         #    to the frontend as state.
         await copilotkit_predict_state({
             "document": {
-                "tool_name": "write_document",
+                "tool_name": "write_document_local",
                 "tool_argument": "document"
             }
         })
@@ -122,7 +122,7 @@ async def chat(self):
             tool_call_name = tool_call["function"]["name"]
             tool_call_args = json.loads(tool_call["function"]["arguments"])
 
-            if tool_call_name == "write_document":
+            if tool_call_name == "write_document_local":
                 self.state.document = tool_call_args["document"]
 
                 # 4.1 Append the result to the messages in state

typescript-sdk/integrations/langgraph/examples/agents/predictive_state_updates/agent.py

@@ -18,7 +18,7 @@
 WRITE_DOCUMENT_TOOL = {
     "type": "function",
     "function": {
-        "name": "write_document",
+        "name": "write_document_local",
         "description": " ".join("""
             Write a document. Use markdown formatting to format the document.
             It's good to format the document extensively so it's easy to read.
@@ -64,25 +64,25 @@ async def chat_node(state: AgentState, config: RunnableConfig):
     """
 
     system_prompt = f"""
-    You are a helpful assistant for writing documents. 
-    To write the document, you MUST use the write_document tool.
+    You are a helpful assistant for writing documents.
+    To write the document, you MUST use the write_document_local tool.
     You MUST write the full document, even when changing only a few words.
-    When you wrote the document, DO NOT repeat it as a message. 
+    When you wrote the document, DO NOT repeat it as a message.
     Just briefly summarize the changes you made. 2 sentences max.
     This is the current state of the document: ----\n {state.get('document')}\n-----
     """
 
     # Define the model
     model = ChatOpenAI(model="gpt-4o")
-    
+
     # Define config for the model with emit_intermediate_state to stream tool calls to frontend
     if config is None:
         config = RunnableConfig(recursion_limit=25)
 
-    # Use "predict_state" metadata to set up streaming for the write_document tool
+    # Use "predict_state" metadata to set up streaming for the write_document_local tool
     config["metadata"]["predict_state"] = [{
         "state_key": "document",
-        "tool": "write_document",
+        "tool": "write_document_local",
         "tool_argument": "document"
     }]
 
@@ -104,11 +104,11 @@ async def chat_node(state: AgentState, config: RunnableConfig):
 
     # Update messages with the response
     messages = state["messages"] + [response]
-    
+
     # Extract any tool calls from the response
     if hasattr(response, "tool_calls") and response.tool_calls:
         tool_call = response.tool_calls[0]
-        
+
         # Handle tool_call as a dictionary or an object
         if isinstance(tool_call, dict):
             tool_call_id = tool_call["id"]
@@ -120,14 +120,14 @@ async def chat_node(state: AgentState, config: RunnableConfig):
             tool_call_name = tool_call.name
             tool_call_args = tool_call.args
 
-        if tool_call_name == "write_document":
+        if tool_call_name == "write_document_local":
             # Add the tool response to messages
             tool_response = {
                 "role": "tool",
                 "content": "Document written.",
                 "tool_call_id": tool_call_id
             }
-            
+
             # Add confirmation tool call
             confirm_tool_call = {
                 "role": "assistant",
@@ -140,9 +140,9 @@ async def chat_node(state: AgentState, config: RunnableConfig):
                     }
                 }]
             }
-            
+
             messages = messages + [tool_response, confirm_tool_call]
-            
+
             # Return Command to route to end
             return Command(
                 goto=END,
@@ -151,7 +151,7 @@ async def chat_node(state: AgentState, config: RunnableConfig):
                     "document": tool_call_args["document"]
                 }
             )
-    
+
     # If no tool was called, go to end
     return Command(
         goto=END,