agent_stream_kit/
context.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use serde::{Deserialize, Serialize};
4
5use super::data::AgentValue;
6
7#[derive(Clone, Debug, Default, Serialize, Deserialize)]
8pub struct AgentContext {
9    port: String,
10
11    #[serde(skip_serializing_if = "Option::is_none")]
12    vars: Option<Arc<BTreeMap<String, AgentValue>>>,
13}
14
15impl AgentContext {
16    pub fn new() -> Self {
17        Self::default()
18    }
19
20    // Port
21
22    pub fn new_with_port(port: impl Into<String>) -> Self {
23        Self {
24            port: port.into(),
25            vars: None,
26        }
27    }
28
29    pub fn with_port(&self, port: impl Into<String>) -> Self {
30        Self {
31            port: port.into(),
32            vars: self.vars.clone(),
33        }
34    }
35
36    pub fn port(&self) -> &str {
37        &self.port
38    }
39
40    // Variables
41
42    pub fn get_var(&self, key: &str) -> Option<&AgentValue> {
43        self.vars.as_ref().and_then(|vars| vars.get(key))
44    }
45
46    pub fn with_var(&self, key: String, value: AgentValue) -> Self {
47        let mut vars = if let Some(vars) = &self.vars {
48            vars.as_ref().clone()
49        } else {
50            BTreeMap::new()
51        };
52        vars.insert(key, value);
53        Self {
54            port: self.port.clone(),
55            vars: Some(Arc::new(vars)),
56        }
57    }
58}