use candid::CandidType;
use serde::Serialize;
pub use crate::types::{EntityRelationship, RelationshipType};
pub trait RagEntity: CandidType + Serialize + Clone {
fn entity_type() -> &'static str
where
Self: Sized;
fn entity_id(&self) -> String;
fn to_context_map(&self) -> Vec<(String, String)>;
fn relationships(&self) -> Vec<EntityRelationship>;
fn to_text(&self) -> String {
let mut lines = vec![
format!("Entity: {}", Self::entity_type()),
format!("ID: {}", self.entity_id()),
String::from("---"),
];
for (key, value) in self.to_context_map() {
lines.push(format!("{}: {}", key, value));
}
lines.join("\n")
}
fn to_summary(&self, max_length: usize) -> String {
let text = self.to_text();
if text.len() <= max_length {
text
} else {
format!("{}...", &text[..max_length])
}
}
}
pub fn flatten_json_to_context(
value: &serde_json::Value,
prefix: &str,
) -> Vec<(String, String)> {
let mut result = vec![];
match value {
serde_json::Value::Object(map) => {
for (key, val) in map {
let new_prefix = if prefix.is_empty() {
key.clone()
} else {
format!("{}.{}", prefix, key)
};
result.extend(flatten_json_to_context(val, &new_prefix));
}
}
serde_json::Value::Array(arr) => {
let items: Vec<String> = arr
.iter()
.map(|v| match v {
serde_json::Value::String(s) => s.clone(),
_ => v.to_string(),
})
.collect();
result.push((prefix.to_string(), items.join(", ")));
}
serde_json::Value::Null => {
result.push((prefix.to_string(), String::new()));
}
_ => {
result.push((prefix.to_string(), value.to_string().trim_matches('"').to_string()));
}
}
result
}
#[macro_export]
macro_rules! impl_rag_entity_auto {
($struct_name:ident, $entity_type:expr, $id_field:ident) => {
impl RagEntity for $struct_name {
fn entity_type() -> &'static str {
$entity_type
}
fn entity_id(&self) -> String {
self.$id_field.clone()
}
fn to_context_map(&self) -> Vec<(String, String)> {
let json = serde_json::to_value(self).unwrap();
$crate::entity::flatten_json_to_context(&json, "")
}
fn relationships(&self) -> Vec<EntityRelationship> {
vec![]
}
}
};
}