agent_tk/
knowledge_base.rs

1//! interface to the knowledge base of an agent
2use crate::{Error, Result};
3
4//  _  __                    _          _            ____
5// | |/ /_ __   _____      _| | ___  __| | __ _  ___| __ )  __ _ ___  ___
6// | ' /| '_ \ / _ \ \ /\ / / |/ _ \/ _` |/ _` |/ _ \  _ \ / _` / __|/ _ \
7// | . \| | | | (_) \ V  V /| |  __/ (_| | (_| |  __/ |_) | (_| \__ \  __/
8// |_|\_\_| |_|\___/ \_/\_/ |_|\___|\__,_|\__, |\___|____/ \__,_|___/\___|
9//                                        |___/
10
11/// Trait for the low level interface to the knowledge base
12pub trait KnowledgeBase: Sync + Send + dyn_clone::DynClone + 'static
13{
14  /// Insert data
15  fn insert_data(&self, collection: String, key: String, data: Vec<u8>) -> Result<()>;
16  /// Retrieve data
17  fn retrieve_data(&self, collection: String, key: String) -> Result<Vec<u8>>;
18}
19
20dyn_clone::clone_trait_object!(KnowledgeBase);
21
22//  _  __                    _          _            ____                 ___       _             __
23// | |/ /_ __   _____      _| | ___  __| | __ _  ___| __ )  __ _ ___  ___|_ _|_ __ | |_ ___ _ __ / _| __ _  ___ ___
24// | ' /| '_ \ / _ \ \ /\ / / |/ _ \/ _` |/ _` |/ _ \  _ \ / _` / __|/ _ \| || '_ \| __/ _ \ '__| |_ / _` |/ __/ _ \
25// | . \| | | | (_) \ V  V /| |  __/ (_| | (_| |  __/ |_) | (_| \__ \  __/| || | | | ||  __/ |  |  _| (_| | (_|  __/
26// |_|\_\_| |_|\___/ \_/\_/ |_|\___|\__,_|\__, |\___|____/ \__,_|___/\___|___|_| |_|\__\___|_|  |_|  \__,_|\___\___|
27//                                        |___/
28
29/// High level interface to the knowledge base
30pub trait KnowledgeBaseInterface: KnowledgeBase
31{
32  /// Insert a serializable object
33  fn insert(
34    &self,
35    collection: impl Into<String>,
36    key: impl Into<String>,
37    data: &impl serde::Serialize,
38  ) -> Result<()>;
39  /// Retrieve and deserialize object
40  fn retrieve<T: serde::de::DeserializeOwned>(
41    &self,
42    collection: impl Into<String>,
43    key: impl Into<String>,
44  ) -> Result<T>;
45}
46
47impl<TKnowledgeBase: ?Sized + KnowledgeBase> KnowledgeBaseInterface for TKnowledgeBase
48{
49  /// Insert a serializable object
50  fn insert(
51    &self,
52    collection: impl Into<String>,
53    key: impl Into<String>,
54    data: &impl serde::Serialize,
55  ) -> Result<()>
56  {
57    let mut data_u8 = Vec::<u8>::new();
58    ciborium::into_writer(data, &mut data_u8)?;
59    self.insert_data(collection.into(), key.into(), data_u8)
60  }
61  fn retrieve<T: serde::de::DeserializeOwned>(
62    &self,
63    collection: impl Into<String>,
64    key: impl Into<String>,
65  ) -> Result<T>
66  {
67    let data = self.retrieve_data(collection.into(), key.into())?;
68    ciborium::from_reader::<T, &[u8]>(data.as_ref())
69      .map_err(|e| Error::DeserializationError(e.to_string()))
70  }
71}
72
73//  ___       __  __                                 ____
74// |_ _|_ __ |  \/  | ___ _ __ ___   ___  _ __ _   _| __ )  __ _ ___  ___
75//  | || '_ \| |\/| |/ _ \ '_ ` _ \ / _ \| '__| | | |  _ \ / _` / __|/ _ \
76//  | || | | | |  | |  __/ | | | | | (_) | |  | |_| | |_) | (_| \__ \  __/
77// |___|_| |_|_|  |_|\___|_| |_| |_|\___/|_|   \__, |____/ \__,_|___/\___|
78//                                             |___/
79
80/// An implementation of KnowledgeBase where the data is stored in memory
81#[derive(Clone)]
82pub struct InMemoryBase
83{
84  storage: ccutils::sync::ArcMutex<
85    std::collections::HashMap<String, std::collections::HashMap<String, Vec<u8>>>,
86  >,
87}
88
89impl Default for InMemoryBase
90{
91  fn default() -> Self
92  {
93    Self::new()
94  }
95}
96
97impl InMemoryBase
98{
99  /// Create a new in memory database
100  pub fn new() -> Self
101  {
102    Self {
103      storage: Default::default(),
104    }
105  }
106}
107
108impl KnowledgeBase for InMemoryBase
109{
110  fn insert_data(&self, collection: String, key: String, data: Vec<u8>) -> Result<()>
111  {
112    let mut storage = self.storage.lock()?;
113    if !storage.contains_key(&collection)
114    {
115      storage.insert(collection.clone(), Default::default());
116    }
117    if let Some(v) = storage.get_mut(&collection)
118    {
119      v.insert(key, data);
120      Ok(())
121    }
122    else
123    {
124      Err(Error::InvalidCollection(collection))
125    }
126  }
127  fn retrieve_data(&self, collection: String, key: String) -> Result<Vec<u8>>
128  {
129    let storage = self.storage.lock()?;
130    if let Some(v) = storage.get(&collection)
131      && let Some(v) = v.get(&key)
132    {
133      return Ok(v.to_owned());
134    }
135
136    Err(Error::UnknownValue(collection, key))
137  }
138}
agent_tk/knowledge_base.rs

agent_tk/
knowledge_base.rs
