agent_stream_kit/
context.rs

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