WIP

Author: wdoppenberg

rust-sdk/Cargo.lock

@@ -11,3 +11,10 @@ serde_json = { workspace = true }
 async-trait = "0.1.88"
 uuid = { version = "1.17.0", features = ["v4"] }
 futures = "0.3.31"
+json-patch = "4.0.0"
+log = "0.4.27"
+reqwest = { version = "0.12.22" , features = ["json", "stream"]}
+sse-client = "1.1.1"
+
+[dev-dependencies]
+tokio = { version = "1.36.0", features = ["full"] }

rust-sdk/crates/ag-ui-client/Cargo.toml

@@ -0,0 +1,36 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#   "uvicorn == 0.34.3",
+#   "pydantic-ai==0.4.10"
+# ]
+# ///
+
+import uvicorn
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.openai import OpenAIProvider
+
+model = OpenAIModel(
+    model_name="llama3.1",
+    provider=OpenAIProvider(
+        base_url="http://localhost:11434/v1", api_key="ollama"
+    ),
+)
+agent = Agent(model)
+
+
+@agent.tool_plain
+def temperature_celsius(city: str) -> float:
+    return 21.0
+
+
+@agent.tool_plain
+def temperature_fahrenheit(city: str) -> float:
+    return 69.8
+
+
+app = agent.to_ag_ui()
+
+if __name__ == "__main__":
+    uvicorn.run(app, port=3001)

rust-sdk/crates/ag-ui-client/scripts/pydantic_ai_server.py

@@ -3,16 +3,16 @@ use std::collections::HashSet;
 use std::sync::Arc;
 use thiserror::Error;
 
-use crate::event::EventExt;
 use crate::stream::EventStream;
 use crate::subscriber::AgentSubscriber;
-use ag_ui_core::event::Event;
+
 use ag_ui_core::types::context::Context;
 use ag_ui_core::types::ids::{AgentId, MessageId, RunId, ThreadId};
 use ag_ui_core::types::input::RunAgentInput;
 use ag_ui_core::types::message::Message;
 use ag_ui_core::types::tool::Tool;
