contrag_core/
entity.rs

1use candid::CandidType;
2use serde::Serialize;
3pub use crate::types::{EntityRelationship, RelationshipType};
4
5/// Trait that marks a struct as a RAG entity
6/// 
7/// Implement this trait for your canister data structures to enable
8/// automatic context building and relationship mapping.
9/// 
10/// # Example
11/// 
12/// ```rust
13/// use contrag_core::prelude::*;
14/// use candid::{CandidType, Deserialize};
15/// use serde::Serialize;
16/// 
17/// #[derive(CandidType, Deserialize, Serialize, Clone)]
18/// pub struct User {
19///     pub id: String,
20///     pub name: String,
21///     pub email: String,
22/// }
23/// 
24/// impl RagEntity for User {
25///     fn entity_type() -> &'static str {
26///         "User"
27///     }
28///     
29///     fn entity_id(&self) -> String {
30///         self.id.clone()
31///     }
32///     
33///     fn to_context_map(&self) -> Vec<(String, String)> {
34///         vec![
35///             ("id".to_string(), self.id.clone()),
36///             ("name".to_string(), self.name.clone()),
37///             ("email".to_string(), self.email.clone()),
38///         ]
39///     }
40///     
41///     fn relationships(&self) -> Vec<EntityRelationship> {
42///         vec![]
43///     }
44/// }
45/// ```
46pub trait RagEntity: CandidType + Serialize + Clone {
47    /// Returns the entity type identifier
48    fn entity_type() -> &'static str
49    where
50        Self: Sized;
51
52    /// Returns the unique identifier for this entity instance
53    fn entity_id(&self) -> String;
54
55    /// Converts the entity to a flat key-value representation
56    /// 
57    /// Use dot notation for nested fields: "profile.age"
58    fn to_context_map(&self) -> Vec<(String, String)>;
59
60    /// Returns relationships to other entities
61    fn relationships(&self) -> Vec<EntityRelationship>;
62
63    /// Converts the entity to a human-readable text representation
64    /// 
65    /// Override this for custom formatting. Default implementation
66    /// joins the context map entries.
67    fn to_text(&self) -> String {
68        let mut lines = vec![
69            format!("Entity: {}", Self::entity_type()),
70            format!("ID: {}", self.entity_id()),
71            String::from("---"),
72        ];
73
74        for (key, value) in self.to_context_map() {
75            lines.push(format!("{}: {}", key, value));
76        }
77
78        lines.join("\n")
79    }
80
81    /// Returns a summary of the entity (first N characters)
82    fn to_summary(&self, max_length: usize) -> String {
83        let text = self.to_text();
84        if text.len() <= max_length {
85            text
86        } else {
87            format!("{}...", &text[..max_length])
88        }
89    }
90}
91
92/// Helper function to flatten nested JSON-like structures using serde_json
93/// 
94/// This is useful when you have complex nested structs and want to
95/// automatically flatten them for context building.
96/// 
97/// # Example
98/// 
99/// ```rust
100/// use serde_json::json;
101/// use contrag_core::entity::flatten_json_to_context;
102/// 
103/// let data = json!({
104///     "user": {
105///         "name": "Alice",
106///         "profile": {
107///             "age": 30,
108///             "location": "NYC"
109///         }
110///     }
111/// });
112/// 
113/// let flattened = flatten_json_to_context(&data, "");
114/// // Results in:
115/// // [
116/// //   ("user.name", "Alice"),
117/// //   ("user.profile.age", "30"),
118/// //   ("user.profile.location", "NYC")
119/// // ]
120/// ```
121pub fn flatten_json_to_context(
122    value: &serde_json::Value,
123    prefix: &str,
124) -> Vec<(String, String)> {
125    let mut result = vec![];
126
127    match value {
128        serde_json::Value::Object(map) => {
129            for (key, val) in map {
130                let new_prefix = if prefix.is_empty() {
131                    key.clone()
132                } else {
133                    format!("{}.{}", prefix, key)
134                };
135                result.extend(flatten_json_to_context(val, &new_prefix));
136            }
137        }
138        serde_json::Value::Array(arr) => {
139            let items: Vec<String> = arr
140                .iter()
141                .map(|v| match v {
142                    serde_json::Value::String(s) => s.clone(),
143                    _ => v.to_string(),
144                })
145                .collect();
146            result.push((prefix.to_string(), items.join(", ")));
147        }
148        serde_json::Value::Null => {
149            result.push((prefix.to_string(), String::new()));
150        }
151        _ => {
152            result.push((prefix.to_string(), value.to_string().trim_matches('"').to_string()));
153        }
154    }
155
156    result
157}
158
159/// Helper macro to implement RagEntity with automatic flattening
160/// 
161/// This macro uses serde_json to automatically flatten your struct,
162/// which is useful for rapid prototyping or complex nested structures.
163#[macro_export]
164macro_rules! impl_rag_entity_auto {
165    ($struct_name:ident, $entity_type:expr, $id_field:ident) => {
166        impl RagEntity for $struct_name {
167            fn entity_type() -> &'static str {
168                $entity_type
169            }
170
171            fn entity_id(&self) -> String {
172                self.$id_field.clone()
173            }
174
175            fn to_context_map(&self) -> Vec<(String, String)> {
176                let json = serde_json::to_value(self).unwrap();
177                $crate::entity::flatten_json_to_context(&json, "")
178            }
179
180            fn relationships(&self) -> Vec<EntityRelationship> {
181                vec![]
182            }
183        }
184    };
185}
contrag_core/entity.rs

contrag_core/
entity.rs
