1use crate::{Agent, Result, types::Content};
2use async_trait::async_trait;
3use serde_json::Value;
4use std::collections::HashMap;
5use std::sync::Arc;
6
7#[async_trait]
8pub trait ReadonlyContext: Send + Sync {
9 fn invocation_id(&self) -> &str;
10 fn agent_name(&self) -> &str;
11 fn user_id(&self) -> &str;
12 fn app_name(&self) -> &str;
13 fn session_id(&self) -> &str;
14 fn branch(&self) -> &str;
15 fn user_content(&self) -> &Content;
16}
17
18pub trait State: Send + Sync {
20 fn get(&self, key: &str) -> Option<Value>;
21 fn set(&mut self, key: String, value: Value);
22 fn all(&self) -> HashMap<String, Value>;
23}
24
25pub trait ReadonlyState: Send + Sync {
26 fn get(&self, key: &str) -> Option<Value>;
27 fn all(&self) -> HashMap<String, Value>;
28}
29
30pub trait Session: Send + Sync {
32 fn id(&self) -> &str;
33 fn app_name(&self) -> &str;
34 fn user_id(&self) -> &str;
35 fn state(&self) -> &dyn State;
36 fn conversation_history(&self) -> Vec<Content>;
38 fn append_to_history(&self, _content: Content) {
40 }
42}
43
44#[async_trait]
45pub trait CallbackContext: ReadonlyContext {
46 fn artifacts(&self) -> Option<Arc<dyn Artifacts>>;
47}
48
49#[async_trait]
50pub trait InvocationContext: CallbackContext {
51 fn agent(&self) -> Arc<dyn Agent>;
52 fn memory(&self) -> Option<Arc<dyn Memory>>;
53 fn session(&self) -> &dyn Session;
54 fn run_config(&self) -> &RunConfig;
55 fn end_invocation(&self);
56 fn ended(&self) -> bool;
57}
58
59#[async_trait]
61pub trait Artifacts: Send + Sync {
62 async fn save(&self, name: &str, data: &crate::Part) -> Result<i64>;
63 async fn load(&self, name: &str) -> Result<crate::Part>;
64 async fn list(&self) -> Result<Vec<String>>;
65}
66
67#[async_trait]
68pub trait Memory: Send + Sync {
69 async fn search(&self, query: &str) -> Result<Vec<MemoryEntry>>;
70}
71
72#[derive(Debug, Clone)]
73pub struct MemoryEntry {
74 pub content: Content,
75 pub author: String,
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
79pub enum StreamingMode {
80 #[default]
81 Auto,
82 Enabled,
83 Disabled,
84}
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
88pub enum IncludeContents {
89 None,
91 #[default]
93 Default,
94}
95
96#[derive(Debug, Clone)]
97pub struct RunConfig {
98 pub streaming_mode: StreamingMode,
99}
100
101impl Default for RunConfig {
102 fn default() -> Self {
103 Self { streaming_mode: StreamingMode::Auto }
104 }
105}
106
107#[cfg(test)]
108mod tests {
109 use super::*;
110
111 #[test]
112 fn test_run_config_default() {
113 let config = RunConfig::default();
114 assert_eq!(config.streaming_mode, StreamingMode::Auto);
115 }
116
117 #[test]
118 fn test_streaming_mode() {
119 assert_eq!(StreamingMode::Auto, StreamingMode::Auto);
120 assert_ne!(StreamingMode::Auto, StreamingMode::Enabled);
121 }
122}