adk_core/
context.rs

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
18// State management traits
19pub 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
30// Session trait
31pub 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    /// Returns the conversation history from this session as Content items
37    fn conversation_history(&self) -> Vec<Content>;
38    /// Append content to conversation history (for sequential agent support)
39    fn append_to_history(&self, _content: Content) {
40        // Default no-op - implementations can override to track history
41    }
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// Placeholder service traits
60#[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/// Controls what parts of prior conversation history is received by llmagent
87#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
88pub enum IncludeContents {
89    /// The llmagent operates solely on its current turn (latest user input + any following agent events)
90    None,
91    /// Default - The llmagent receives the relevant conversation history
92    #[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}