yoagent_state/
adapter.rs

1use crate::{ActorRef, Event, EventId, EventStore, RunId, StateError, YoAgentState};
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4use serde_json::{Value as JsonValue, json};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct YoAgentRunStarted {
8    pub run_id: RunId,
9    pub task: String,
10    pub metadata: JsonValue,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct YoAgentRunFinished {
15    pub run_id: RunId,
16    pub outcome: String,
17    pub metadata: JsonValue,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct YoAgentModelCalled {
22    pub run_id: RunId,
23    pub model: String,
24    pub prompt_summary: String,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct YoAgentModelFinished {
29    pub run_id: RunId,
30    pub model: String,
31    pub output_summary: String,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct YoAgentToolCalled {
36    pub run_id: RunId,
37    pub tool: String,
38    pub input_summary: String,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct YoAgentToolFinished {
43    pub run_id: RunId,
44    pub tool: String,
45    pub output_summary: String,
46    pub success: bool,
47}
48
49#[async_trait]
50pub trait YoAgentStateSink: Send + Sync {
51    async fn on_run_started(&self, event: YoAgentRunStarted) -> Result<EventId, StateError>;
52    async fn on_run_finished(&self, event: YoAgentRunFinished) -> Result<EventId, StateError>;
53    async fn on_model_called(&self, event: YoAgentModelCalled) -> Result<EventId, StateError>;
54    async fn on_model_finished(&self, event: YoAgentModelFinished) -> Result<EventId, StateError>;
55    async fn on_tool_called(&self, event: YoAgentToolCalled) -> Result<EventId, StateError>;
56    async fn on_tool_finished(&self, event: YoAgentToolFinished) -> Result<EventId, StateError>;
57}
58
59#[derive(Debug, Clone)]
60pub struct YoAgentStateAdapter<S: EventStore> {
61    state: YoAgentState<S>,
62    actor: ActorRef,
63}
64
65impl<S: EventStore> YoAgentStateAdapter<S> {
66    pub fn new(state: YoAgentState<S>, actor: ActorRef) -> Self {
67        Self { state, actor }
68    }
69
70    async fn record<T: Serialize + Send + Sync>(
71        &self,
72        kind: &'static str,
73        event: T,
74    ) -> Result<EventId, StateError> {
75        self.state
76            .record_event(Event::new(
77                self.actor.clone(),
78                kind,
79                serde_json::to_value(event)?,
80            ))
81            .await
82    }
83}
84
85#[async_trait]
86impl<S: EventStore> YoAgentStateSink for YoAgentStateAdapter<S> {
87    async fn on_run_started(&self, event: YoAgentRunStarted) -> Result<EventId, StateError> {
88        self.record("run.started", event).await
89    }
90
91    async fn on_run_finished(&self, event: YoAgentRunFinished) -> Result<EventId, StateError> {
92        self.record("run.finished", event).await
93    }
94
95    async fn on_model_called(&self, event: YoAgentModelCalled) -> Result<EventId, StateError> {
96        self.record("model.called", event).await
97    }
98
99    async fn on_model_finished(&self, event: YoAgentModelFinished) -> Result<EventId, StateError> {
100        self.record("model.finished", event).await
101    }
102
103    async fn on_tool_called(&self, event: YoAgentToolCalled) -> Result<EventId, StateError> {
104        self.record("tool.called", event).await
105    }
106
107    async fn on_tool_finished(&self, event: YoAgentToolFinished) -> Result<EventId, StateError> {
108        if !event.success {
109            self.state
110                .record_event(Event::new(
111                    self.actor.clone(),
112                    "failure.observed",
113                    json!({
114                        "run_id": event.run_id,
115                        "tool": event.tool,
116                        "output_summary": event.output_summary,
117                    }),
118                ))
119                .await?;
120            return self.record("tool.finished", event).await;
121        }
122
123        self.record("tool.finished", event).await
124    }
125}
yoagent_state/adapter.rs

yoagent_state/
adapter.rs
