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/// Streaming mode for agent responses.
79/// Matches ADK Python/Go specification.
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
81pub enum StreamingMode {
82    /// No streaming; responses delivered as complete units.
83    /// Agent collects all chunks internally and yields a single final event.
84    None,
85    /// Server-Sent Events streaming; one-way streaming from server to client.
86    /// Agent yields each chunk as it arrives with stable event ID.
87    #[default]
88    SSE,
89    /// Bidirectional streaming; simultaneous communication in both directions.
90    /// Used for realtime audio/video agents.
91    Bidi,
92}
93
94/// Controls what parts of prior conversation history is received by llmagent
95#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
96pub enum IncludeContents {
97    /// The llmagent operates solely on its current turn (latest user input + any following agent events)
98    None,
99    /// Default - The llmagent receives the relevant conversation history
100    #[default]
101    Default,
102}
103
104#[derive(Debug, Clone)]
105pub struct RunConfig {
106    pub streaming_mode: StreamingMode,
107}
108
109impl Default for RunConfig {
110    fn default() -> Self {
111        Self { streaming_mode: StreamingMode::SSE }
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn test_run_config_default() {
121        let config = RunConfig::default();
122        assert_eq!(config.streaming_mode, StreamingMode::SSE);
123    }
124
125    #[test]
126    fn test_streaming_mode() {
127        assert_eq!(StreamingMode::SSE, StreamingMode::SSE);
128        assert_ne!(StreamingMode::SSE, StreamingMode::None);
129        assert_ne!(StreamingMode::None, StreamingMode::Bidi);
130    }
131}