adk_core/
context.rs

1use crate::{types::Content, Agent, Result};
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}
39
40#[async_trait]
41pub trait CallbackContext: ReadonlyContext {
42    fn artifacts(&self) -> Option<Arc<dyn Artifacts>>;
43}
44
45#[async_trait]
46pub trait InvocationContext: CallbackContext {
47    fn agent(&self) -> Arc<dyn Agent>;
48    fn memory(&self) -> Option<Arc<dyn Memory>>;
49    fn session(&self) -> &dyn Session;
50    fn run_config(&self) -> &RunConfig;
51    fn end_invocation(&self);
52    fn ended(&self) -> bool;
53}
54
55// Placeholder service traits
56#[async_trait]
57pub trait Artifacts: Send + Sync {
58    async fn save(&self, name: &str, data: &crate::Part) -> Result<i64>;
59    async fn load(&self, name: &str) -> Result<crate::Part>;
60    async fn list(&self) -> Result<Vec<String>>;
61}
62
63#[async_trait]
64pub trait Memory: Send + Sync {
65    async fn search(&self, query: &str) -> Result<Vec<MemoryEntry>>;
66}
67
68#[derive(Debug, Clone)]
69pub struct MemoryEntry {
70    pub content: Content,
71    pub author: String,
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
75pub enum StreamingMode {
76    #[default]
77    Auto,
78    Enabled,
79    Disabled,
80}
81
82/// Controls what parts of prior conversation history is received by llmagent
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
84pub enum IncludeContents {
85    /// The llmagent operates solely on its current turn (latest user input + any following agent events)
86    None,
87    /// Default - The llmagent receives the relevant conversation history
88    #[default]
89    Default,
90}
91
92#[derive(Debug, Clone)]
93pub struct RunConfig {
94    pub streaming_mode: StreamingMode,
95}
96
97impl Default for RunConfig {
98    fn default() -> Self {
99        Self { streaming_mode: StreamingMode::Auto }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn test_run_config_default() {
109        let config = RunConfig::default();
110        assert_eq!(config.streaming_mode, StreamingMode::Auto);
111    }
112
113    #[test]
114    fn test_streaming_mode() {
115        assert_eq!(StreamingMode::Auto, StreamingMode::Auto);
116        assert_ne!(StreamingMode::Auto, StreamingMode::Enabled);
117    }
118}
adk_core/context.rs

adk_core/
context.rs
