next_web_ai/chat/observation/
model_observation_context.rs1use 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 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}