diff --git a/api/build.rs b/api/build.rs index da489a8..191dcc3 100644 --- a/api/build.rs +++ b/api/build.rs @@ -17,8 +17,8 @@ fn main() { .map(|desc| desc.trim().to_string()) .unwrap_or_else(|| "unknown".to_string()); - println!("cargo::rustc-env=GIT_COMMIT={}", git_hash); - println!("cargo::rustc-env=GIT_TAG={}", git_tag); + println!("cargo::rustc-env=GIT_COMMIT={git_hash}"); + println!("cargo::rustc-env=GIT_TAG={git_tag}"); // Always rerun if any git changes occur println!("cargo::rerun-if-changed=.git/HEAD"); diff --git a/api/src/query_mapping.rs b/api/src/query_mapping.rs index 43a991b..2b94d25 100644 --- a/api/src/query_mapping.rs +++ b/api/src/query_mapping.rs @@ -28,7 +28,7 @@ impl QueryMapper { self.node_counter += 1; self.match_statements - .push(format!("MATCH ({} {{id: \"{id}\"}})", node_var)); + .push(format!("MATCH ({node_var} {{id: \"{id}\"}})")); self.return_statement_vars.insert(node_var.clone()); selection @@ -50,10 +50,8 @@ impl QueryMapper { self.relation_counter += 1; self.node_counter += 1; - self.match_statements.push(format!( - "MATCH ({}) -[{}]-> ({})", - node_var, relation_var, to_var - )); + self.match_statements + .push(format!("MATCH ({node_var}) -[{relation_var}]-> ({to_var})")); self.return_statement_vars.insert(self.relation_var()); selection diff --git a/grc20-core/src/ids/base58.rs b/grc20-core/src/ids/base58.rs index 6bb21b9..35a68ad 100644 --- a/grc20-core/src/ids/base58.rs +++ b/grc20-core/src/ids/base58.rs @@ -34,7 +34,7 @@ pub fn decode_base58_to_uuid(encoded: &str) -> Result { } } - let hex_str = format!("{:032x}", decoded); + let hex_str = format!("{decoded:032x}"); Ok(format!( "{}-{}-{}-{}-{}", &hex_str[0..8], diff --git a/grc20-core/src/ids/id.rs b/grc20-core/src/ids/id.rs index f1864f5..edf0413 100644 --- a/grc20-core/src/ids/id.rs +++ b/grc20-core/src/ids/id.rs @@ -43,11 +43,11 @@ pub fn create_merged_version_id(merged_version_ids: Vec<&str>) -> String { } pub fn create_version_id(space_id: &str, proposal_id: &str) -> String { - create_id_from_unique_string(format!("{}:{}", space_id, proposal_id)) + create_id_from_unique_string(format!("{space_id}:{proposal_id}")) } pub fn create_version_id_from_block(space_id: &str, block: u64) -> String { - create_id_from_unique_string(format!("{}:{}", space_id, block)) + create_id_from_unique_string(format!("{space_id}:{block}")) } /** @@ -56,7 +56,7 @@ pub fn create_version_id_from_block(space_id: &str, block: u64) -> String { * the new one that they're creating. */ pub fn create_space_id(network: &str, address: &str) -> String { - create_id_from_unique_string(format!("{}:{}", network, address)) + create_id_from_unique_string(format!("{network}:{address}")) } pub fn create_id_from_unique_string(text: impl Into) -> String { diff --git a/grc20-core/src/mapping/entity/mod.rs b/grc20-core/src/mapping/entity/mod.rs index 3e67560..93b8efb 100644 --- a/grc20-core/src/mapping/entity/mod.rs +++ b/grc20-core/src/mapping/entity/mod.rs @@ -6,6 +6,7 @@ pub mod find_path; pub mod insert_many; pub mod insert_one; pub mod models; +pub mod search_with_traversals; pub mod semantic_search; pub mod utils; @@ -15,6 +16,7 @@ pub use find_one::FindOneQuery; pub use find_path::FindPathQuery; pub use insert_one::InsertOneQuery; pub use models::{Entity, EntityNode, EntityNodeRef, SystemProperties}; +pub use search_with_traversals::SearchWithTraversals; pub use semantic_search::SemanticSearchQuery; pub use utils::{EntityFilter, EntityRelationFilter, TypesFilter}; @@ -130,6 +132,13 @@ pub fn search(neo4j: &neo4rs::Graph, vector: Vec) -> SemanticSearchQuery SemanticSearchQuery::new(neo4j, vector) } +pub fn search_from_restictions( + neo4j: &neo4rs::Graph, + vector: Vec, +) -> SearchWithTraversals { + SearchWithTraversals::new(neo4j, vector) +} + // TODO: add docs for use via GraphQL pub fn find_path(neo4j: &neo4rs::Graph, id1: String, id2: String) -> FindPathQuery { FindPathQuery::new(neo4j, id1, id2) diff --git a/grc20-core/src/mapping/entity/search_with_traversals.rs b/grc20-core/src/mapping/entity/search_with_traversals.rs new file mode 100644 index 0000000..ade6741 --- /dev/null +++ b/grc20-core/src/mapping/entity/search_with_traversals.rs @@ -0,0 +1,206 @@ +use futures::{Stream, StreamExt, TryStreamExt}; + +use crate::{ + entity::utils::MatchEntity, + error::DatabaseError, + mapping::{ + query_utils::VersionFilter, AttributeNode, FromAttributes, PropFilter, QueryBuilder, + QueryStream, Subquery, EFFECTIVE_SEARCH_RATIO, + }, +}; + +use super::{Entity, EntityFilter, EntityNode}; + +pub struct SearchWithTraversals { + neo4j: neo4rs::Graph, + vector: Vec, + filters: Vec, + space_id: Option>, + version: VersionFilter, + limit: usize, + skip: Option, + threshold: f64, + + _marker: std::marker::PhantomData, +} + +impl SearchWithTraversals { + pub fn new(neo4j: &neo4rs::Graph, vector: Vec) -> Self { + Self { + neo4j: neo4j.clone(), + vector, + filters: Vec::new(), + space_id: None, + version: VersionFilter::default(), + limit: 100, + skip: None, + threshold: 0.75, + + _marker: std::marker::PhantomData, + } + } + + pub fn filter(mut self, filter: EntityFilter) -> Self { + self.filters.push(filter); + self + } + + pub fn space_id(mut self, filter: PropFilter) -> Self { + self.space_id = Some(filter); + self + } + + pub fn version(mut self, version: impl Into) -> Self { + self.version.version_mut(version.into()); + self + } + + pub fn limit(mut self, limit: usize) -> Self { + self.limit = limit; + self + } + + pub fn limit_opt(mut self, limit: Option) -> Self { + if let Some(limit) = limit { + self.limit = limit; + } + self + } + + pub fn skip(mut self, skip: usize) -> Self { + self.skip = Some(skip); + self + } + + pub fn skip_opt(mut self, skip: Option) -> Self { + self.skip = skip; + self + } + + pub fn threshold(mut self, threshold: f64) -> Self { + if (0.0..=1.0).contains(&threshold) { + self.threshold = threshold + } + self + } + + fn subquery(&self) -> QueryBuilder { + const QUERY: &str = r#" + CALL db.index.vector.queryNodes('vector_index', $limit * $effective_search_ratio, $vector) + YIELD node AS n, score AS score + WHERE score > $threshold + MATCH (e:Entity) -[r:ATTRIBUTE]-> (n) + "#; + + self.filters + .iter() + .fold(QueryBuilder::default().subquery(QUERY), |query, filter| { + query.subquery(filter.subquery("e")) + }) + .limit(self.limit) + .skip_opt(self.skip) + .params("vector", self.vector.clone()) + .params("effective_search_ratio", EFFECTIVE_SEARCH_RATIO) + .params("limit", self.limit as i64) + .params("threshold", self.threshold) + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct SearchWithTraversalsResult { + pub entity: T, +} + +impl QueryStream> for SearchWithTraversals { + async fn send( + self, + ) -> Result< + impl Stream, DatabaseError>>, + DatabaseError, + > { + let query = self.subquery().r#return("DISTINCT e"); + + if cfg!(debug_assertions) || cfg!(test) { + tracing::info!( + "entity_node::SearchWithTraversals:::\n{}\nparams:{:?}", + query.compile(), + query.params() + ); + }; + + #[derive(Debug, serde::Deserialize)] + struct RowResult { + e: EntityNode, + } + + Ok(self + .neo4j + .execute(query.build()) + .await? + .into_stream_as::() + .map_err(DatabaseError::from) + .and_then(|row| async move { Ok(SearchWithTraversalsResult { entity: row.e }) })) + } +} + +impl QueryStream>> + for SearchWithTraversals> +{ + async fn send( + self, + ) -> Result< + impl Stream>, DatabaseError>>, + DatabaseError, + > { + let match_entity = MatchEntity::new(&self.space_id, &self.version); + + let query = self.subquery().with( + vec!["e".to_string()], + match_entity.chain( + "e", + "attrs", + "types", + Some(vec![]), + "RETURN e{.*, attrs: attrs, types: types}", + ), + ); + + if cfg!(debug_assertions) || cfg!(test) { + tracing::info!( + "entity_node::SearchWithTraversals::>:\n{}\nparams:{:?}", + query.compile(), + query.params + ); + }; + + #[derive(Debug, serde::Deserialize)] + struct RowResult { + #[serde(flatten)] + node: EntityNode, + attrs: Vec, + types: Vec, + } + + let stream = self + .neo4j + .execute(query.build()) + .await? + .into_stream_as::() + .map_err(DatabaseError::from) + .map(|row_result| { + row_result.and_then(|row| { + T::from_attributes(row.attrs.into()) + .map(|data| SearchWithTraversalsResult { + entity: Entity { + node: row.node, + attributes: data, + types: row.types.into_iter().map(|t| t.id).collect(), + }, + }) + .map_err(DatabaseError::from) + }) + }); + + Ok(stream) + } +} diff --git a/grc20-core/src/mapping/entity/semantic_search.rs b/grc20-core/src/mapping/entity/semantic_search.rs index f8607f9..22c181a 100644 --- a/grc20-core/src/mapping/entity/semantic_search.rs +++ b/grc20-core/src/mapping/entity/semantic_search.rs @@ -5,7 +5,7 @@ use crate::{ error::DatabaseError, mapping::{ query_utils::VersionFilter, AttributeNode, FromAttributes, PropFilter, QueryBuilder, - QueryStream, Subquery, + QueryStream, Subquery, EFFECTIVE_SEARCH_RATIO, }, }; @@ -112,9 +112,6 @@ pub struct SemanticSearchResult { pub entity: T, pub score: f64, } - -const EFFECTIVE_SEARCH_RATIO: f64 = 10000.0; // Adjust this ratio based on your needs - impl QueryStream> for SemanticSearchQuery { async fn send( self, diff --git a/grc20-core/src/mapping/entity/utils.rs b/grc20-core/src/mapping/entity/utils.rs index 5dc450e..faadd08 100644 --- a/grc20-core/src/mapping/entity/utils.rs +++ b/grc20-core/src/mapping/entity/utils.rs @@ -3,8 +3,8 @@ use rand::distributions::DistString; use crate::{ mapping::{ query_utils::{ - query_builder::{MatchQuery, QueryBuilder, Subquery}, - VersionFilter, + query_builder::{MatchQuery, NamePair, QueryBuilder, Rename, Subquery}, + RelationDirection, VersionFilter, }, AttributeFilter, PropFilter, }, @@ -17,6 +17,8 @@ pub struct EntityFilter { pub(crate) id: Option>, pub(crate) attributes: Vec, pub(crate) relations: Option, + /// traverse relation now in entity directly but eventually for modularity will be standalone to be chained + pub(crate) traverse_relation: Option, /// Used to check if the entity exists in the space (i.e.: the entity /// has at least one attribute in the space). pub(crate) space_id: Option>, @@ -51,6 +53,11 @@ impl EntityFilter { self } + pub fn traverse_relation(mut self, traverse_relation: impl Into) -> Self { + self.traverse_relation = Some(traverse_relation.into()); + self + } + /// Used to check if the entity exists in the space. pub fn space_id(mut self, space_id: PropFilter) -> Self { self.space_id = Some(space_id.clone()); @@ -79,6 +86,12 @@ impl EntityFilter { .as_ref() .map(|relations| relations.subquery(&node_var)), ) + // Apply relation traversal + .subquery_opt( + self.traverse_relation + .as_ref() + .map(|traverse| traverse.subquery(&node_var)), + ) } } @@ -134,7 +147,7 @@ impl EntityRelationFilter { let node_var = node_var.into(); let random_suffix: String = rand::distributions::Alphanumeric.sample_string(&mut rand::thread_rng(), 4); - let rel_edge_var = format!("r_{node_var}_{}", random_suffix); + let rel_edge_var = format!("r_{node_var}_{random_suffix}"); let to_node_var = format!("r_{node_var}_to"); MatchQuery::new(format!( @@ -163,6 +176,90 @@ impl EntityRelationFilter { } } +/// Filter used to: +/// - Traverse to inbound or outbound relation +#[derive(Clone, Debug, Default)] +pub struct TraverseRelation { + relation_type_id: Option>, + destination_id: Option>, + direction: RelationDirection, + space_id: Option>, + version: VersionFilter, +} + +impl TraverseRelation { + pub fn relation_type_id(mut self, relation_type_id: impl Into>) -> Self { + self.relation_type_id = Some(relation_type_id.into()); + self + } + + pub fn destination_id(mut self, destination_id: impl Into>) -> Self { + self.destination_id = Some(destination_id.into()); + self + } + + pub fn direction(mut self, direction: RelationDirection) -> Self { + self.direction = direction; + self + } + + pub fn space_id(mut self, space_id: impl Into>) -> Self { + self.space_id = Some(space_id.into()); + self + } + + pub fn version(mut self, version: impl Into) -> Self { + self.version.version_mut(version.into()); + self + } + + pub fn is_empty(&self) -> bool { + self.relation_type_id.is_none() && self.destination_id.is_none() + } + + pub(crate) fn subquery(&self, node_var: impl Into) -> MatchQuery { + let node_var_curr = node_var.into(); + let random_suffix: String = + rand::distributions::Alphanumeric.sample_string(&mut rand::thread_rng(), 4); + let rel_edge_var = format!("r_{node_var_curr}_{random_suffix}"); + let node_var_dest = format!("r_{node_var_curr}_{random_suffix}_to"); + + MatchQuery::new(match self.direction { + RelationDirection::From => { + format!("({node_var_curr}) -[{rel_edge_var}:RELATION]-> ({node_var_dest})") + } + RelationDirection::To => { + format!("({node_var_dest}) -[{rel_edge_var}:RELATION]-> ({node_var_curr})") + } + }) + // rename to change direction of relation + .rename(Rename::new(NamePair::new( + node_var_curr.clone(), + node_var_dest.clone(), + ))) + // Apply the version filter to the relation + .r#where(self.version.subquery(&rel_edge_var)) + // Apply the relation_type filter to the relation (if any) + .where_opt( + self.relation_type_id + .as_ref() + .map(|relation_type| relation_type.subquery(&rel_edge_var, "relation_type", None)), + ) + // Apply the from_id filter to the relation (if any) + .where_opt( + self.destination_id + .as_ref() + .map(|dest_id| dest_id.subquery(&node_var_curr, "id", None)), + ) + // Apply the space_id filter to the relation (if any) + .where_opt( + self.space_id + .as_ref() + .map(|space_id| space_id.subquery(&rel_edge_var, "space_id", None)), + ) + } +} + #[derive(Clone, Debug, Default)] pub struct TypesFilter { types_contains: Vec, diff --git a/grc20-core/src/mapping/mod.rs b/grc20-core/src/mapping/mod.rs index aa7d43d..dc9f1fd 100644 --- a/grc20-core/src/mapping/mod.rs +++ b/grc20-core/src/mapping/mod.rs @@ -28,8 +28,10 @@ pub use value::{Options, Value, ValueType}; use crate::{error::DatabaseError, indexer_ids}; +pub const EFFECTIVE_SEARCH_RATIO: f64 = 100000.0; + pub fn new_version_index(block_number: u64, idx: usize) -> String { - format!("{:016}:{:04}", block_number, idx) + format!("{block_number:016}:{idx:04}") } pub async fn get_version_index( diff --git a/grc20-core/src/mapping/query_utils/mod.rs b/grc20-core/src/mapping/query_utils/mod.rs index af743b5..af27e5a 100644 --- a/grc20-core/src/mapping/query_utils/mod.rs +++ b/grc20-core/src/mapping/query_utils/mod.rs @@ -6,12 +6,14 @@ pub mod order_by; pub mod prop_filter; pub mod query_builder; pub mod query_part; +pub mod relation_direction; pub mod version_filter; pub use attributes_filter::AttributeFilter; pub use order_by::{FieldOrderBy, OrderDirection}; pub use prop_filter::PropFilter; pub use query_part::QueryPart; +pub use relation_direction::RelationDirection; pub use version_filter::VersionFilter; pub trait Query: Sized { diff --git a/grc20-core/src/mapping/query_utils/prop_filter.rs b/grc20-core/src/mapping/query_utils/prop_filter.rs index 9eb816f..44d14f4 100644 --- a/grc20-core/src/mapping/query_utils/prop_filter.rs +++ b/grc20-core/src/mapping/query_utils/prop_filter.rs @@ -197,7 +197,7 @@ impl> PropFilter { let expr = expr .map(|e| e.to_string()) - .unwrap_or(format!("{}.`{}`", node_var, key)); + .unwrap_or(format!("{node_var}.`{key}`")); if let Some(value) = &self.value { let param_key = format!("{node_var}_{key}_value"); diff --git a/grc20-core/src/mapping/query_utils/query_builder.rs b/grc20-core/src/mapping/query_utils/query_builder.rs index 16d8564..a5ce4d4 100644 --- a/grc20-core/src/mapping/query_utils/query_builder.rs +++ b/grc20-core/src/mapping/query_utils/query_builder.rs @@ -130,6 +130,7 @@ pub struct MatchQuery { pub(crate) match_clause: String, pub(crate) optional: bool, pub(crate) where_clauses: Vec, + pub(crate) rename: Option, pub(crate) params: HashMap, } @@ -138,6 +139,7 @@ impl MatchQuery { Self { match_clause: match_clause.into(), optional: false, + rename: None, where_clauses: Vec::new(), params: HashMap::new(), } @@ -147,6 +149,7 @@ impl MatchQuery { Self { match_clause: match_clause.into(), optional: true, + rename: None, where_clauses: Vec::new(), params: HashMap::new(), } @@ -157,6 +160,12 @@ impl MatchQuery { self } + pub fn rename(mut self, rename: impl Into) -> Self { + let rename_clause: Rename = rename.into(); + self.rename = Some(rename_clause.name_pair); + self + } + pub fn r#where(mut self, clause: impl Into) -> Self { let where_clause: WhereClause = clause.into(); self.where_clauses.extend(where_clause.clauses); @@ -181,18 +190,24 @@ impl MatchQuery { impl Subquery for MatchQuery { fn statements(&self) -> Vec { - let mut statements = if self.optional { - vec![format!("OPTIONAL MATCH {}", self.match_clause)] - } else { - vec![format!("MATCH {}", self.match_clause)] + let mut statements = Vec::new(); + + if let Some(rename) = &self.rename { + statements.push(format!("WITH {} AS {}", rename.from_name, rename.to_name)) }; + statements.push(if self.optional { + format!("OPTIONAL MATCH {}", self.match_clause) + } else { + format!("MATCH {}", self.match_clause) + }); + match &self.where_clauses.as_slice() { [] => (), [clause, rest @ ..] => { - statements.push(format!("WHERE {}", clause)); + statements.push(format!("WHERE {clause}")); for rest_clause in rest { - statements.push(format!("AND {}", rest_clause)); + statements.push(format!("AND {rest_clause}")); } } } @@ -205,6 +220,67 @@ impl Subquery for MatchQuery { } } +#[derive(Clone, Debug, Default, PartialEq)] +pub struct NamePair { + from_name: String, + to_name: String, +} + +impl NamePair { + pub fn new(from_name: impl Into, to_name: impl Into) -> Self { + Self { + from_name: from_name.into(), + to_name: to_name.into(), + } + } +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Rename { + name_pair: NamePair, +} + +impl Rename { + pub fn new(name_pair: impl Into) -> Self { + Self { + name_pair: name_pair.into(), + } + } +} + +impl Subquery for Rename { + fn statements(&self) -> Vec { + vec![format!( + "{} AS {}", + self.name_pair.from_name, self.name_pair.to_name + )] + } + + fn params(&self) -> HashMap { + HashMap::new() + } +} + +impl Rename { + pub fn name_pair(mut self, name_pair: impl Into) -> Self { + self.name_pair = name_pair.into(); + self + } + + pub fn name_pair_opt(mut self, name_pair: Option>) -> Self { + if let Some(name_pair) = name_pair { + self.name_pair = name_pair.into(); + } + self + } +} + +impl From for Rename { + fn from(rename: NamePair) -> Self { + Self { name_pair: rename } + } +} + #[derive(Clone, Debug, Default, PartialEq)] pub struct WhereClause { pub clauses: Vec, @@ -247,9 +323,9 @@ impl Subquery for WhereClause { match &self.clauses.as_slice() { [] => vec![], [clause, rest @ ..] => { - let mut statements = vec![format!("WHERE {}", clause)]; + let mut statements = vec![format!("WHERE {clause}")]; for rest_clause in rest { - statements.push(format!("AND {}", rest_clause)); + statements.push(format!("AND {rest_clause}")); } statements } diff --git a/grc20-core/src/mapping/query_utils/query_part.rs b/grc20-core/src/mapping/query_utils/query_part.rs index 0c8b7c0..1ffeaa9 100644 --- a/grc20-core/src/mapping/query_utils/query_part.rs +++ b/grc20-core/src/mapping/query_utils/query_part.rs @@ -217,11 +217,11 @@ impl QueryPart { } if let Some(skip) = self.skip { - query.push_str(&format!("SKIP {}\n", skip)); + query.push_str(&format!("SKIP {skip}\n")); } if let Some(limit) = self.limit { - query.push_str(&format!("LIMIT {}\n", limit)); + query.push_str(&format!("LIMIT {limit}\n")); } if let Some((clause, other)) = &self.with_clauses { diff --git a/grc20-core/src/mapping/query_utils/relation_direction.rs b/grc20-core/src/mapping/query_utils/relation_direction.rs new file mode 100644 index 0000000..d1cc35b --- /dev/null +++ b/grc20-core/src/mapping/query_utils/relation_direction.rs @@ -0,0 +1,6 @@ +#[derive(Clone, Debug, Default)] +pub enum RelationDirection { + From, + #[default] + To, +} diff --git a/grc20-core/src/mapping/query_utils/version_filter.rs b/grc20-core/src/mapping/query_utils/version_filter.rs index 615c066..16cc119 100644 --- a/grc20-core/src/mapping/query_utils/version_filter.rs +++ b/grc20-core/src/mapping/query_utils/version_filter.rs @@ -25,7 +25,7 @@ impl VersionFilter { pub fn subquery(&self, var: &str) -> WhereClause { if let Some(version) = &self.version { - let param_key = format!("{}_version", var); + let param_key = format!("{var}_version"); WhereClause::new(format!("{var}.min_version <= ${param_key} AND ({var}.max_version IS NULL OR {var}.max_version > ${param_key})")) .set_param(param_key, version.clone()) diff --git a/grc20-core/src/mapping/triple.rs b/grc20-core/src/mapping/triple.rs index 61d504b..c917c9b 100644 --- a/grc20-core/src/mapping/triple.rs +++ b/grc20-core/src/mapping/triple.rs @@ -5,8 +5,11 @@ use neo4rs::{BoltMap, BoltType}; use serde::Deserialize; use crate::{ - block::BlockMetadata, error::DatabaseError, ids, indexer_ids, - mapping::query_utils::query_builder::Subquery, pb, + block::BlockMetadata, + error::DatabaseError, + ids, indexer_ids, + mapping::{query_utils::query_builder::Subquery, EFFECTIVE_SEARCH_RATIO}, + pb, }; use super::{ @@ -678,8 +681,6 @@ pub struct SemanticSearchResult { pub space_version: String, } -const EFFECTIVE_SEARCH_RATIO: f64 = 10000.0; // Adjust this ratio based on your needs - impl QueryStream for SemanticSearchQuery { async fn send( self, diff --git a/grc20-core/src/mapping/value.rs b/grc20-core/src/mapping/value.rs index 78a10c9..1d279da 100644 --- a/grc20-core/src/mapping/value.rs +++ b/grc20-core/src/mapping/value.rs @@ -253,10 +253,7 @@ impl TryFrom for DateTime { fn try_from(value: Value) -> Result { Ok(DateTime::parse_from_rfc3339(&value.value) .map_err(|e| { - TriplesConversionError::InvalidValue(format!( - "Failed to parse DateTime value: {}", - e - )) + TriplesConversionError::InvalidValue(format!("Failed to parse DateTime value: {e}")) })? .with_timezone(&Utc)) } diff --git a/grc20-macros/src/entity.rs b/grc20-macros/src/entity.rs index 2bddf3e..1cfba9c 100644 --- a/grc20-macros/src/entity.rs +++ b/grc20-macros/src/entity.rs @@ -120,7 +120,7 @@ pub(crate) fn generate_from_attributes_impl(opts: &EntityOpts) -> TokenStream2 { pub(crate) fn generate_builder_impl(opts: &EntityOpts) -> TokenStream2 { let struct_name = &opts.ident; - let builder_name = Ident::new(&format!("{}Builder", struct_name), Span::call_site()); + let builder_name = Ident::new(&format!("{struct_name}Builder"), Span::call_site()); let fields = opts.data.as_ref().take_struct().expect("Expected struct"); let schema_type = opts.schema_type.as_ref().map(|s| quote!(#s)); @@ -136,7 +136,7 @@ pub(crate) fn generate_builder_impl(opts: &EntityOpts) -> TokenStream2 { // Generate setter methods for each field let setter_methods = fields.iter().map(|field| { let field_name = field.ident.as_ref().expect("Expected named field"); - let mut_name = Ident::new(&format!("{}_mut", field_name), Span::call_site()); + let mut_name = Ident::new(&format!("{field_name}_mut"), Span::call_site()); let field_type = &field.ty; // For Option types, we don't need impl Into @@ -346,7 +346,7 @@ pub(crate) fn generate_query_impls(opts: &EntityOpts) -> TokenStream2 { .iter() .zip(field_types.iter()) .map(|(field_name, field_type)| { - let doc_comment = format!("Filter by {}", field_name); + let doc_comment = format!("Filter by {field_name}"); let filter_type = if let syn::Type::Path(type_path) = field_type { if type_path .path diff --git a/grc20-macros/src/relation.rs b/grc20-macros/src/relation.rs index 313b700..dc65d94 100644 --- a/grc20-macros/src/relation.rs +++ b/grc20-macros/src/relation.rs @@ -120,7 +120,7 @@ pub(crate) fn generate_from_attributes_impl(opts: &RelationOpts) -> TokenStream2 pub(crate) fn generate_builder_impl(opts: &RelationOpts) -> TokenStream2 { let struct_name = &opts.ident; - let builder_name = Ident::new(&format!("{}Builder", struct_name), Span::call_site()); + let builder_name = Ident::new(&format!("{struct_name}Builder"), Span::call_site()); let fields = opts.data.as_ref().take_struct().expect("Expected struct"); let relation_type = opts.relation_type.as_ref().map(|s| quote!(#s)); @@ -136,7 +136,7 @@ pub(crate) fn generate_builder_impl(opts: &RelationOpts) -> TokenStream2 { // Generate setter methods for each field let setter_methods = fields.iter().map(|field| { let field_name = field.ident.as_ref().expect("Expected named field"); - let mut_name = Ident::new(&format!("{}_mut", field_name), Span::call_site()); + let mut_name = Ident::new(&format!("{field_name}_mut"), Span::call_site()); let field_type = &field.ty; // For Option types, we don't need impl Into diff --git a/grc20-sdk/src/models/edit.rs b/grc20-sdk/src/models/edit.rs index 62fe140..4e2181a 100644 --- a/grc20-sdk/src/models/edit.rs +++ b/grc20-sdk/src/models/edit.rs @@ -41,7 +41,7 @@ pub struct Edits; impl Edits { pub fn gen_id(space_id: &str, edit_id: &str) -> String { - ids::create_id_from_unique_string(format!("{}:{}", space_id, edit_id)) + ids::create_id_from_unique_string(format!("{space_id}:{edit_id}")) } pub fn new( @@ -70,7 +70,7 @@ pub struct ProposedEdit; impl ProposedEdit { pub fn gen_id(proposal_id: &str, edit_id: &str) -> String { - ids::create_id_from_unique_string(format!("{}:{}", proposal_id, edit_id)) + ids::create_id_from_unique_string(format!("{proposal_id}:{edit_id}")) } pub fn new( diff --git a/grc20-sdk/src/models/proposal.rs b/grc20-sdk/src/models/proposal.rs index 07770c8..2e2678e 100644 --- a/grc20-sdk/src/models/proposal.rs +++ b/grc20-sdk/src/models/proposal.rs @@ -139,7 +139,7 @@ impl TryFrom for ProposalType { pb::ipfs::ActionType::RemoveSubspace => Ok(Self::RemoveSubspace), pb::ipfs::ActionType::ImportSpace => Ok(Self::ImportSpace), pb::ipfs::ActionType::ArchiveSpace => Ok(Self::ArchiveSpace), - _ => Err(format!("Invalid action type: {:?}", action_type)), + _ => Err(format!("Invalid action type: {action_type:?}")), } } } @@ -422,10 +422,7 @@ pub struct ProposedAccount; impl ProposedAccount { pub fn gen_id(proposal_id: &str, account_id: &str) -> String { - ids::create_id_from_unique_string(format!( - "PROPOSED_ACCOUNT:{}:{}", - proposal_id, account_id - )) + ids::create_id_from_unique_string(format!("PROPOSED_ACCOUNT:{proposal_id}:{account_id}")) } pub fn new(proposal_id: &str, account_id: &str) -> Relation { diff --git a/grc20-sdk/src/models/space/parent_spaces_query.rs b/grc20-sdk/src/models/space/parent_spaces_query.rs index 4218703..c720b8f 100644 --- a/grc20-sdk/src/models/space/parent_spaces_query.rs +++ b/grc20-sdk/src/models/space/parent_spaces_query.rs @@ -61,7 +61,7 @@ impl ParentSpacesQuery { indexer_ids::INDEXER_SPACE_ID, )) .subquery("WHERE size(s) = size(COLLECT { WITH s UNWIND s AS _ RETURN DISTINCT _ })") - .subquery_opt(self.max_depth.map(|depth| format!("AND size(s) <= {}", depth))) + .subquery_opt(self.max_depth.map(|depth| format!("AND size(s) <= {depth}"))) .subquery("WITH {space_id: LAST([start] + s).id, depth: SIZE(s)} AS parent_spaces") .limit(self.limit) .skip_opt(self.skip) diff --git a/grc20-sdk/src/models/space/subspaces_query.rs b/grc20-sdk/src/models/space/subspaces_query.rs index d4a4eec..8cfabd3 100644 --- a/grc20-sdk/src/models/space/subspaces_query.rs +++ b/grc20-sdk/src/models/space/subspaces_query.rs @@ -61,7 +61,7 @@ impl SubspacesQuery { indexer_ids::INDEXER_SPACE_ID, )) .subquery("WHERE size(s) = size(COLLECT { WITH s UNWIND s AS _ RETURN DISTINCT _ })") - .subquery_opt(self.max_depth.map(|depth| format!("AND size(s) <= {}", depth))) + .subquery_opt(self.max_depth.map(|depth| format!("AND size(s) <= {depth}"))) .subquery("WITH {space_id: LAST([start] + s).id, depth: SIZE(s)} AS subspaces") .limit(self.limit) .skip_opt(self.skip) diff --git a/grc20-sdk/src/models/vote.rs b/grc20-sdk/src/models/vote.rs index d1b3ee9..8a5f388 100644 --- a/grc20-sdk/src/models/vote.rs +++ b/grc20-sdk/src/models/vote.rs @@ -51,7 +51,7 @@ impl TryFrom for VoteType { match vote { 2 => Ok(Self::Accept), 3 => Ok(Self::Reject), - _ => Err(format!("Invalid vote type: {}", vote)), + _ => Err(format!("Invalid vote type: {vote}")), } } } @@ -73,8 +73,7 @@ impl TryFrom for VoteType { (mapping::ValueType::Text, "ACCEPT") => Ok(Self::Accept), (mapping::ValueType::Text, "REJECT") => Ok(Self::Reject), (value_type, _) => Err(TriplesConversionError::InvalidValue(format!( - "Invalid vote type value_type: {:?}", - value_type + "Invalid vote type value_type: {value_type:?}" ))), } } diff --git a/mcp-server/resources/get_entity_info_description.md b/mcp-server/resources/get_entity_info_description.md new file mode 100644 index 0000000..97bf061 --- /dev/null +++ b/mcp-server/resources/get_entity_info_description.md @@ -0,0 +1,123 @@ +This request allows you to get the detailed information about an Entity with it's ID. You will get the name, description, other attributes, inbound relations and outbound relations of the Entity. + +The id for San Francisco is: 3qayfdjYyPv1dAYf8gPL5r + +ToolCall> get_entity_info("3qayfdjYyPv1dAYf8gPL5r") +ToolResult> +``` +{ + "all_attributes": [ + { + "attribute_name": "Description", + "attribute_value": "A vibrant city known for its iconic Golden Gate Bridge, steep rolling hills, historic cable cars, and a rich cultural tapestry including diverse neighborhoods like the Castro and the Mission District." + }, + { + "attribute_name": "Name", + "attribute_value": "San Francisco" + } + ], + "id": "3qayfdjYyPv1dAYf8gPL5r", + "inbound_relations": [ + { + "id": "NAMA1uDMzBQTvPYV9N92BV", + "name": "SF Mayor Lurie launching police task force to counter crime in core downtown areas", + "relation_id": "8ESicJHiNJ28VGL5u34A5q", + "relation_type": "Related spaces" + }, + { + "id": "6wAoNdGVbweKi2JRPZP4bX", + "name": "San Francisco Independent Film Festival", + "relation_id": "TH5Tu5Y5nacvREvAQRvcR2", + "relation_type": "Related spaces" + }, + { + "id": "8VCHYDURDStwuTCUBjWLQa", + "name": "Product Engineer at Geo", + "relation_id": "KPTqdNpCusxfM37KbKPX8w", + "relation_type": "Related spaces" + }, + { + "id": "NcQ3h9jeJSavVd8iFsUxvD", + "name": "Senior Civil Engineer @ Golden Gate Bridge, Highway & Transportation District", + "relation_id": "AqpNtJ3XxaY4fqRCyoXbdt", + "relation_type": "Cities" + }, + { + "id": "4ojV4dS1pV2tRnzXTpcMKJ", + "name": "Senior Plan Check Engineer (FT - Hybrid) @ CSG Consultants, Inc.", + "relation_id": "3AX4j43nywT5eBRV3s6AXi", + "relation_type": "Cities" + }, + { + "id": "QoakYWCuv85FVuYdSmonxr", + "name": "Senior Civil Engineer - Land Development (FT - Hybrid) @ CSG Consultants, Inc.", + "relation_id": "8GEF1i3LK4Z56THjE8dVku", + "relation_type": "Cities" + }, + { + "id": "JuV7jLoypebzLhkma6oZoU", + "name": "Lead Django Backend Engineer @ Textme Inc", + "relation_id": "46aBsQyBq15DimJ2i1DX4a", + "relation_type": "Cities" + }, + { + "id": "RTmcYhLVmmfgUn9L3D1J3y", + "name": "Chief Engineer @ Wyndham Hotels & Resorts", + "relation_id": "8uYxjzkkdjskDQAeTQomvc", + "relation_type": "Cities" + } + ], + "outbound_relations": [ + { + "id": "CUoEazCD7EmzXPTFFY8gGY", + "name": "No name", + "relation_id": "5WeSkkE1XXvGJGmXj9VUQ8", + "relation_type": "Cover" + }, + { + "id": "7gzF671tq5JTZ13naG4tnr", + "name": "Space", + "relation_id": "WUZCXE1UGRtxdNQpGug8Tf", + "relation_type": "Types" + }, + { + "id": "D6Wy4bdtdoUrG3PDZceHr", + "name": "City", + "relation_id": "ARMj8fjJtdCwbtZa1f3jwe", + "relation_type": "Types" + }, + { + "id": "AhidiWYnQ8fAbHqfzdU74k", + "name": "Upcoming events", + "relation_id": "V1ikGW9riu7dAP8rMgZq3u", + "relation_type": "Blocks" + }, + { + "id": "T6iKbwZ17iv4dRdR9Qw7qV", + "name": "Trending restaurants", + "relation_id": "CvGXCmGXE7ofsgZeWad28p", + "relation_type": "Blocks" + }, + { + "id": "X18WRE36mjwQ7gu3LKaLJS", + "name": "Neighborhoods", + "relation_id": "Uxpsee9LoTgJqMFfAQyJP6", + "relation_type": "Blocks" + }, + { + "id": "HeC2pygci2tnvjTt5aEnBV", + "name": "Top goals", + "relation_id": "5WMTAzCnZH9Bsevou9GQ3K", + "relation_type": "Blocks" + }, + { + "id": "5YtYFsnWq1jupvh5AjM2ni", + "name": "Culture", + "relation_id": "5TmxfepRr1THMRkGWenj5G", + "relation_type": "Tabs" + } + ] +} +``` + +Any of the given field can be further queried by using get_entity_info with that id since all information in the Knowledge Graph(KG) is an Entity. diff --git a/mcp-server/resources/get_properties_description.md b/mcp-server/resources/get_properties_description.md deleted file mode 100644 index e69de29..0000000 diff --git a/mcp-server/resources/get_relations_between_entities_description.md b/mcp-server/resources/get_relations_between_entities_description.md new file mode 100644 index 0000000..70b0160 --- /dev/null +++ b/mcp-server/resources/get_relations_between_entities_description.md @@ -0,0 +1,34 @@ +This request allows you to find the direct and distant relationships between 2 entities by their ID. + +Crypto Briefing (a Crypto company) id: 9xRruQhSfAuJjHwKnvTjma +Bullish (a Crypto company) id: ESShPFkqfFnDzYkSwGGVuR + +ToolCall> search_types("9xRruQhSfAuJjHwKnvTjma", "ESShPFkqfFnDzYkSwGGVuR") +ToolResult> +``` +[ + { + "nodes": [ + "Crypto Briefing", + "Crypto", + "Bullish" + ], + "relations": [ + "Related spaces", + "Related spaces" + ] + }, + { + "nodes": [ + "Crypto Briefing", + "Featured", + "Bullish" + ], + "relations": [ + "Tags", + "Tags" + ] + } +] +``` + diff --git a/mcp-server/resources/instructions.md b/mcp-server/resources/instructions.md index e251cd8..9aeec45 100644 --- a/mcp-server/resources/instructions.md +++ b/mcp-server/resources/instructions.md @@ -1,13 +1,186 @@ -This server provides tools to query the Knowledge Graph (KG), a database of wide-ranging structured information (similar to wikidata). The KG organizes information using entities and relations. Entities can have 0, 1 or many types, while relations have exactly one relation type. Both entities and relations can have properties. +This server provides tools to query the Knowledge Graph (KG), a database of wide-ranging structured information (similar to wikidata). The KG organizes information using entities and relations. -Importantly, types, relation types and properties are themselves entities that can be queried. In other words, the KG contains both the property graph of the data as well as the data itself! +You should use it for every request to get the informations for your answers since it covers a wide range of subject like internet would. -The tools defined in the MCP server are made to be used in combination with each other. All except the most trivial user requests will require the use of multiple tools. +The tools defined in the MCP server are made to be used in combination with each other. All except the most trivial requests will require the use of multiple tools. Here is an example: -User> What are the properties of the Person type? +User> Can you give me information about San Francisco? -ToolCall> search_type("person") +ToolCall> search_entity({"query": "San Francisco"}) ToolResult> ``` -``` \ No newline at end of file +{ + "entities": [ + { + "description": "A vibrant city known for its iconic Golden Gate Bridge, steep rolling hills, historic cable cars, and a rich cultural tapestry including diverse neighborhoods like the Castro and the Mission District.", + "id": "3qayfdjYyPv1dAYf8gPL5r", + "name": "San Francisco" + }, + { + "description": null, + "id": "W5ZEpuy3Tij1XSXtJLruQ5", + "name": "SF Bay Area" + }, + { + "description": null, + "id": "RHoJT3hNVaw7m5fLLtZ8WQ", + "name": "California" + }, + { + "description": null, + "id": "Sh1qtjr4i92ZD6YGPeu5a2", + "name": "Abundant housing in San Francisco" + }, + { + "description": null, + "id": "UqLf9fTVKHkDs3LzP9zHpH", + "name": "Public safety in San Francisco" + }, + { + "description": null, + "id": "BeyiZ6oLqLMaSXiG41Yxtf", + "name": "City" + }, + { + "description": null, + "id": "D6Wy4bdtdoUrG3PDZceHr", + "name": "City" + }, + { + "description": null, + "id": "JWVrgUXmjS75PqNX2hry5q", + "name": "Clean streets in San Francisco" + }, + { + "description": null, + "id": "DcA2c7ooFTgEdtaRcaj7Z1", + "name": "Revitalizing downtown San Francisco" + }, + { + "description": null, + "id": "KWBLj9czHBBmYUT98rnxVM", + "name": "Location" + } + ] +} +``` +Let's get more info about San Francisco (id: 3qayfdjYyPv1dAYf8gPL5r) + +ToolCall> get_entity_info("3qayfdjYyPv1dAYf8gPL5r") +ToolResult> +``` +{ + "all_attributes": [ + { + "attribute_name": "Description", + "attribute_value": "A vibrant city known for its iconic Golden Gate Bridge, steep rolling hills, historic cable cars, and a rich cultural tapestry including diverse neighborhoods like the Castro and the Mission District." + }, + { + "attribute_name": "Name", + "attribute_value": "San Francisco" + } + ], + "id": "3qayfdjYyPv1dAYf8gPL5r", + "inbound_relations": [ + { + "id": "NAMA1uDMzBQTvPYV9N92BV", + "name": "SF Mayor Lurie launching police task force to counter crime in core downtown areas", + "relation_id": "8ESicJHiNJ28VGL5u34A5q", + "relation_type": "Related spaces" + }, + { + "id": "6wAoNdGVbweKi2JRPZP4bX", + "name": "San Francisco Independent Film Festival", + "relation_id": "TH5Tu5Y5nacvREvAQRvcR2", + "relation_type": "Related spaces" + }, + { + "id": "8VCHYDURDStwuTCUBjWLQa", + "name": "Product Engineer at Geo", + "relation_id": "KPTqdNpCusxfM37KbKPX8w", + "relation_type": "Related spaces" + }, + { + "id": "NcQ3h9jeJSavVd8iFsUxvD", + "name": "Senior Civil Engineer @ Golden Gate Bridge, Highway & Transportation District", + "relation_id": "AqpNtJ3XxaY4fqRCyoXbdt", + "relation_type": "Cities" + }, + { + "id": "4ojV4dS1pV2tRnzXTpcMKJ", + "name": "Senior Plan Check Engineer (FT - Hybrid) @ CSG Consultants, Inc.", + "relation_id": "3AX4j43nywT5eBRV3s6AXi", + "relation_type": "Cities" + }, + { + "id": "QoakYWCuv85FVuYdSmonxr", + "name": "Senior Civil Engineer - Land Development (FT - Hybrid) @ CSG Consultants, Inc.", + "relation_id": "8GEF1i3LK4Z56THjE8dVku", + "relation_type": "Cities" + }, + { + "id": "JuV7jLoypebzLhkma6oZoU", + "name": "Lead Django Backend Engineer @ Textme Inc", + "relation_id": "46aBsQyBq15DimJ2i1DX4a", + "relation_type": "Cities" + }, + { + "id": "RTmcYhLVmmfgUn9L3D1J3y", + "name": "Chief Engineer @ Wyndham Hotels & Resorts", + "relation_id": "8uYxjzkkdjskDQAeTQomvc", + "relation_type": "Cities" + } + ], + "outbound_relations": [ + { + "id": "CUoEazCD7EmzXPTFFY8gGY", + "name": "No name", + "relation_id": "5WeSkkE1XXvGJGmXj9VUQ8", + "relation_type": "Cover" + }, + { + "id": "7gzF671tq5JTZ13naG4tnr", + "name": "Space", + "relation_id": "WUZCXE1UGRtxdNQpGug8Tf", + "relation_type": "Types" + }, + { + "id": "D6Wy4bdtdoUrG3PDZceHr", + "name": "City", + "relation_id": "ARMj8fjJtdCwbtZa1f3jwe", + "relation_type": "Types" + }, + { + "id": "AhidiWYnQ8fAbHqfzdU74k", + "name": "Upcoming events", + "relation_id": "V1ikGW9riu7dAP8rMgZq3u", + "relation_type": "Blocks" + }, + { + "id": "T6iKbwZ17iv4dRdR9Qw7qV", + "name": "Trending restaurants", + "relation_id": "CvGXCmGXE7ofsgZeWad28p", + "relation_type": "Blocks" + }, + { + "id": "X18WRE36mjwQ7gu3LKaLJS", + "name": "Neighborhoods", + "relation_id": "Uxpsee9LoTgJqMFfAQyJP6", + "relation_type": "Blocks" + }, + { + "id": "HeC2pygci2tnvjTt5aEnBV", + "name": "Top goals", + "relation_id": "5WMTAzCnZH9Bsevou9GQ3K", + "relation_type": "Blocks" + }, + { + "id": "5YtYFsnWq1jupvh5AjM2ni", + "name": "Culture", + "relation_id": "5TmxfepRr1THMRkGWenj5G", + "relation_type": "Tabs" + } + ] +} +``` diff --git a/mcp-server/resources/name_search_entity_description.md b/mcp-server/resources/name_search_entity_description.md new file mode 100644 index 0000000..8377684 --- /dev/null +++ b/mcp-server/resources/name_search_entity_description.md @@ -0,0 +1,81 @@ +This request allows you to get Entities from a name/description search and traversal from that query by using relation name. + +Example Query: Find employees that works at The Graph. + +ToolCall> +``` +name_search_entity( + { + "query": "The Graph", + "traversal_filter": { + "relation_type_id": "Works at", + "direction": "From" + } + } +) +``` + +ToolResult> +``` +{ + "entities": [ + { + "description": "Founder & CEO of Geo. Cofounder of The Graph, Edge & Node, House of Web3. Building a vibrant decentralized future.", + "id": "9HsfMWYHr9suYdMrtssqiX", + "name": "Yaniv Tal" + }, + { + "description": "Developer Relations Engineer", + "id": "22MGz47c9WHtRiHuSEPkcG", + "name": "Kevin Jones" + }, + { + "description": "Description will go here", + "id": "JYTfEcdmdjiNzBg469gE83", + "name": "Pedro Diogo" + } + ] +} +``` + +Example Query: Find all the articles written by employees that works at The Graph. + +ToolCall> +``` +name_search_entity( + { + "query": "The Graph", + "traversal_filter": { + "relation_type_id": "Works at", + "direction": "From", + "traversal_filter": { + "relation_type_id": "Author", + "direction": "From" + } + } + } +) +``` + +ToolResult> +``` +{ + "entities": [ + { + "description": "A fresh look at what web3 is and what the missing pieces have been for making it a reality.", + "id": "XYo6aR3VqFQSEcf6AeTikW", + "name": "Knowledge graphs are web3" + }, + { + "description": "A new standard is here for structuring knowledge. GRC-20 will reshape how we make applications composable and redefine web3.", + "id": "5FkVvS4mTz6Ge7wHkAUMRk", + "name": "Introducing GRC-20: A knowledge graph standard for web3" + }, + { + "description": "How do you know what is true? Who do you trust? Everybody has a point of view, but no one is an authority. As humanity we need a way to aggregate our knowledge into something we can trust. We need a system.", + "id": "5WHP8BuoCdSiqtfy87SYWG", + "name": "Governing public knowledge" + } + ] +} +``` diff --git a/mcp-server/resources/search_entity_description.md b/mcp-server/resources/search_entity_description.md index e69de29..6ba6e33 100644 --- a/mcp-server/resources/search_entity_description.md +++ b/mcp-server/resources/search_entity_description.md @@ -0,0 +1,104 @@ +This request allows you to get the Entities from a name/description search and traversal from that query if needed. + + +Example Query: Can you give me information about San Francisco? + +ToolCall> +``` +search_entity({ +"query": "San Francisco" +}) +``` +Tool Result> +``` +{ + "entities": [ + { + "description": "A vibrant city known for its iconic Golden Gate Bridge, steep rolling hills, historic cable cars, and a rich cultural tapestry including diverse neighborhoods like the Castro and the Mission District.", + "id": "3qayfdjYyPv1dAYf8gPL5r", + "name": "San Francisco" + }, + { + "description": null, + "id": "W5ZEpuy3Tij1XSXtJLruQ5", + "name": "SF Bay Area" + }, + { + "description": null, + "id": "RHoJT3hNVaw7m5fLLtZ8WQ", + "name": "California" + }, + { + "description": null, + "id": "Sh1qtjr4i92ZD6YGPeu5a2", + "name": "Abundant housing in San Francisco" + }, + { + "description": null, + "id": "UqLf9fTVKHkDs3LzP9zHpH", + "name": "Public safety in San Francisco" + }, + { + "description": null, + "id": "BeyiZ6oLqLMaSXiG41Yxtf", + "name": "City" + }, + { + "description": null, + "id": "D6Wy4bdtdoUrG3PDZceHr", + "name": "City" + }, + { + "description": null, + "id": "JWVrgUXmjS75PqNX2hry5q", + "name": "Clean streets in San Francisco" + }, + { + "description": null, + "id": "DcA2c7ooFTgEdtaRcaj7Z1", + "name": "Revitalizing downtown San Francisco" + }, + { + "description": null, + "id": "KWBLj9czHBBmYUT98rnxVM", + "name": "Location" + } + ] +} +``` + +Another Query: Give me the employees that work at The Graph? + +Work_at id: U1uCAzXsRSTP4vFwo1JwJG +ToolCall> +``` +search_entity({ +"query": "The Graph", +"traversal_filter": { + "relation_type_id": "U1uCAzXsRSTP4vFwo1JwJG", + "direction": "From" +} +}) +``` +ToolResult> +``` +{ + "entities": [ + { + "description": "Founder & CEO of Geo. Cofounder of The Graph, Edge & Node, House of Web3. Building a vibrant decentralized future.", + "id": "9HsfMWYHr9suYdMrtssqiX", + "name": "Yaniv Tal" + }, + { + "description": "Developer Relations Engineer", + "id": "22MGz47c9WHtRiHuSEPkcG", + "name": "Kevin Jones" + }, + { + "description": "Description will go here", + "id": "JYTfEcdmdjiNzBg469gE83", + "name": "Pedro Diogo" + } + ] +} +``` diff --git a/mcp-server/resources/search_properties_description.md b/mcp-server/resources/search_properties_description.md new file mode 100644 index 0000000..7b1c41f --- /dev/null +++ b/mcp-server/resources/search_properties_description.md @@ -0,0 +1,25 @@ +This request allows you to search by name for the ATTRIBUTES (properties) that can be used to describe an Entity. + + +ToolCall> search_properties("Authors") +ToolResult> +``` +[ + [ + { + "attribute_name": "Name", + "attribute_value": "Authors", + "entity_id": "JzFpgguvcCaKhbQYPHsrNT" + } + ], + [ + { + "attribute_name": "Name", + "attribute_value": "Owners", + "entity_id": "RwDfM3vUvyLwSNYv6sWhc9" + } + ] +] +``` + +Since all the Relations are also of the type Entity. they can be queried by their id for more information. diff --git a/mcp-server/resources/search_relation_type_description.md b/mcp-server/resources/search_relation_type_description.md index e69de29..02134b4 100644 --- a/mcp-server/resources/search_relation_type_description.md +++ b/mcp-server/resources/search_relation_type_description.md @@ -0,0 +1,44 @@ +This request allows you to search by name for information of the relations between entities in the Knowledge Graph like works at. + +ToolCall> search_relation_types("works at") +ToolResult> +``` +[ + [ + { + "attribute_name": "Name", + "attribute_value": "Works at", + "entity_id": "U1uCAzXsRSTP4vFwo1JwJG" + }, + { + "attribute_name": "Is type property", + "attribute_value": "0", + "entity_id": "U1uCAzXsRSTP4vFwo1JwJG" + } + ], + [ + { + "attribute_name": "Name", + "attribute_value": "Worked at", + "entity_id": "8fvqALeBDwEExJsDeTcvnV" + }, + { + "attribute_name": "Is type property", + "attribute_value": "0", + "entity_id": "8fvqALeBDwEExJsDeTcvnV" + }, + { + "attribute_name": "Name", + "attribute_value": "Worked at", + "entity_id": "8fvqALeBDwEExJsDeTcvnV" + }, + { + "attribute_name": "Description", + "attribute_value": "A project that someone worked at in the past. Details about the role can be added as properties on the relation.", + "entity_id": "8fvqALeBDwEExJsDeTcvnV" + } + ] +] +``` + +Since all the relation types are also of the type Entity. they can be queried by their id for more information. diff --git a/mcp-server/resources/search_space_description.md b/mcp-server/resources/search_space_description.md new file mode 100644 index 0000000..3224922 --- /dev/null +++ b/mcp-server/resources/search_space_description.md @@ -0,0 +1,40 @@ +This request allows you to find a Space from it's name or description. The spaces are where the attributes and relations are and may be useful to specify when querying entities and relations. + +ToolCall> +``` +search_space("San Francisco") +``` + +ToolResult> +``` +[ + [ + { + "attribute_name": "Description", + "attribute_value": "A vibrant city known for its iconic Golden Gate Bridge, steep rolling hills, historic cable cars, and a rich cultural tapestry including diverse neighborhoods like the Castro and the Mission District.", + "entity_id": "3qayfdjYyPv1dAYf8gPL5r" + }, + { + "attribute_name": "Name", + "attribute_value": "San Francisco", + "entity_id": "3qayfdjYyPv1dAYf8gPL5r" + } + ], + [ + { + "attribute_name": "Name", + "attribute_value": "SF Bay Area", + "entity_id": "W5ZEpuy3Tij1XSXtJLruQ5" + } + ], + [ + { + "attribute_name": "Name", + "attribute_value": "California", + "entity_id": "RHoJT3hNVaw7m5fLLtZ8WQ" + } + ] +] +``` + +Eventually, space will be used to narrow research or help format result diff --git a/mcp-server/resources/search_type_description.md b/mcp-server/resources/search_type_description.md index e69de29..478367f 100644 --- a/mcp-server/resources/search_type_description.md +++ b/mcp-server/resources/search_type_description.md @@ -0,0 +1,35 @@ +This request allows you to search by name for a basic type of the Knowledge Graph(KG) like Person or Event. This will give back the type with it's name, id and description. + +ToolCall> search_type("University") +ToolResult> +``` +[ + [ + { + "attribute_name": "Description", + "attribute_value": "An institution of higher education offering undergraduate and graduate degrees, research opportunities, and specialized academic programs.", + "entity_id": "L8iozarUyS8bkcUiS6kPqV" + }, + { + "attribute_name": "Name", + "attribute_value": "University", + "entity_id": "L8iozarUyS8bkcUiS6kPqV" + } + ], + [ + { + "attribute_name": "Description", + "attribute_value": "An educational institution where students acquire knowledge, skills, and credentials through structured learning programs.", + "entity_id": "M89C7wwdJVaCW9rAVQpJbY" + }, + { + "attribute_name": "Name", + "attribute_value": "School", + "entity_id": "M89C7wwdJVaCW9rAVQpJbY" + } + ] +] +``` + + +Since all the types are also of the type Entity. they can be queried by their id for more information. diff --git a/mcp-server/src/input_types.rs b/mcp-server/src/input_types.rs new file mode 100644 index 0000000..4fa59af --- /dev/null +++ b/mcp-server/src/input_types.rs @@ -0,0 +1,50 @@ +#[derive(Debug, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] +pub struct SearchTraversalInputFilter { + pub query: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub traversal_filter: Option, +} + +#[derive(Debug, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] +pub struct TraversalFilter { + pub direction: RelationDirection, + pub relation_type_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub traversal_filter: Option>, +} + +#[derive(Debug, serde::Deserialize, serde::Serialize, schemars::JsonSchema, Clone)] +pub enum RelationDirection { + From, + To, +} + +/// Struct returned by call to `OneOrMany::into_iter()`. +pub struct IntoIter { + // Owned. + next_filter: Option, +} + +/// Implement `IntoIterator` for `TraversalFilter`. +impl IntoIterator for TraversalFilter { + type Item = TraversalFilter; + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter { + next_filter: Some(self), + } + } +} + +/// Implement `Iterator` for `IntoIter`. +impl Iterator for IntoIter { + type Item = TraversalFilter; + + fn next(&mut self) -> Option { + self.next_filter.take().map(|mut current| { + self.next_filter = current.traversal_filter.take().map(|boxed| *boxed); + current + }) + } +} diff --git a/mcp-server/src/lib.rs b/mcp-server/src/lib.rs new file mode 100644 index 0000000..f728dbe --- /dev/null +++ b/mcp-server/src/lib.rs @@ -0,0 +1 @@ +pub mod input_types; diff --git a/mcp-server/src/main.rs b/mcp-server/src/main.rs index 19267da..3147abe 100644 --- a/mcp-server/src/main.rs +++ b/mcp-server/src/main.rs @@ -1,12 +1,21 @@ use clap::{Args, Parser}; use fastembed::{EmbeddingModel, InitOptions, TextEmbedding}; -use futures::{TryStreamExt, future::join_all}; +use futures::{TryStreamExt, future::join_all, pin_mut}; use grc20_core::{ - entity::{self, Entity, EntityFilter, EntityNode, EntityRelationFilter, TypesFilter}, - mapping::{Attributes, Query, QueryStream, RelationEdge, prop_filter}, - neo4rs, relation, system_ids, + entity::{ + self, Entity, EntityFilter, EntityNode, EntityRelationFilter, utils::TraverseRelation, + }, + mapping::{ + Query, QueryStream, RelationEdge, Triple, prop_filter, + query_utils::RelationDirection, + triple::{self, SemanticSearchResult}, + }, + neo4rs, + relation::{self, RelationFilter}, + system_ids, }; use grc20_sdk::models::BaseEntity; +use mcp_server::input_types::{self, SearchTraversalInputFilter}; use rmcp::{ Error as McpError, RoleServer, ServerHandler, model::*, @@ -14,8 +23,8 @@ use rmcp::{ tool, transport::sse_server::{SseServer, SseServerConfig}, }; -use serde_json::json; -use std::sync::Arc; +use serde_json::{Value, json}; +use std::{collections::HashSet, sync::Arc}; use tracing_subscriber::{ layer::SubscriberExt, util::SubscriberInitExt, @@ -77,12 +86,6 @@ async fn main() -> anyhow::Result<()> { Ok(()) } -#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] -pub struct StructRequest { - pub a: i32, - pub b: i32, -} - const EMBEDDING_MODEL: EmbeddingModel = EmbeddingModel::AllMiniLML6V2; #[derive(Clone)] @@ -110,13 +113,11 @@ impl KnowledgeGraph { RawResource::new(uri, name.to_string()).no_annotation() } - #[tool(description = "Search Types")] - async fn search_types( + async fn search( &self, - #[tool(param)] - #[schemars(description = "The query string to search for types")] query: String, - ) -> Result { + limit: Option, + ) -> Result, McpError> { let embedding = self .embedding_model .embed(vec![&query], None) @@ -127,12 +128,10 @@ impl KnowledgeGraph { .map(|v| v as f64) .collect::>(); - let results = entity::search::>(&self.neo4j, embedding) - .filter( - entity::EntityFilter::default() - .relations(TypesFilter::default().r#type(system_ids::SCHEMA_TYPE)), - ) - .limit(10) + let limit = limit.unwrap_or(10); + + let semantic_search_triples = triple::search(&self.neo4j, embedding) + .limit(limit) .send() .await .map_err(|e| { @@ -149,35 +148,315 @@ impl KnowledgeGraph { Some(json!({ "error": e.to_string() })), ) })?; + Ok(semantic_search_triples) + } + + async fn get_ids_from_search( + &self, + search_triples: Vec, + create_relation_filter: impl Fn(SemanticSearchResult) -> RelationFilter, + ) -> Result, McpError> { + let mut seen_ids: HashSet = HashSet::new(); + let mut result_ids: Vec = Vec::new(); + + for semantic_search_triple in search_triples { + let filtered_for_types = relation::find_many::>(&self.neo4j) + .filter(create_relation_filter(semantic_search_triple)) + .send() + .await; + + //We only need to get the first relation since they would share the same entity id + if let Ok(stream) = filtered_for_types { + pin_mut!(stream); + if let Some(edge) = stream.try_next().await.ok().flatten() { + let id = edge.from.id; + if seen_ids.insert(id.clone()) { + result_ids.push(id); + } + } + } + } + Ok(result_ids) + } - tracing::info!("Found {} results for query '{}'", results.len(), query); + async fn format_triples_detailled( + &self, + triples: Result, ErrorData>, + ) -> Vec { + if let Ok(triples) = triples { + join_all(triples.into_iter().map(|triple| async move {json!({ + "entity_id": triple.entity, + "attribute_name": self.get_name_of_id(triple.attribute).await.unwrap_or("No attribute name".to_string()), + "attribute_value": String::try_from(triple.value).unwrap_or("No value".to_string()) + })})).await.to_vec() + } else { + Vec::new() + } + } + + #[tool(description = include_str!("../resources/search_type_description.md"))] + async fn search_types( + &self, + #[tool(param)] + #[schemars(description = "The query string to search for types")] + query: String, + ) -> Result { + let semantic_search_triples = self.search(query, Some(10)).await.unwrap_or_default(); + + let create_relation_filter = |search_result: SemanticSearchResult| { + RelationFilter::default() + .from_(EntityFilter::default().id(prop_filter::value(search_result.triple.entity))) + .relation_type( + EntityFilter::default().id(prop_filter::value(system_ids::TYPES_ATTRIBUTE)), + ) + .to_(EntityFilter::default().id(prop_filter::value(system_ids::SCHEMA_TYPE))) + }; + + let result_types = self + .get_ids_from_search(semantic_search_triples, &create_relation_filter) + .await + .unwrap_or_default(); + + let entities: Vec, McpError>> = + join_all(result_types.into_iter().map(|id| async { + triple::find_many(&self.neo4j) + .entity_id(prop_filter::value(id)) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + })? + .try_collect::>() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + }) + })) + .await + .to_vec(); Ok(CallToolResult::success( - results - .into_iter() - .map(|result| { - Content::json(json!({ - "id": result.entity.id(), - "name": result.entity.attributes.name, - "description": result.entity.attributes.description, - "types": result.entity.types, - })) - .expect("Failed to create JSON content") - }) - .collect(), + join_all( + entities + .into_iter() + .map(|result: Result, _>| async { + Content::json(self.format_triples_detailled(result).await) + .expect("Failed to create JSON content") + }), + ) + .await + .to_vec(), )) } - #[tool(description = "Search Relation Types")] + #[tool(description = include_str!("../resources/search_relation_type_description.md"))] async fn search_relation_types( &self, #[tool(param)] #[schemars(description = "The query string to search for relation types")] query: String, ) -> Result { + let semantic_search_triples = self.search(query, Some(10)).await.unwrap_or_default(); + + let create_relation_filter = |search_result: SemanticSearchResult| { + RelationFilter::default() + .from_(EntityFilter::default().id(prop_filter::value(search_result.triple.entity))) + .relation_type( + EntityFilter::default() + .id(prop_filter::value(system_ids::VALUE_TYPE_ATTRIBUTE)), + ) + .to_( + EntityFilter::default() + .id(prop_filter::value(system_ids::RELATION_SCHEMA_TYPE)), + ) + }; + + let result_types = self + .get_ids_from_search(semantic_search_triples, &create_relation_filter) + .await + .unwrap_or_default(); + + let entities: Vec, McpError>> = + join_all(result_types.into_iter().map(|id| async { + triple::find_many(&self.neo4j) + .entity_id(prop_filter::value(id)) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + })? + .try_collect::>() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + }) + })) + .await + .to_vec(); + + Ok(CallToolResult::success( + join_all( + entities + .into_iter() + .map(|result: Result, _>| async { + Content::json(self.format_triples_detailled(result).await) + .expect("Failed to create JSON content") + }), + ) + .await + .to_vec(), + )) + } + + #[tool(description = include_str!("../resources/search_space_description.md"))] + async fn search_space( + &self, + #[tool(param)] + #[schemars(description = "The query string to search for space")] + query: String, + ) -> Result { + let semantic_search_triples = self.search(query, Some(10)).await.unwrap_or_default(); + + let create_relation_filter = |search_result: SemanticSearchResult| { + RelationFilter::default() + .from_(EntityFilter::default().id(prop_filter::value(search_result.triple.entity))) + .relation_type( + EntityFilter::default().id(prop_filter::value(system_ids::TYPES_ATTRIBUTE)), + ) + .to_(EntityFilter::default().id(prop_filter::value(system_ids::SPACE_TYPE))) + }; + + let result_types = self + .get_ids_from_search(semantic_search_triples, &create_relation_filter) + .await + .unwrap_or_default(); + + let entities: Vec, McpError>> = + join_all(result_types.into_iter().map(|id| async { + triple::find_many(&self.neo4j) + .entity_id(prop_filter::value(id)) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + })? + .try_collect::>() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + }) + })) + .await + .to_vec(); + + Ok(CallToolResult::success( + join_all( + entities + .into_iter() + .map(|result: Result, _>| async { + Content::json(self.format_triples_detailled(result).await) + .expect("Failed to create JSON content") + }), + ) + .await + .to_vec(), + )) + } + + #[tool(description = include_str!("../resources/search_properties_description.md"))] + async fn search_properties( + &self, + #[tool(param)] + #[schemars(description = "The query string to search for properties")] + query: String, + ) -> Result { + let semantic_search_triples = self.search(query, Some(10)).await.unwrap_or_default(); + + let create_relation_filter = |search_result: SemanticSearchResult| { + RelationFilter::default() + .from_(EntityFilter::default().id(prop_filter::value(search_result.triple.entity))) + .relation_type( + EntityFilter::default().id(prop_filter::value(system_ids::TYPES_ATTRIBUTE)), + ) + .to_(EntityFilter::default().id(prop_filter::value(system_ids::ATTRIBUTE))) + }; + + let result_types = self + .get_ids_from_search(semantic_search_triples, &create_relation_filter) + .await + .unwrap_or_default(); + + let entities: Vec, McpError>> = + join_all(result_types.into_iter().map(|id| async { + triple::find_many(&self.neo4j) + .entity_id(prop_filter::value(id)) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + })? + .try_collect::>() + .await + .map_err(|e| { + McpError::internal_error( + "search_types_failed", + Some(json!({ "error": e.to_string() })), + ) + }) + })) + .await + .to_vec(); + + Ok(CallToolResult::success( + join_all( + entities + .into_iter() + .map(|result: Result, _>| async { + Content::json(self.format_triples_detailled(result).await) + .expect("Failed to create JSON content") + }), + ) + .await + .to_vec(), + )) + } + + #[tool(description = include_str!("../resources/search_entity_description.md"))] + async fn search_entity( + &self, + #[tool(param)] + #[schemars(description = "A filter of the relation(s) to traverse from the query")] + search_traversal_filter: SearchTraversalInputFilter, + ) -> Result { + tracing::info!( + "SearchTraversalFilter query: {}", + search_traversal_filter.query + ); + let embedding = self .embedding_model - .embed(vec![&query], None) + .embed(vec![&search_traversal_filter.query], None) .expect("Failed to get embedding") .pop() .expect("Embedding is empty") @@ -185,16 +464,37 @@ impl KnowledgeGraph { .map(|v| v as f64) .collect::>(); - let results = entity::search::>(&self.neo4j, embedding) - .filter(entity::EntityFilter::default().relations( - EntityRelationFilter::default().relation_type(system_ids::RELATION_SCHEMA_TYPE), - )) + let traversal_filters: Vec<_> = search_traversal_filter + .traversal_filter + .map(|relation_filter| relation_filter.into_iter().collect()) + .unwrap_or_default(); + + let results_search = traversal_filters + .into_iter() + .fold( + entity::search_from_restictions::>( + &self.neo4j, + embedding.clone(), + ), + |query, filter| { + query.filter( + EntityFilter::default().traverse_relation( + TraverseRelation::default() + .relation_type_id(filter.relation_type_id) + .direction(match filter.direction { + input_types::RelationDirection::From => RelationDirection::From, + input_types::RelationDirection::To => RelationDirection::To, + }), + ), + ) + }, + ) .limit(10) .send() .await .map_err(|e| { McpError::internal_error( - "search_relation_types", + "search_properties", Some(json!({ "error": e.to_string() })), ) })? @@ -202,39 +502,42 @@ impl KnowledgeGraph { .await .map_err(|e| { McpError::internal_error( - "search_relation_types", + "search_properties", Some(json!({ "error": e.to_string() })), ) })?; - tracing::info!("Found {} results for query '{}'", results.len(), query); - - Ok(CallToolResult::success( - results - .into_iter() - .map(|result| { - Content::json(json!({ - "id": result.entity.id(), - "name": result.entity.attributes.name, - "description": result.entity.attributes.description, - "types": result.entity.types, - })) - .expect("Failed to create JSON content") + let entities_vec: Vec<_> = results_search + .into_iter() + .map(|result| { + json!({ + "id": result.entity.id(), + "name": result.entity.attributes.name, + "description": result.entity.attributes.description, }) - .collect(), - )) + }) + .collect::>(); + + Ok(CallToolResult::success(vec![ + Content::json(json!({ + "entities": entities_vec, + })) + .expect("Failed to create JSON content"), + ])) } - #[tool(description = "Search Properties")] - async fn search_properties( + #[tool(description = include_str!("../resources/name_search_entity_description.md"))] + async fn name_search_entity( &self, #[tool(param)] - #[schemars(description = "The query string to search for properties")] - query: String, + #[schemars(description = "A filter of the relation(s) to traverse from the query")] + search_traversal_filter: SearchTraversalInputFilter, ) -> Result { + tracing::info!("SearchTraversalFilter query: {:?}", search_traversal_filter); + let embedding = self .embedding_model - .embed(vec![&query], None) + .embed(vec![&search_traversal_filter.query], None) .expect("Failed to get embedding") .pop() .expect("Embedding is empty") @@ -242,10 +545,73 @@ impl KnowledgeGraph { .map(|v| v as f64) .collect::>(); - let results = entity::search::>(&self.neo4j, embedding) - .filter( - entity::EntityFilter::default() - .relations(TypesFilter::default().r#type(system_ids::ATTRIBUTE)), + let traversal_filters: Vec> = + match search_traversal_filter.traversal_filter { + Some(traversal_filter) => { + join_all(traversal_filter.into_iter().map(|filter| async move { + let rel_embedding = self + .embedding_model + .embed(vec![&filter.relation_type_id], None) + .expect("Failed to get embedding") + .pop() + .expect("Embedding is empty") + .into_iter() + .map(|v| v as f64) + .collect::>(); + + let rel_results = entity::search::(&self.neo4j, rel_embedding) + .filter( + entity::EntityFilter::default().relations( + EntityRelationFilter::default() + .relation_type(system_ids::VALUE_TYPE_ATTRIBUTE) + .to_id(system_ids::RELATION_SCHEMA_TYPE), + ), + ) + .limit(10) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "search_relation_types", + Some(json!({ "error": e.to_string() })), + ) + })? + .try_collect::>() + .await + .map_err(|e| { + McpError::internal_error( + "search_relation_types", + Some(json!({ "error": e.to_string() })), + ) + })?; + let relation_ids: Vec = rel_results + .into_iter() + .map(|sem_search| sem_search.entity.id) + .collect(); + Ok(TraverseRelation::default() + .direction(match filter.direction { + input_types::RelationDirection::From => RelationDirection::From, + input_types::RelationDirection::To => RelationDirection::To, + }) + .relation_type_id(prop_filter::value_in(relation_ids))) + })) + .await + .to_vec() + } + None => Vec::new(), + }; + + let results_search = traversal_filters + .into_iter() + .fold( + entity::search_from_restictions::>( + &self.neo4j, + embedding.clone(), + ), + |query, result_ids: Result<_, McpError>| match result_ids { + Ok(ids) => query.filter(EntityFilter::default().traverse_relation(ids)), + Err(_) => query, + }, ) .limit(10) .send() @@ -265,25 +631,26 @@ impl KnowledgeGraph { ) })?; - tracing::info!("Found {} results for query '{}'", results.len(), query); - - Ok(CallToolResult::success( - results - .into_iter() - .map(|result| { - Content::json(json!({ - "id": result.entity.id(), - "name": result.entity.attributes.name, - "description": result.entity.attributes.description, - "types": result.entity.types, - })) - .expect("Failed to create JSON content") + let entities_vec: Vec<_> = results_search + .into_iter() + .map(|result| { + json!({ + "id": result.entity.id(), + "name": result.entity.attributes.name, + "description": result.entity.attributes.description, }) - .collect(), - )) + }) + .collect::>(); + + Ok(CallToolResult::success(vec![ + Content::json(json!({ + "entities": entities_vec, + })) + .expect("Failed to create JSON content"), + ])) } - #[tool(description = "Get entity by ID with it's attributes and relations")] + #[tool(description = include_str!("../resources/get_entity_info_description.md"))] async fn get_entity_info( &self, #[tool(param)] @@ -292,14 +659,17 @@ impl KnowledgeGraph { )] id: String, ) -> Result { - let entity = entity::find_one::>(&self.neo4j, &id) + let entity_attributes = triple::find_many(&self.neo4j) + .entity_id(prop_filter::value(&id)) .send() .await .map_err(|e| { - McpError::internal_error("get_entity", Some(json!({ "error": e.to_string() }))) + McpError::internal_error("get_entity_info", Some(json!({ "error": e.to_string() }))) })? - .ok_or_else(|| { - McpError::internal_error("entity_not_found", Some(json!({ "id": id }))) + .try_collect::>() + .await + .map_err(|e| { + McpError::internal_error("get_entity_info", Some(json!({ "error": e.to_string() }))) })?; let out_relations = relation::find_many::>(&self.neo4j) @@ -350,31 +720,27 @@ impl KnowledgeGraph { tracing::info!("Found entity with ID '{}'", id); - let clean_up_relations = |relations: Vec>| async { + let clean_up_relations = |relations: Vec>, is_inbound: bool| async move { join_all(relations .into_iter() .map(|result| async move { - Content::json(json!({ + json!({ "relation_id": result.id, "relation_type": self.get_name_of_id(result.relation_type).await.unwrap_or("No relation type".to_string()), - "from_id": result.from.id, - "from_name": self.get_name_of_id(result.from.id).await.unwrap_or("No name".to_string()), - "to_id": result.to.id, - "to_name": self.get_name_of_id(result.to.id).await.unwrap_or("No name".to_string()), - })) - .expect("Failed to create JSON content") + "id": if is_inbound {result.from.id.clone()} else {result.to.id.clone()}, + "name": self.get_name_of_id(if is_inbound {result.from.id.clone()} else {result.to.id.clone()}).await.unwrap_or("No name".to_string()), + }) })).await.to_vec() }; - let inbound_relations = clean_up_relations(in_relations).await; - let outbound_relations = clean_up_relations(out_relations).await; - - let attributes_vec: Vec<_> = join_all(entity.attributes.0.clone().into_iter().map( - |(key, attr)| async { - Content::json(json!({ - "attribute_name": self.get_name_of_id(key).await.unwrap_or("No attribute name".to_string()), - "attribute_value": String::try_from(attr).unwrap_or("No attributes".to_string()), - })) - .expect("Failed to create JSON content") + let inbound_relations = clean_up_relations(in_relations, true).await; + let outbound_relations = clean_up_relations(out_relations, false).await; + + let attributes_vec: Vec<_> = join_all(entity_attributes.into_iter().map( + |attr| async { + json!({ + "attribute_name": self.get_name_of_id(attr.attribute).await.unwrap_or("No attribute name".to_string()), + "attribute_value": String::try_from(attr.value).unwrap_or("No attributes".to_string()), + }) }, )) .await @@ -382,18 +748,16 @@ impl KnowledgeGraph { Ok(CallToolResult::success(vec![ Content::json(json!({ - "id": entity.id(), - "name": entity.attributes.get::(system_ids::NAME_ATTRIBUTE).unwrap_or("No name".to_string()), - "description": entity.attributes.get::(system_ids::DESCRIPTION_ATTRIBUTE).unwrap_or("No description".to_string()), - "types": entity.types, + "id": id, "all_attributes": attributes_vec, "inbound_relations": inbound_relations, "outbound_relations": outbound_relations, - })).expect("Failed to create JSON content"), + })) + .expect("Failed to create JSON content"), ])) } - #[tool(description = "Search for distant or close Relations between 2 entities")] + #[tool(description = include_str!("../resources/get_relations_between_entities_description.md"))] async fn get_relations_between_entities( &self, #[tool(param)] @@ -433,57 +797,6 @@ impl KnowledgeGraph { )) } - #[tool(description = "Get Entity by Attribute")] - async fn get_entity_by_attribute( - &self, - #[tool(param)] - #[schemars(description = "The value of the attribute of an Entity")] - attribute_value: String, - ) -> Result { - let embedding = self - .embedding_model - .embed(vec![&attribute_value], None) - .expect("Failed to get embedding") - .pop() - .expect("Embedding is empty") - .into_iter() - .map(|v| v as f64) - .collect::>(); - - let entities = entity::search::>(&self.neo4j, embedding) - .filter(entity::EntityFilter::default()) - .limit(10) - .send() - .await - .map_err(|e| { - McpError::internal_error("get_entity", Some(json!({ "error": e.to_string() }))) - })? - .try_collect::>() - .await - .map_err(|e| { - McpError::internal_error( - "get_relation_by_id_not_found", - Some(json!({ "error": e.to_string() })), - ) - })?; - - tracing::info!("Found {} entities with given attributes", entities.len()); - - Ok(CallToolResult::success( - entities - .into_iter() - .map(|result| { - Content::json(json!({ - "id": result.entity.id(), - "name": result.entity.attributes.name, - "description": result.entity.attributes.description, - })) - .expect("Failed to create JSON content") - }) - .collect(), - )) - } - async fn get_name_of_id(&self, id: String) -> Result { let entity = entity::find_one::>(&self.neo4j, &id) .send() @@ -494,7 +807,7 @@ impl KnowledgeGraph { .ok_or_else(|| { McpError::internal_error("entity_name_not_found", Some(json!({ "id": id }))) })?; - Ok(entity.attributes.name.unwrap()) + Ok(entity.attributes.name.unwrap_or("No name".to_string())) } } @@ -525,6 +838,53 @@ impl ServerHandler for KnowledgeGraph { } Ok(self.get_info()) } + + //TODO: make prompt examples to use on data + async fn list_prompts( + &self, + _request: Option, + _: RequestContext, + ) -> Result { + Ok(ListPromptsResult { + next_cursor: None, + prompts: vec![Prompt::new( + "example_prompt", + Some("This is an example prompt that takes one required argument, message"), + Some(vec![PromptArgument { + name: "message".to_string(), + description: Some("A message to put in the prompt".to_string()), + required: Some(true), + }]), + )], + }) + } + + async fn get_prompt( + &self, + GetPromptRequestParam { name, arguments }: GetPromptRequestParam, + _: RequestContext, + ) -> Result { + match name.as_str() { + "example_prompt" => { + let message = arguments + .and_then(|json| json.get("message")?.as_str().map(|s| s.to_string())) + .ok_or_else(|| { + McpError::invalid_params("No message provided to example_prompt", None) + })?; + + let prompt = + format!("This is an example prompt with your message here: '{message}'"); + Ok(GetPromptResult { + description: None, + messages: vec![PromptMessage { + role: PromptMessageRole::User, + content: PromptMessageContent::text(prompt), + }], + }) + } + _ => Err(McpError::invalid_params("prompt not found", None)), + } + } } #[derive(Debug, Parser)] diff --git a/sink/build.rs b/sink/build.rs index da489a8..191dcc3 100644 --- a/sink/build.rs +++ b/sink/build.rs @@ -17,8 +17,8 @@ fn main() { .map(|desc| desc.trim().to_string()) .unwrap_or_else(|| "unknown".to_string()); - println!("cargo::rustc-env=GIT_COMMIT={}", git_hash); - println!("cargo::rustc-env=GIT_TAG={}", git_tag); + println!("cargo::rustc-env=GIT_COMMIT={git_hash}"); + println!("cargo::rustc-env=GIT_TAG={git_tag}"); // Always rerun if any git changes occur println!("cargo::rerun-if-changed=.git/HEAD"); diff --git a/sink/examples/seed_data.rs b/sink/examples/seed_data.rs index 48b25f9..d5640d3 100644 --- a/sink/examples/seed_data.rs +++ b/sink/examples/seed_data.rs @@ -1,11 +1,13 @@ +use chrono::Local; use fastembed::{EmbeddingModel, InitOptions, TextEmbedding}; use grc20_core::{ block::BlockMetadata, entity::EntityNodeRef, - ids, - mapping::{triple, Query, RelationEdge, Triple}, + ids, indexer_ids, + mapping::{triple, Query, RelationEdge, Triple, Value}, neo4rs, relation, system_ids, }; +use grc20_sdk::models::space; const EMBEDDING_MODEL: EmbeddingModel = EmbeddingModel::AllMiniLML6V2; @@ -32,17 +34,19 @@ const RUST_ASYNC_WORKSHOP_SIDEEVENT: &str = "QPZnckrRUebWjdwQZTR7Ka"; const RUST_HACKATHON_SIDEEVENT: &str = "ReJ5RRMqTer9qfr87Yjexp"; const JOE_ID: &str = "MpR7wuVWyXV988F5NWZ21r"; const CHRIS_ID: &str = "ScHYh4PpRpyuvY2Ab4Znf5"; -const _: &str = "Mu7ddiBnwZH1LvpDTpKcvq"; -const _: &str = "DVurPdLUZi7Ajfv9BC3ADm"; -const _: &str = "MPxRvh35rnDeRJNEJLU1YF"; -const _: &str = "JjoWPp8LiCKVZiWtE5iZaJ"; -const _: &str = "8bCuTuWqL3dxALLff1Awdb"; -const _: &str = "9Bj46RXQzHQq25WNPY4Lw"; -const _: &str = "RkTkM28NSx3WZuW33vZUjx"; -const _: &str = "Lc9L7StPfXMFGWw45utaTY"; -const _: &str = "G49gECRJmW6BwqHaENF5nS"; -const _: &str = "GfugZRvoWmQhkjMcFJHg49"; -const _: &str = "5bwj7yNukCHoJnW8ksgZY"; +const POLYMTL_ID: &str = "Mu7ddiBnwZH1LvpDTpKcvq"; +const MAUD_COHEN_ID: &str = "DVurPdLUZi7Ajfv9BC3ADm"; +const CIVIL_ENGINEERING_ID: &str = "YEZVCYJTudKVreLEWuxFXV"; +const SOFTWARE_ENGINEERING_ID: &str = "MPxRvh35rnDeRJNEJLU1YF"; +const COMPUTER_ENGINEERING_ID: &str = "JjoWPp8LiCKVZiWtE5iZaJ"; +const MECANICAL_ENGINEERING_ID: &str = "8bCuTuWqL3dxALLff1Awdb"; +const OLIVIER_GENDREAU_ID: &str = "9Bj46RXQzHQq25WNPY4Lw"; +const FR_SPACE_ID: &str = "RkTkM28NSx3WZuW33vZUjx"; +const FR_QC_SPACE_ID: &str = "Lc9L7StPfXMFGWw45utaTY"; +const DIRECTOR_PROP: &str = "G49gECRJmW6BwqHaENF5nS"; +const PROGRAM_TYPE: &str = "GfugZRvoWmQhkjMcFJHg49"; +const SCHOOL_TYPE: &str = "M89C7wwdJVaCW9rAVQpJbY"; +const PROGRAM_PROP: &str = "5bwj7yNukCHoJnW8ksgZY"; const _: &str = "GKXfCXBAJ2oAufgETPcFK7"; const _: &str = "X6q73SFySo5u2BuQrYUxR5"; const _: &str = "S2etHTe7W92QbXz32QWimW"; @@ -80,6 +84,102 @@ async fn main() -> anyhow::Result<()> { reset_db(&neo4j).await?; bootstrap(&neo4j, &embedding_model).await?; + let dt = Local::now(); + + let block = BlockMetadata { + cursor: "random_cursor".to_string(), + block_number: 0, + timestamp: dt.to_utc(), + request_id: "request_id".to_string(), + }; + + create_type( + &neo4j, + &embedding_model, + "Space", + [], + [ + system_ids::NAME_ATTRIBUTE, + system_ids::DESCRIPTION_ATTRIBUTE, + ], + Some(system_ids::SPACE_TYPE), + None, + ) + .await?; + + space::builder(FR_SPACE_ID, "0x5a0b54d5dc17e0aadc383d2db43b0a0d3e029c4c") + .build() + .insert(&neo4j, &block, indexer_ids::INDEXER_SPACE_ID, "0") + .send() + .await?; + + space::builder(FR_QC_SPACE_ID, "0x5a0b54d5dc17e0aadc383d2db43b0a0d3e029c4c") + .build() + .insert(&neo4j, &block, indexer_ids::INDEXER_SPACE_ID, "0") + .send() + .await?; + + space::builder( + system_ids::ROOT_SPACE_ID, + "0x5a0b54d5dc17e0aadc383d2db43b0a0d3e029c4c", + ) + .build() + .insert(&neo4j, &block, indexer_ids::INDEXER_SPACE_ID, "0") + .send() + .await?; + + insert_relation( + &neo4j, + FR_QC_SPACE_ID, + system_ids::TYPES_ATTRIBUTE, + system_ids::SPACE_TYPE, + indexer_ids::INDEXER_SPACE_ID, + ) + .await?; + insert_relation( + &neo4j, + FR_SPACE_ID, + system_ids::TYPES_ATTRIBUTE, + system_ids::SPACE_TYPE, + indexer_ids::INDEXER_SPACE_ID, + ) + .await?; + insert_relation( + &neo4j, + system_ids::ROOT_SPACE_ID, + system_ids::TYPES_ATTRIBUTE, + system_ids::SPACE_TYPE, + indexer_ids::INDEXER_SPACE_ID, + ) + .await?; + insert_attribute_with_embedding( + &neo4j, + &embedding_model, + FR_QC_SPACE_ID, + system_ids::NAME_ATTRIBUTE, + "Quebec", + FR_QC_SPACE_ID, + ) + .await?; + insert_attribute_with_embedding( + &neo4j, + &embedding_model, + FR_QC_SPACE_ID, + system_ids::DESCRIPTION_ATTRIBUTE, + "The space for Quebec related content", + FR_QC_SPACE_ID, + ) + .await?; + insert_attribute_with_embedding( + &neo4j, + &embedding_model, + FR_SPACE_ID, + system_ids::NAME_ATTRIBUTE, + "Francophonie", + FR_SPACE_ID, + ) + .await?; + // Create some common types create_type( &neo4j, @@ -91,6 +191,7 @@ async fn main() -> anyhow::Result<()> { system_ids::DESCRIPTION_ATTRIBUTE, ], Some(system_ids::PERSON_TYPE), + None, ) .await?; @@ -104,6 +205,7 @@ async fn main() -> anyhow::Result<()> { system_ids::DESCRIPTION_ATTRIBUTE, ], Some(EVENT_TYPE), + None, ) .await?; @@ -117,6 +219,35 @@ async fn main() -> anyhow::Result<()> { system_ids::DESCRIPTION_ATTRIBUTE, ], Some(CITY_TYPE), + None, + ) + .await?; + + create_type( + &neo4j, + &embedding_model, + "Program", + [], + [ + system_ids::NAME_ATTRIBUTE, + system_ids::DESCRIPTION_ATTRIBUTE, + ], + Some(PROGRAM_TYPE), + None, + ) + .await?; + + create_type( + &neo4j, + &embedding_model, + "School", + [], + [ + system_ids::NAME_ATTRIBUTE, + system_ids::DESCRIPTION_ATTRIBUTE, + ], + Some(SCHOOL_TYPE), + None, ) .await?; @@ -127,6 +258,7 @@ async fn main() -> anyhow::Result<()> { system_ids::RELATION_SCHEMA_TYPE, Some(CITY_TYPE), Some(EVENT_LOCATION_PROP), + None, ) .await?; @@ -137,6 +269,7 @@ async fn main() -> anyhow::Result<()> { system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::PERSON_TYPE), Some(SPEAKERS_PROP), + None, ) .await?; @@ -147,6 +280,29 @@ async fn main() -> anyhow::Result<()> { system_ids::RELATION_SCHEMA_TYPE, Some(EVENT_TYPE), Some(SIDE_EVENTS), + None, + ) + .await?; + + create_property( + &neo4j, + &embedding_model, + "Director", + system_ids::RELATION_SCHEMA_TYPE, + Some(system_ids::PERSON_TYPE), + Some(DIRECTOR_PROP), + None, + ) + .await?; + + create_property( + &neo4j, + &embedding_model, + "Program", + system_ids::RELATION_SCHEMA_TYPE, + Some(system_ids::PERSON_TYPE), + Some(PROGRAM_PROP), + None, ) .await?; @@ -166,6 +322,7 @@ async fn main() -> anyhow::Result<()> { ], [], Some(ALICE_ID), + None, ) .await?; @@ -178,6 +335,7 @@ async fn main() -> anyhow::Result<()> { [], [], Some(BOB_ID), + None, ) .await?; @@ -190,6 +348,7 @@ async fn main() -> anyhow::Result<()> { [], [], Some(CAROL_ID), + None, ) .await?; @@ -202,6 +361,7 @@ async fn main() -> anyhow::Result<()> { [], [], Some(DAVE_ID), + None, ) .await?; @@ -214,6 +374,7 @@ async fn main() -> anyhow::Result<()> { [], [], Some(JOE_ID), + None, ) .await?; @@ -226,6 +387,125 @@ async fn main() -> anyhow::Result<()> { [], [], Some(CHRIS_ID), + None, + ) + .await?; + + create_entity( + &neo4j, + &embedding_model, + "Maud Cohen", + None, + [system_ids::PERSON_TYPE], + [], + [], + Some(MAUD_COHEN_ID), + None, + ) + .await?; + + create_entity( + &neo4j, + &embedding_model, + "Olivier Gendreau", + None, + [system_ids::PERSON_TYPE], + [], + [], + Some(OLIVIER_GENDREAU_ID), + None, + ) + .await?; + + //Create programs entities + create_entity( + &neo4j, + &embedding_model, + "Software Engineering", + None, + [PROGRAM_TYPE], + [], + [(DIRECTOR_PROP, OLIVIER_GENDREAU_ID)], + Some(SOFTWARE_ENGINEERING_ID), + None, + ) + .await?; + + insert_attribute_with_embedding( + &neo4j, + &embedding_model, + SOFTWARE_ENGINEERING_ID, + system_ids::NAME_ATTRIBUTE, + "Génie logiciel", + FR_SPACE_ID, + ) + .await?; + + create_entity( + &neo4j, + &embedding_model, + "Computer Engineering", + None, + [PROGRAM_TYPE], + [], + [], + Some(COMPUTER_ENGINEERING_ID), + None, + ) + .await?; + + insert_attribute_with_embedding( + &neo4j, + &embedding_model, + SOFTWARE_ENGINEERING_ID, + system_ids::NAME_ATTRIBUTE, + "Génie informatique", + FR_SPACE_ID, + ) + .await?; + + create_entity( + &neo4j, + &embedding_model, + "Civil Engineering", + None, + [PROGRAM_TYPE], + [], + [], + Some(CIVIL_ENGINEERING_ID), + None, + ) + .await?; + + create_entity( + &neo4j, + &embedding_model, + "Mecanical Engineering", + None, + [PROGRAM_TYPE], + [], + [], + Some(MECANICAL_ENGINEERING_ID), + None, + ) + .await?; + + create_entity( + &neo4j, + &embedding_model, + "Polytechnique Montreal", + None, + [SCHOOL_TYPE], + [], + [ + (DIRECTOR_PROP, MAUD_COHEN_ID), + (PROGRAM_PROP, CIVIL_ENGINEERING_ID), + (PROGRAM_PROP, SOFTWARE_ENGINEERING_ID), + (PROGRAM_PROP, COMPUTER_ENGINEERING_ID), + (PROGRAM_PROP, MECANICAL_ENGINEERING_ID), + ], + Some(POLYMTL_ID), + None, ) .await?; @@ -239,6 +519,7 @@ async fn main() -> anyhow::Result<()> { [], [], Some(SAN_FRANCISCO_ID), + None, ) .await?; @@ -251,6 +532,7 @@ async fn main() -> anyhow::Result<()> { [], [], Some(NEW_YORK_ID), + None, ) .await?; @@ -268,6 +550,7 @@ async fn main() -> anyhow::Result<()> { (SPEAKERS_PROP, JOE_ID), ], Some(RUST_ASYNC_WORKSHOP_SIDEEVENT), + None, ) .await?; @@ -283,6 +566,7 @@ async fn main() -> anyhow::Result<()> { (SPEAKERS_PROP, CHRIS_ID), ], Some(RUST_HACKATHON_SIDEEVENT), + None, ) .await?; @@ -301,6 +585,7 @@ async fn main() -> anyhow::Result<()> { (SIDE_EVENTS, RUST_HACKATHON_SIDEEVENT), // RustConf Hackathon ], Some(RUSTCONF_2023), + None, ) .await?; @@ -317,6 +602,7 @@ async fn main() -> anyhow::Result<()> { (EVENT_LOCATION_PROP, NEW_YORK_ID), // New York ], Some(JSCONF_2024), + None, ) .await?; @@ -427,6 +713,7 @@ pub async fn bootstrap( system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::ATTRIBUTE), Some(system_ids::PROPERTIES), + None, ) .await?; @@ -437,6 +724,7 @@ pub async fn bootstrap( system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::SCHEMA_TYPE), Some(system_ids::TYPES_ATTRIBUTE), + None, ) .await?; @@ -447,6 +735,7 @@ pub async fn bootstrap( system_ids::RELATION_SCHEMA_TYPE, None::<&str>, Some(system_ids::VALUE_TYPE_ATTRIBUTE), + None, ) .await?; @@ -457,6 +746,7 @@ pub async fn bootstrap( system_ids::RELATION_SCHEMA_TYPE, None::<&str>, Some(system_ids::RELATION_TYPE_ATTRIBUTE), + None, ) .await?; @@ -467,6 +757,7 @@ pub async fn bootstrap( system_ids::TEXT, None::<&str>, Some(system_ids::RELATION_INDEX), + None, ) .await?; @@ -477,6 +768,7 @@ pub async fn bootstrap( system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::SCHEMA_TYPE), Some(system_ids::RELATION_TYPE_ATTRIBUTE), + None, ) .await?; @@ -487,6 +779,7 @@ pub async fn bootstrap( system_ids::TEXT, None::<&str>, Some(system_ids::NAME_ATTRIBUTE), + None, ) .await?; @@ -497,6 +790,7 @@ pub async fn bootstrap( system_ids::TEXT, None::<&str>, Some(system_ids::DESCRIPTION_ATTRIBUTE), + None, ) .await?; @@ -513,6 +807,7 @@ pub async fn bootstrap( system_ids::DESCRIPTION_ATTRIBUTE, ], Some(system_ids::SCHEMA_TYPE), + None, ) .await?; @@ -523,6 +818,7 @@ pub async fn bootstrap( [system_ids::RELATION_SCHEMA_TYPE], [system_ids::RELATION_VALUE_RELATIONSHIP_TYPE], Some(system_ids::RELATION_SCHEMA_TYPE), + None, ) .await?; @@ -537,6 +833,7 @@ pub async fn bootstrap( system_ids::DESCRIPTION_ATTRIBUTE, ], Some(system_ids::ATTRIBUTE), + None, ) .await?; @@ -550,6 +847,7 @@ pub async fn bootstrap( system_ids::RELATION_INDEX, ], Some(system_ids::RELATION_TYPE), + None, ) .await?; @@ -565,13 +863,16 @@ pub async fn create_entity( properties: impl IntoIterator, relations: impl IntoIterator, id: Option<&str>, + space_id: Option, ) -> anyhow::Result { let block = BlockMetadata::default(); let entity_id = id.map(Into::into).unwrap_or_else(|| ids::create_geo_id()); let name = name.into(); + let space_id = space_id.as_deref().unwrap_or(system_ids::ROOT_SPACE_ID); + // Set: Entity.name - triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_ID, DEFAULT_VERSION) + triple::insert_many(neo4j, &block, space_id, DEFAULT_VERSION) .triples(vec![Triple::with_embedding( &entity_id, system_ids::NAME_ATTRIBUTE, @@ -590,7 +891,7 @@ pub async fn create_entity( // Set: Entity.description if let Some(description) = description { - triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_ID, DEFAULT_VERSION) + triple::insert_many(neo4j, &block, space_id, DEFAULT_VERSION) .triples(vec![Triple::new( &entity_id, system_ids::DESCRIPTION_ATTRIBUTE, @@ -604,7 +905,7 @@ pub async fn create_entity( set_types(neo4j, &entity_id, types).await?; // Set: Entity.* - triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_ID, DEFAULT_VERSION) + triple::insert_many(neo4j, &block, space_id, DEFAULT_VERSION) .triples( properties .into_iter() @@ -614,24 +915,113 @@ pub async fn create_entity( .await?; // Set: Entity > RELATIONS > Relation[] - relation::insert_many::>( + relation::insert_many::>(neo4j, &block, space_id, DEFAULT_VERSION) + .relations(relations.into_iter().map(|(relation_type, target_id)| { + RelationEdge::new( + ids::create_geo_id(), + &entity_id, + target_id, + relation_type, + "0", + ) + })) + .send() + .await?; + + Ok(entity_id) +} + +pub async fn insert_attribute( + neo4j: &neo4rs::Graph, + entity_id: impl Into, + attribute_id: impl Into, + attribute_value: impl Into, + space_id: impl Into, +) -> anyhow::Result { + let block = BlockMetadata::default(); + let attribute_id = attribute_id.into(); + let attribute_value = attribute_value.into(); + let space_id = space_id.into(); + let entity_id = entity_id.into(); + + triple::insert_one( neo4j, &block, - system_ids::ROOT_SPACE_ID, + space_id, DEFAULT_VERSION, + Triple::new(entity_id.clone(), attribute_id, attribute_value), ) - .relations(relations.into_iter().map(|(relation_type, target_id)| { + .send() + .await?; + Ok(entity_id) +} + +pub async fn insert_relation( + neo4j: &neo4rs::Graph, + entity_from_id: impl Into, + relation_id: impl Into, + entity_to_id: impl Into, + space_id: impl Into, +) -> anyhow::Result { + let block = BlockMetadata::default(); + let entity_from_id = entity_from_id.into(); + let relation_id = relation_id.into(); + let space_id = space_id.into(); + let entity_to_id = entity_to_id.into(); + + relation::insert_one( + neo4j, + &block, + space_id, + DEFAULT_VERSION, RelationEdge::new( - ids::create_geo_id(), - &entity_id, - target_id, - relation_type, - "0", - ) - })) + "id".to_string(), + entity_from_id, + entity_to_id, + relation_id.clone(), + Value::text(relation_id.clone()), + ), + ) .send() .await?; + Ok(relation_id) +} + +pub async fn insert_attribute_with_embedding( + neo4j: &neo4rs::Graph, + embedding_model: &TextEmbedding, + entity_id: impl Into, + attribute_id: impl Into, + attribute_value: impl Into, + space_id: impl Into, +) -> anyhow::Result { + let block = BlockMetadata::default(); + let attribute_id = attribute_id.into(); + let attribute_value = attribute_value.into(); + let space_id = space_id.into(); + let entity_id = entity_id.into(); + triple::insert_one( + neo4j, + &block, + space_id, + DEFAULT_VERSION, + Triple::with_embedding( + &entity_id, + attribute_id, + attribute_value.clone(), + embedding_model + .embed(vec![attribute_value], Some(1)) + .unwrap_or(vec![Vec::::new()]) + .get(0) + .unwrap_or(&Vec::::new()) + .iter() + .map(|&x| x as f64) + .collect(), + ), + ) + .send() + .await?; Ok(entity_id) } @@ -643,6 +1033,7 @@ pub async fn create_type( types: impl IntoIterator, properties: impl IntoIterator, id: Option<&str>, + space_id: Option, ) -> anyhow::Result { let block = BlockMetadata::default(); let type_id = id.map(Into::into).unwrap_or_else(|| ids::create_geo_id()); @@ -653,8 +1044,10 @@ pub async fn create_type( types_vec.push(system_ids::SCHEMA_TYPE); } + let space_id = space_id.as_deref().unwrap_or(system_ids::ROOT_SPACE_ID); + // Set: Type.name - triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_TYPE, DEFAULT_VERSION) + triple::insert_many(neo4j, &block, space_id, DEFAULT_VERSION) .triples(vec![Triple::with_embedding( &type_id, system_ids::NAME_ATTRIBUTE, @@ -675,23 +1068,18 @@ pub async fn create_type( set_types(neo4j, &type_id, types_vec).await?; // Set: Type > PROPERTIES > Property[] - relation::insert_many::>( - neo4j, - &block, - system_ids::ROOT_SPACE_ID, - DEFAULT_VERSION, - ) - .relations(properties.into_iter().map(|property_id| { - RelationEdge::new( - ids::create_geo_id(), - &type_id, - system_ids::PROPERTIES, - property_id, - "0", - ) - })) - .send() - .await?; + relation::insert_many::>(neo4j, &block, space_id, DEFAULT_VERSION) + .relations(properties.into_iter().map(|property_id| { + RelationEdge::new( + ids::create_geo_id(), + &type_id, + system_ids::PROPERTIES, + property_id, + "0", + ) + })) + .send() + .await?; Ok(type_id) } @@ -706,14 +1094,17 @@ pub async fn create_property( value_type: impl Into, relation_value_type: Option>, id: Option>, + space_id: Option, ) -> anyhow::Result { let block = BlockMetadata::default(); let property_id = id.map(Into::into).unwrap_or_else(|| ids::create_geo_id()); let string_name = name.into(); + let space_id = space_id.as_deref().unwrap_or(system_ids::ROOT_SPACE_ID); + // Set: Property.name - triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_ID, DEFAULT_VERSION) + triple::insert_many(neo4j, &block, space_id, DEFAULT_VERSION) .triples(vec![Triple::with_embedding( &property_id, system_ids::NAME_ATTRIBUTE, @@ -734,7 +1125,7 @@ pub async fn create_property( relation::insert_one::>( neo4j, &block, - system_ids::ROOT_SPACE_ID, + space_id, DEFAULT_VERSION, RelationEdge::new( ids::create_geo_id(), diff --git a/sink/src/events/handler.rs b/sink/src/events/handler.rs index 7cfd3d4..49edabd 100644 --- a/sink/src/events/handler.rs +++ b/sink/src/events/handler.rs @@ -96,15 +96,13 @@ impl EventHandler { InitOptions::new(EMBEDDING_MODEL).with_show_download_progress(true), ) .map_err(|e| { - tracing::error!("Error initializing embedding model: {:?}", e); - HandlerError::Other(format!("Error initializing embedding model: {:?}", e).into()) + tracing::error!("Error initializing embedding model: {e:?}"); + HandlerError::Other(format!("Error initializing embedding model: {e:?}").into()) })?, embedding_model_dim: TextEmbedding::get_model_info(&EMBEDDING_MODEL) .map_err(|e| { - tracing::error!("Error getting embedding model info: {:?}", e); - HandlerError::Other( - format!("Error getting embedding model info: {:?}", e).into(), - ) + tracing::error!("Error getting embedding model info: {e:?}"); + HandlerError::Other(format!("Error getting embedding model info: {e:?}").into()) })? .dim, versioning: false, diff --git a/sink/src/main.rs b/sink/src/main.rs index 7913559..ccd8828 100644 --- a/sink/src/main.rs +++ b/sink/src/main.rs @@ -81,10 +81,10 @@ async fn main() -> Result<(), Error> { MODULE_NAME, start_block .parse() - .unwrap_or_else(|_| panic!("Invalid start block: {}! Must be integer", start_block)), + .unwrap_or_else(|_| panic!("Invalid start block: {start_block}! Must be integer")), end_block .parse() - .unwrap_or_else(|_| panic!("Invalid end block: {}! Must be integer", end_block)), + .unwrap_or_else(|_| panic!("Invalid end block: {end_block}! Must be integer")), Some(64), ) .await?; diff --git a/substreams-utils/src/sink.rs b/substreams-utils/src/sink.rs index 8dd6d23..793a9a2 100644 --- a/substreams-utils/src/sink.rs +++ b/substreams-utils/src/sink.rs @@ -128,7 +128,7 @@ pub trait Sink: Send + Sync { Some(Err(err)) => { println!(); println!("Stream terminated with error"); - println!("{:?}", err); + println!("{err:?}"); exit(1); } } diff --git a/substreams-utils/src/substreams_stream.rs b/substreams-utils/src/substreams_stream.rs index 35178d5..f041743 100644 --- a/substreams-utils/src/substreams_stream.rs +++ b/substreams-utils/src/substreams_stream.rs @@ -119,7 +119,7 @@ fn stream_blocks( return Err(anyhow::Error::new(status.clone()))?; } - println!("Received tonic error {:#}", status); + println!("Received tonic error {status:#}"); encountered_error = true; break; }, @@ -136,7 +136,7 @@ fn stream_blocks( // case where we actually _want_ to back off in case we keep // having connection errors. - println!("Unable to connect to endpoint: {:#}", e); + println!("Unable to connect to endpoint: {e:#}"); } }