1use crate::{
19 ActorRef, Event, EventId, EventStore, ModelCall, NodeId, RunId, StateError, ToolCall,
20 YoAgentState,
21};
22use async_trait::async_trait;
23use serde::{Deserialize, Serialize};
24use serde_json::{Value as JsonValue, json};
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct YoAgentRunStarted {
28 pub run_id: RunId,
29 pub task: String,
30 pub metadata: JsonValue,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct YoAgentRunFinished {
35 pub run_id: RunId,
36 pub outcome: String,
37 pub metadata: JsonValue,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct YoAgentModelCalled {
42 pub run_id: RunId,
43 pub model: String,
44 pub prompt_summary: String,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct YoAgentModelFinished {
49 pub run_id: RunId,
50 pub model: String,
51 pub output_summary: String,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct YoAgentToolCalled {
56 pub run_id: RunId,
57 pub tool: String,
58 pub input_summary: String,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct YoAgentToolFinished {
63 pub run_id: RunId,
64 pub tool: String,
65 pub output_summary: String,
66 pub success: bool,
67}
68
69#[async_trait]
70pub trait YoAgentStateSink: Send + Sync {
71 async fn on_run_started(&self, event: YoAgentRunStarted) -> Result<EventId, StateError>;
72 async fn on_run_finished(&self, event: YoAgentRunFinished) -> Result<EventId, StateError>;
73 async fn on_model_called(&self, event: YoAgentModelCalled) -> Result<EventId, StateError>;
74 async fn on_model_finished(&self, event: YoAgentModelFinished) -> Result<EventId, StateError>;
75 async fn on_tool_called(&self, event: YoAgentToolCalled) -> Result<EventId, StateError>;
76 async fn on_tool_finished(&self, event: YoAgentToolFinished) -> Result<EventId, StateError>;
77}
78
79#[derive(Debug, Clone)]
80pub struct YoAgentStateAdapter<S: EventStore> {
81 state: YoAgentState<S>,
82 actor: ActorRef,
83}
84
85impl<S: EventStore> YoAgentStateAdapter<S> {
86 pub fn new(state: YoAgentState<S>, actor: ActorRef) -> Self {
87 Self { state, actor }
88 }
89
90 async fn record<T: Serialize + Send + Sync>(
91 &self,
92 kind: &'static str,
93 event: T,
94 ) -> Result<EventId, StateError> {
95 self.state
96 .record_event(Event::new(
97 self.actor.clone(),
98 kind,
99 serde_json::to_value(event)?,
100 ))
101 .await
102 }
103}
104
105#[async_trait]
106impl<S: EventStore> YoAgentStateSink for YoAgentStateAdapter<S> {
107 async fn on_run_started(&self, event: YoAgentRunStarted) -> Result<EventId, StateError> {
108 self.state
109 .record_run_started(self.actor.clone(), event.run_id, event.task)
110 .await
111 }
112
113 async fn on_run_finished(&self, event: YoAgentRunFinished) -> Result<EventId, StateError> {
114 self.state
115 .record_run_finished(self.actor.clone(), event.run_id, event.outcome)
116 .await
117 }
118
119 async fn on_model_called(&self, event: YoAgentModelCalled) -> Result<EventId, StateError> {
120 self.state
121 .record_model_call(
122 self.actor.clone(),
123 ModelCall {
124 id: NodeId::generate(),
125 run_id: event.run_id,
126 model: event.model,
127 prompt_summary: event.prompt_summary,
128 output_summary: None,
129 metadata: json!({}),
130 },
131 )
132 .await
133 }
134
135 async fn on_model_finished(&self, event: YoAgentModelFinished) -> Result<EventId, StateError> {
136 self.record("model.finished", event).await
138 }
139
140 async fn on_tool_called(&self, event: YoAgentToolCalled) -> Result<EventId, StateError> {
141 self.state
142 .record_tool_call(
143 self.actor.clone(),
144 ToolCall {
145 id: NodeId::generate(),
146 run_id: event.run_id,
147 tool: event.tool,
148 input_summary: event.input_summary,
149 output_summary: None,
150 success: None,
151 metadata: json!({}),
152 },
153 )
154 .await
155 }
156
157 async fn on_tool_finished(&self, event: YoAgentToolFinished) -> Result<EventId, StateError> {
158 if !event.success {
159 self.state
161 .record_failure(
162 self.actor.clone(),
163 NodeId::generate(),
164 format!("tool {} failed", event.tool),
165 event.output_summary.clone(),
166 )
167 .await?;
168 }
169 self.record("tool.finished", event).await
170 }
171}