next_web_ai/chat/observation/
model_observation_context.rs

1use std::collections::BTreeMap;
2
3use next_web_core::util::any_map::AnyValue;
4
5use crate::{
6    chat::observation::ai_operation_metadata::AiOperationMetadata, util::key_value::KeyValue,
7};
8
9#[derive(Clone)]
10pub struct ModelObservationContext<Q, R> {
11    pub(crate) request: Q,
12    pub(crate) operation_metadata: AiOperationMetadata,
13    pub(crate) response: Option<R>,
14
15    /// impl context
16    pub(crate) map: dashmap::DashMap<String, AnyValue>,
17    pub(crate) name: Option<String>,
18    pub(crate) contextual_name: Option<String>,
19
20    pub(crate) low_cardinality_key_values: BTreeMap<String, Box<dyn KeyValue>>,
21    pub(crate) high_cardinality_key_values: BTreeMap<String, Box<dyn KeyValue>>,
22}
23
24impl<Q, R> ModelObservationContext<Q, R> {
25    pub fn new(request: Q, operation_metadata: AiOperationMetadata) -> Self {
26        Self {
27            request,
28            operation_metadata,
29            response: None,
30            map: Default::default(),
31            name: Default::default(),
32            contextual_name: Default::default(),
33            low_cardinality_key_values: Default::default(),
34            high_cardinality_key_values: Default::default(),
35        }
36    }
37}
38
39impl<Q, R> ModelObservationContext<Q, R> {
40    pub fn operation_metadata(&self) -> &AiOperationMetadata {
41        &self.operation_metadata
42    }
43
44    pub fn request(&self) -> &Q {
45        &self.request
46    }
47
48    pub fn response(&self) -> Option<&R> {
49        self.response.as_ref()
50    }
51    pub fn set_response(&mut self, response: R) {
52        self.response = Some(response);
53    }
54}
55
56impl<Q, R> ModelObservationContext<Q, R> {
57    pub fn name(&self) -> Option<&str> {
58        self.name.as_ref().map(|x| x.as_str())
59    }
60
61    pub fn set_name(&mut self, name: impl Into<String>) {
62        self.name = Some(name.into());
63    }
64
65    pub fn contextual_name(&self) -> Option<&str> {
66        self.contextual_name.as_ref().map(|x| x.as_str())
67    }
68
69    pub fn set_contextual_name(&mut self, contextual_name: impl Into<String>) {
70        self.contextual_name = Some(contextual_name.into());
71    }
72
73    pub fn put(&self, key: impl Into<String>, value: impl Into<AnyValue>) {
74        self.map.insert(key.into(), value.into());
75    }
76
77    pub fn get(&self, key: impl AsRef<str>) -> Option<AnyValue> {
78        self.map.get(key.as_ref()).map(|v| v.value().clone())
79    }
80
81    pub fn get_or_default(&self, key: impl AsRef<str>, default: impl Into<AnyValue>) -> AnyValue {
82        self.map
83            .get(key.as_ref())
84            .map(|v| v.value().clone())
85            .unwrap_or(default.into())
86    }
87
88    pub fn contains_key(&self, key: impl AsRef<str>) -> bool {
89        self.map.contains_key(key.as_ref())
90    }
91
92    pub fn remove(&self, key: impl AsRef<str>) -> Option<(String, AnyValue)> {
93        self.map.remove(key.as_ref())
94    }
95
96    pub fn clear(&self) {
97        self.map.clear()
98    }
99}