agent_stream_kit/
context.rs

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

agent_stream_kit/
context.rs
