feat: Pydantic AI support

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,

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/examples/agents/predictive_state_updates.py

@@ -19,7 +19,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.
@@ -66,7 +66,7 @@ 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.
+    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.
@@ -80,10 +80,10 @@ async def chat_node(state: AgentState, config: RunnableConfig):
     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"
     }]
 
@@ -121,7 +121,7 @@ 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",

typescript-sdk/integrations/pydantic-ai/.npmignore

@@ -0,0 +1,12 @@
+.turbo
+.DS_Store
+.git
+.gitignore
+.idea
+.vscode
+.env
+__tests__
+src
+tsup.config.ts
+tsconfig.json
+jest.config.js