-use ag_ui_core::{FwdProps, JsonValue, State};
+use ag_ui_core::{AgentState, FwdProps, JsonValue};
+use crate::event_handler::EventHandler;
 
 #[derive(Debug, Clone)]
 pub struct AgentConfig<StateT = JsonValue> {
@@ -42,11 +42,13 @@ where
 
 /// Parameters for running an agent.
 #[derive(Debug, Clone, Default)]
-pub struct RunAgentParams<FwdPropsT = JsonValue> {
+pub struct RunAgentParams<StateT: AgentState, FwdPropsT = JsonValue> {
     pub run_id: Option<RunId>,
     pub tools: Option<Vec<Tool>>,
     pub context: Option<Vec<Context>>,
     pub forwarded_props: Option<FwdPropsT>,
+    pub messages: Vec<Message>,
+    pub state: StateT,
 }
 
 #[derive(Debug, Clone)]
@@ -55,6 +57,8 @@ pub struct RunAgentResult {
     pub new_messages: Vec<Message>,
 }
 
+pub type AgentRunState<StateT, FwdPropsT> = RunAgentInput<StateT, FwdPropsT>;
+
 #[derive(Debug, Clone)]
 pub struct AgentStateMutation<StateT = JsonValue> {
     pub messages: Option<Vec<Message>>,
@@ -90,227 +94,105 @@ pub enum AgentError {
 #[async_trait::async_trait]
 pub trait Agent<StateT = JsonValue, FwdPropsT = JsonValue>: Send + Sync
 where
-    StateT: State,
+    StateT: AgentState,
     FwdPropsT: FwdProps,
 {
-    fn run<'a>(&'a self, input: &'a RunAgentInput<StateT, FwdPropsT>) -> EventStream<'a>;
-
-    // Idiomatic accessors for agent state.
-    fn agent_id(&self) -> Option<&AgentId>;
-    fn agent_id_mut(&mut self) -> &mut Option<AgentId>;
-    fn description(&self) -> &str;
-    fn description_mut(&mut self) -> &mut String;
-    fn thread_id(&self) -> &ThreadId;
-    fn thread_id_mut(&mut self) -> &mut ThreadId;
-    fn messages(&self) -> &[Message];
-    fn messages_mut(&mut self) -> &mut Vec<Message>;
-    fn state(&self) -> &StateT;
-    fn state_mut(&mut self) -> &mut StateT;
-    fn subscribers(&self) -> &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>];
-    fn subscribers_mut(&mut self) -> &mut Vec<Arc<dyn AgentSubscriber<StateT, FwdPropsT>>>;
-
-    /// Adds a subscriber to the agent.
-    fn add_subscriber(&mut self, subscriber: Arc<dyn AgentSubscriber<StateT, FwdPropsT>>) {
-        self.subscribers_mut().push(subscriber);
-    }
+    async fn run(
+        &self,
+        input: &RunAgentInput<StateT, FwdPropsT>,
+    ) -> Result<EventStream<'async_trait, StateT>, AgentError>;
 
     /// The main execution method, containing the full pipeline logic.
     async fn run_agent(
-        &mut self,
-        params: &RunAgentParams<FwdPropsT>,
-        subscriber: Option<Arc<dyn AgentSubscriber<StateT, FwdPropsT>>>,
+        &self,
+        params: &RunAgentParams<StateT, FwdPropsT>,
+        subscribers: Vec<Arc<dyn AgentSubscriber<StateT, FwdPropsT>>>,
     ) -> Result<RunAgentResult, AgentError> {
-        if self.agent_id().is_none() {
-            *self.agent_id_mut() = Some(AgentId::new());
-        }
-
-        let mut subscribers = self.subscribers().to_vec();
-        if let Some(sub) = subscriber {
-            subscribers.push(sub);
-        }
+        // TODO: Use Agent ID?
+        let agent_id = AgentId::random();
+
+        let input = RunAgentInput {
+            thread_id: ThreadId::random(),
+            run_id: params.run_id.clone().unwrap_or_else(RunId::random),
+            state: params.state.clone(),
+            messages: params.messages.clone(),
+            tools: params.tools.clone().unwrap_or_default(),
+            context: params.context.clone().unwrap_or_default(),
+            // TODO: Find suitable default value
+            forwarded_props: params.forwarded_props.clone().unwrap(),
+        };
+        let current_message_ids: HashSet<&MessageId> =
+            params.messages.iter().map(|m| m.id()).collect();
 
-        let input = self.prepare_run_agent_input(params);
-        let messages = self.messages().to_vec();
-        let current_message_ids: HashSet<&MessageId> = messages.iter().map(|m| m.id()).collect();
-        let mut result_val = JsonValue::Null;
+        // Initialize event handler with the current state
+        let mut event_handler = EventHandler::new(
+            params.messages.clone(),
+            params.state.clone(),
+            &input,
+            subscribers,
+        );
 
-        let mut stream = self.run(&input).fuse();
+        let mut stream = self.run(&input).await?.fuse();
 
         while let Some(event_result) = stream.next().await {
             match event_result {
                 Ok(event) => {
-                    let (mutation, value) = event
-                        .apply_and_process_event(&input, &messages, &input.state, &subscribers)
-                        .await?;
-                    result_val = JsonValue::from(value);
+                    let mutation = event_handler.handle_event(&event).await?;
+                    event_handler.apply_mutation(mutation).await?;
                 }
                 Err(e) => {
-                    // self.on_error(&input, &e, &subscribers).await?;
+                    event_handler.on_error(&e).await?;
                     return Err(e);
                 }
             }
         }
 
-        // self.on_finalize(&input, &subscribers).await?;
+        // Finalize the run
+        event_handler.on_finalize().await?;
 
-        let new_messages = self
-            .messages()
+        // Collect new messages
+        let new_messages = event_handler
+            .messages
             .iter()
             .filter(|m| !current_message_ids.contains(&m.id()))
             .cloned()
             .collect();
 
         Ok(RunAgentResult {
-            result: result_val,
+            result: event_handler.result,
             new_messages,
         })
     }
 
-    /// Helper to construct the input for the `run` method.
-    fn prepare_run_agent_input(
-        &self,
-        params: &RunAgentParams<FwdPropsT>,
-    ) -> RunAgentInput<StateT, FwdPropsT> {
-        RunAgentInput {
-            thread_id: self.thread_id().clone(),
-            run_id: params.run_id.clone().unwrap_or_else(|| RunId::new()),
-            state: self.state().clone(),
-            messages: self.messages().to_vec(),
-            tools: params.tools.clone().unwrap_or_default(),
-            context: params.context.clone().unwrap_or_default(),
-            // TODO: Find suitable default value
-            forwarded_props: params.forwarded_props.clone().unwrap(),
-        }
-    }
-
-    /// Processes a single event, applying mutations and notifying subscribers.
-    /// Returns the final result if the event is `Done`.
-    async fn apply_and_process_event(
-        &mut self,
-        event: Event,
-        input: &RunAgentInput<StateT, FwdPropsT>,
-        subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
-    ) -> Result<Option<JsonValue>, AgentError> {
-        // This is a simplified stand-in for the logic from `defaultApplyEvents` in TS.
-        // A full implementation would handle each event type to create the correct state mutation.
-        let (mutation, result) = match event {
-            Event::RunFinished(e) => {
-                for sub in subscribers {
-                    sub.on_run_finished(
-                        &e.result.clone().unwrap(),
-                        self.messages(),
-                        self.state(),
-                        input,
-                    )
-                    .await?;
-                }
-                (AgentStateMutation::default(), e.result)
-            }
-            // In a real implementation, other events like Text, ToolCall, etc.,
-            // would create mutations to update messages and state.
-            _ => (AgentStateMutation::default(), None),
-        };
-
-        self.apply_mutation(mutation, input, subscribers).await?;
-        Ok(result)
-    }
-
-    async fn on_initialize(
-        &mut self,
-        input: &mut RunAgentInput<StateT, FwdPropsT>,
-        subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
-    ) -> Result<(), AgentError> {
-        for subscriber in subscribers {
-            let mutation = subscriber
-                .on_run_initialized(self.messages(), self.state(), input)
-                .await?;
-
-            if mutation.messages.is_some() || mutation.state.is_some() {
-                if let Some(ref messages) = mutation.messages {
-                    input.messages = messages.clone();
-                }
-                if let Some(ref state) = mutation.state {
-                    input.state = state.clone();
-                }
-                self.apply_mutation(mutation, input, subscribers).await?;
-            }
-        }
-        Ok(())
-    }
-
-    async fn on_error(
-        &mut self,
-        input: &RunAgentInput<StateT, FwdPropsT>,
-        error: &AgentError,
-        subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
-    ) -> Result<(), AgentError> {
-        for subscriber in subscribers {
-            let mutation = subscriber
-                .on_run_failed(error, self.messages(), self.state(), input)
-                .await?;
-
-            self.apply_mutation(mutation, input, subscribers).await?;
-        }
-        Ok(())
-    }
-
-    async fn on_finalize(
-        &mut self,
-        input: &RunAgentInput<StateT, FwdPropsT>,
-        subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
-    ) -> Result<(), AgentError> {
-        for subscriber in subscribers {
-            let mutation = subscriber
-                .on_run_finalized(self.messages(), self.state(), input)
-                .await?;
-
-            self.apply_mutation(mutation, input, subscribers).await?;
-        }
-        Ok(())
-    }
-
-    async fn apply_mutation(
-        &mut self,
-        mutation: AgentStateMutation<StateT>,
-        input: &RunAgentInput<StateT, FwdPropsT>,
-        subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
-    ) -> Result<(), AgentError> {
-        if let Some(messages) = mutation.messages {
-            *self.messages_mut() = messages;
-            self.notify_messages_changed(input, subscribers).await?;
-        }
-
-        if let Some(state) = mutation.state {
-            *self.state_mut() = state;
-            self.notify_state_changed(input, subscribers).await?;
-        }
-
-        Ok(())
-    }
-
-    async fn notify_messages_changed(
-        &self,
-        input: &RunAgentInput<StateT, FwdPropsT>,
-        subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
-    ) -> Result<(), AgentError> {
-        for subscriber in subscribers {
-            subscriber
-                .on_messages_changed(self.messages(), self.state(), input)
-                .await?;
-        }
-        Ok(())
-    }
-
-    async fn notify_state_changed(
-        &self,
-        input: &RunAgentInput<StateT, FwdPropsT>,
-        subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
-    ) -> Result<(), AgentError> {
-        for subscriber in subscribers {
-            subscriber
-                .on_state_changed(self.messages(), self.state(), input)
-                .await?;
-        }
-        Ok(())
-    }
+    // Helper function to run subscribers that can return a mutation
+    // async fn run_subscribers_with_mutation<F, Fut>(
+    //     &self,
+    //     subscribers: &[Arc<dyn AgentSubscriber<StateT, FwdPropsT>>],
+    //     mut callback: F,
+    // ) -> Result<AgentStateMutation<StateT>, AgentError>
+    // where
+    //     F: FnMut(&Arc<dyn AgentSubscriber<StateT, FwdPropsT>>) -> Fut + Send,
+    //     Fut: std::future::Future<Output = Result<AgentStateMutation<StateT>, AgentError>>,
+    // {
+    //     let mut result = AgentStateMutation::default();
+    //
+    //     for subscriber in subscribers {
+    //         let mutation = callback(subscriber).await?;
+    //
+    //         if mutation.messages.is_some() {
+    //             result.messages = mutation.messages;
+    //         }
+    //
+    //         if mutation.state.is_some() {
+    //             result.state = mutation.state;
+    //         }
+    //
+    //         if mutation.stop_propagation {
+    //             result.stop_propagation = true;
+    //             break;
+    //         }
+    //     }
+    //
+    //     Ok(result)
+    // }
 }