diff --git a/.gitignore b/.gitignore index 39a1aa6..1867157 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ **/target **/ipfs-cache **/data +**/sample_data *.json .env **/.fastembed_cache \ No newline at end of file diff --git a/grc20-core/src/mapping/entity/find_path.rs b/grc20-core/src/mapping/entity/find_path.rs new file mode 100644 index 0000000..90e9cdd --- /dev/null +++ b/grc20-core/src/mapping/entity/find_path.rs @@ -0,0 +1,161 @@ +use neo4rs::Path; + +use crate::{ + entity::EntityFilter, + error::DatabaseError, + mapping::{ + order_by::FieldOrderBy, + query_utils::{ + query_builder::{MatchQuery, QueryBuilder, Subquery}, + VersionFilter, + }, + AttributeFilter, PropFilter, Query, + }, + system_ids::SCHEMA_TYPE, +}; + +pub struct Relation { + pub nodes_ids: Vec, + pub relations_ids: Vec, +} + +pub struct FindPathQuery { + neo4j: neo4rs::Graph, + id1: String, + id2: String, + filter: EntityFilter, + order_by: Option, + limit: usize, + skip: Option, + space_id: Option>, + version: VersionFilter, +} + +impl FindPathQuery { + pub(super) fn new(neo4j: &neo4rs::Graph, id1: String, id2: String) -> Self { + Self { + neo4j: neo4j.clone(), + id1, + id2, + filter: EntityFilter::default(), + order_by: None, + limit: 100, + skip: None, + space_id: None, + version: VersionFilter::default(), + } + } + + pub fn id(mut self, id: PropFilter) -> Self { + self.filter.id = Some(id); + self + } + + pub fn attribute(mut self, attribute: AttributeFilter) -> Self { + self.filter.attributes.push(attribute); + self + } + + pub fn attribute_mut(&mut self, attribute: AttributeFilter) { + self.filter.attributes.push(attribute); + } + + pub fn attributes(mut self, attributes: impl IntoIterator) -> Self { + self.filter.attributes.extend(attributes); + self + } + + pub fn attributes_mut(&mut self, attributes: impl IntoIterator) { + self.filter.attributes.extend(attributes); + } + + pub fn limit(mut self, limit: usize) -> Self { + self.limit = limit; + self + } + + pub fn skip(mut self, skip: usize) -> Self { + self.skip = Some(skip); + self + } + + /// Overwrite the current filter with a new one + pub fn with_filter(mut self, filter: EntityFilter) -> Self { + self.filter = filter; + self + } + + pub fn order_by(mut self, order_by: FieldOrderBy) -> Self { + self.order_by = Some(order_by); + self + } + + pub fn order_by_mut(&mut self, order_by: FieldOrderBy) { + self.order_by = Some(order_by); + } + + pub fn space_id(mut self, space_id: impl Into>) -> Self { + self.space_id = Some(space_id.into()); + self + } + + pub fn version(mut self, space_version: String) -> Self { + self.version.version_mut(space_version); + self + } + + pub fn version_opt(mut self, space_version: Option) -> Self { + self.version.version_opt(space_version); + self + } + + fn subquery(&self) -> QueryBuilder { + QueryBuilder::default() + .subquery(MatchQuery::new( + "p = allShortestPaths((e1:Entity {id: $id1}) -[:RELATION*1..10]-(e2:Entity {id: $id2}))", + ).r#where(format!("NONE(n IN nodes(p) WHERE EXISTS((n)-[:RELATION]-(:Entity {{id: \"{SCHEMA_TYPE}\"}})))")))//makes sure to not use primitive types + .limit(self.limit) + .params("id1", self.id1.clone()) + .params("id2", self.id2.clone()) + } +} + +impl Query> for FindPathQuery { + async fn send(self) -> Result, DatabaseError> { + let query = self.subquery().r#return("p"); + + if cfg!(debug_assertions) || cfg!(test) { + println!( + "entity::FindPathQuery:::\n{}\nparams:{:?}", + query.compile(), + [self.id1, self.id2] + ); + } + + let mut result = self.neo4j.execute(query.build()).await?; + let mut all_relationship_data = Vec::new(); + + // Process each row + while let Some(row) = result.next().await? { + let path: Path = row.get("p")?; + tracing::info!("This is the info for Path: {:?}", path); + + let relationship_data: Relation = Relation { + nodes_ids: (path + .nodes() + .iter() + .filter_map(|rel| rel.get("id").ok()) + .collect()), + relations_ids: (path + .rels() + .iter() + .filter_map(|rel| rel.get("relation_type").ok()) + .collect()), + }; + + all_relationship_data.push(relationship_data); + } + + Ok(all_relationship_data) + } +} diff --git a/grc20-core/src/mapping/entity/mod.rs b/grc20-core/src/mapping/entity/mod.rs index 5f8192b..3e67560 100644 --- a/grc20-core/src/mapping/entity/mod.rs +++ b/grc20-core/src/mapping/entity/mod.rs @@ -2,6 +2,7 @@ pub mod delete_many; pub mod delete_one; pub mod find_many; pub mod find_one; +pub mod find_path; pub mod insert_many; pub mod insert_one; pub mod models; @@ -11,6 +12,7 @@ pub mod utils; pub use delete_one::DeleteOneQuery; pub use find_many::FindManyQuery; pub use find_one::FindOneQuery; +pub use find_path::FindPathQuery; pub use insert_one::InsertOneQuery; pub use models::{Entity, EntityNode, EntityNodeRef, SystemProperties}; pub use semantic_search::SemanticSearchQuery; @@ -128,6 +130,11 @@ pub fn search(neo4j: &neo4rs::Graph, vector: Vec) -> SemanticSearchQuery SemanticSearchQuery::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) +} + pub fn insert_one( neo4j: &neo4rs::Graph, block: &BlockMetadata, diff --git a/grc20-core/src/mapping/query_utils/attributes_filter.rs b/grc20-core/src/mapping/query_utils/attributes_filter.rs index 7860342..3b1ff61 100644 --- a/grc20-core/src/mapping/query_utils/attributes_filter.rs +++ b/grc20-core/src/mapping/query_utils/attributes_filter.rs @@ -104,9 +104,10 @@ impl AttributeFilter { pub fn subquery(&self, node_var: &str) -> MatchQuery { let attr_rel_var = format!("r_{node_var}_{}", self.attribute); let attr_node_var = format!("{node_var}_{}", self.attribute); + let attr_id_var = format!("a_{node_var}_{}", self.attribute); MatchQuery::new( - format!("({node_var}) -[{attr_rel_var}:ATTRIBUTE]-> ({attr_node_var}:Attribute {{id: $attribute}})") + format!("({node_var}) -[{attr_rel_var}:ATTRIBUTE]-> ({attr_node_var}:Attribute {{id: ${attr_id_var}}})") ) .r#where(self.version.subquery(&attr_rel_var)) .where_opt( @@ -118,6 +119,6 @@ impl AttributeFilter { .where_opt( self.value_type.as_ref().map(|value_type| value_type.subquery(&attr_node_var, "value_type", None)) ) - .params("attribute", self.attribute.clone()) + .params(attr_id_var, self.attribute.clone()) } } diff --git a/grc20-core/src/mapping/relation/utils.rs b/grc20-core/src/mapping/relation/utils.rs index 15082d5..fba8398 100644 --- a/grc20-core/src/mapping/relation/utils.rs +++ b/grc20-core/src/mapping/relation/utils.rs @@ -85,6 +85,6 @@ impl<'a> MatchOneRelation<'a> { .subquery(format!("ORDER BY {edge_var}.index")) .params("id", self.id) .params("space_id", self.space_id) - .with(vec![from_node_var, edge_var, to_node_var], next) + .with(vec![from_node_var, edge_var.to_string(), to_node_var], next) } } diff --git a/mcp-server/src/main.rs b/mcp-server/src/main.rs index 7e55c45..5e4bcd3 100644 --- a/mcp-server/src/main.rs +++ b/mcp-server/src/main.rs @@ -1,10 +1,10 @@ use clap::{Args, Parser}; use fastembed::{EmbeddingModel, InitOptions, TextEmbedding}; -use futures::TryStreamExt; +use futures::{TryStreamExt, future::join_all}; use grc20_core::{ - entity::{self, Entity, EntityRelationFilter, TypesFilter}, - mapping::{Query, QueryStream}, - neo4rs, system_ids, + entity::{self, Entity, EntityFilter, EntityNode, EntityRelationFilter, TypesFilter}, + mapping::{Attributes, Query, QueryStream, RelationEdge, prop_filter}, + neo4rs, relation, system_ids, }; use grc20_sdk::models::BaseEntity; use rmcp::{ @@ -132,7 +132,7 @@ impl KnowledgeGraph { entity::EntityFilter::default() .relations(TypesFilter::default().r#type(system_ids::SCHEMA_TYPE)), ) - .limit(8) + .limit(10) .send() .await .map_err(|e| { @@ -186,14 +186,10 @@ impl KnowledgeGraph { .collect::>(); let results = entity::search::>(&self.neo4j, embedding) - .filter( - entity::EntityFilter::default().relations( - EntityRelationFilter::default() - .relation_type(system_ids::VALUE_TYPE_ATTRIBUTE) - .to_id(system_ids::RELATION_SCHEMA_TYPE), - ), - ) - .limit(8) + .filter(entity::EntityFilter::default().relations( + EntityRelationFilter::default().relation_type(system_ids::RELATION_SCHEMA_TYPE), + )) + .limit(10) .send() .await .map_err(|e| { @@ -251,7 +247,7 @@ impl KnowledgeGraph { entity::EntityFilter::default() .relations(TypesFilter::default().r#type(system_ids::ATTRIBUTE)), ) - .limit(8) + .limit(10) .send() .await .map_err(|e| { @@ -287,24 +283,16 @@ impl KnowledgeGraph { )) } - // #[tool(description = "Search Properties")] - // async fn get_entities( - // &self, - // #[tool(param)] - // #[schemars(description = "The query string to search for properties")] - // query: String, - // ) - - #[tool(description = "Get entity by ID")] - async fn get_entity( + #[tool(description = "Get entity by ID with it's attributes and relations")] + async fn get_entity_info( &self, #[tool(param)] #[schemars( - description = "Return an entity by its ID along with its attributes (name, description, etc.) and types" + description = "Return an entity by its ID along with its attributes (name, description, etc.), relations and types" )] id: String, ) -> Result { - let entity = entity::find_one::>(&self.neo4j, &id) + let entity = entity::find_one::>(&self.neo4j, &id) .send() .await .map_err(|e| { @@ -314,18 +302,200 @@ impl KnowledgeGraph { McpError::internal_error("entity_not_found", Some(json!({ "id": id }))) })?; + let out_relations = relation::find_many::>(&self.neo4j) + .filter( + relation::RelationFilter::default() + .from_(EntityFilter::default().id(prop_filter::value(id.clone()))), + ) + .limit(10) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "get_relation_by_id", + 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() })), + ) + })?; + + let in_relations = relation::find_many::>(&self.neo4j) + .filter( + relation::RelationFilter::default() + .to_(EntityFilter::default().id(prop_filter::value(id.clone()))), + ) + .limit(10) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "get_relation_by_id", + 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 entity with ID '{}'", id); + let clean_up_relations = |relations: Vec>| async { + join_all(relations + .into_iter() + .map(|result| async move { + Content::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") + })).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") + }, + )) + .await + .to_vec(); + Ok(CallToolResult::success(vec![ Content::json(json!({ "id": entity.id(), - "name": entity.attributes.name, - "description": entity.attributes.description, + "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, - })) - .expect("Failed to create JSON content"), + "all_attributes": attributes_vec, + "inbound_relations": inbound_relations, + "outbound_relations": outbound_relations, + })).expect("Failed to create JSON content"), ])) } + + #[tool(description = "Search for distant or close Relations between 2 entities")] + async fn get_relations_between_entities( + &self, + #[tool(param)] + #[schemars(description = "The id of the first Entity to find relations")] + entity1_id: String, + #[tool(param)] + #[schemars(description = "The id of the second Entity to find relations")] + entity2_id: String, + ) -> Result { + let relations = entity::find_path(&self.neo4j, entity1_id.clone(), entity2_id.clone()) + .limit(10) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "get_relation_by_ids", + Some(json!({ "error": e.to_string() })), + ) + })? + .into_iter() + .collect::>(); + + tracing::info!("Found {} relations", relations.len()); + + Ok(CallToolResult::success( + join_all(relations + .into_iter() + .map(|result| async { + Content::json(json!({ + "nodes": join_all(result.nodes_ids.into_iter().map(|node_id| async {self.get_name_of_id(node_id).await.unwrap_or("No attribute name".to_string())})).await.to_vec(), + "relations": join_all(result.relations_ids.into_iter().map(|node_id| async {self.get_name_of_id(node_id).await.unwrap_or("No attribute name".to_string())})).await.to_vec(), + })) + .expect("Failed to create JSON content") + })) + .await + .to_vec(), + )) + } + + #[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() + .await + .map_err(|e| { + McpError::internal_error("get_entity_name", Some(json!({ "error": e.to_string() }))) + })? + .ok_or_else(|| { + McpError::internal_error("entity_name_not_found", Some(json!({ "id": id }))) + })?; + Ok(entity.attributes.name.unwrap()) + } } #[tool(tool_box)] diff --git a/sink/examples/seed_data.rs b/sink/examples/seed_data.rs index 57d9439..48b25f9 100644 --- a/sink/examples/seed_data.rs +++ b/sink/examples/seed_data.rs @@ -72,13 +72,18 @@ async fn main() -> anyhow::Result<()> { .await .expect("Failed to connect to Neo4j"); + let embedding_model = TextEmbedding::try_new( + InitOptions::new(EMBEDDING_MODEL).with_show_download_progress(true), + )?; + // Reset and bootstrap the database reset_db(&neo4j).await?; - bootstrap(&neo4j).await?; + bootstrap(&neo4j, &embedding_model).await?; // Create some common types create_type( &neo4j, + &embedding_model, "Person", [], [ @@ -91,6 +96,7 @@ async fn main() -> anyhow::Result<()> { create_type( &neo4j, + &embedding_model, "Event", [], [ @@ -103,6 +109,7 @@ async fn main() -> anyhow::Result<()> { create_type( &neo4j, + &embedding_model, "City", [], [ @@ -115,6 +122,7 @@ async fn main() -> anyhow::Result<()> { create_property( &neo4j, + &embedding_model, "Event location", system_ids::RELATION_SCHEMA_TYPE, Some(CITY_TYPE), @@ -124,6 +132,7 @@ async fn main() -> anyhow::Result<()> { create_property( &neo4j, + &embedding_model, "Speakers", system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::PERSON_TYPE), @@ -133,6 +142,7 @@ async fn main() -> anyhow::Result<()> { create_property( &neo4j, + &embedding_model, "Side events", system_ids::RELATION_SCHEMA_TYPE, Some(EVENT_TYPE), @@ -143,6 +153,7 @@ async fn main() -> anyhow::Result<()> { // Create person entities create_entity( &neo4j, + &embedding_model, "Alice", None, [system_ids::PERSON_TYPE], @@ -160,6 +171,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "Bob", None, [system_ids::PERSON_TYPE], @@ -171,6 +183,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "Carol", None, [system_ids::PERSON_TYPE], @@ -182,6 +195,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "Dave", None, [system_ids::PERSON_TYPE], @@ -193,6 +207,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "Joe", None, [system_ids::PERSON_TYPE], @@ -204,6 +219,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "Chris", None, [system_ids::PERSON_TYPE], @@ -216,6 +232,7 @@ async fn main() -> anyhow::Result<()> { // Create city entities create_entity( &neo4j, + &embedding_model, "San Francisco", Some("City in California"), [CITY_TYPE], @@ -227,6 +244,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "New York", Some("City in New York State"), [CITY_TYPE], @@ -240,6 +258,7 @@ async fn main() -> anyhow::Result<()> { // Create side event entities for RustConf 2023 create_entity( &neo4j, + &embedding_model, "Rust Async Workshop", Some("A hands-on workshop about async programming in Rust"), [EVENT_TYPE], @@ -254,6 +273,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "RustConf Hackathon", Some("A hackathon for RustConf 2023 attendees"), [EVENT_TYPE], @@ -268,6 +288,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "Rust Conference 2023", Some("A conference about Rust programming language"), [EVENT_TYPE], @@ -285,6 +306,7 @@ async fn main() -> anyhow::Result<()> { create_entity( &neo4j, + &embedding_model, "JavaScript Summit 2024", Some("A summit for JavaScript enthusiasts and professionals"), [EVENT_TYPE], @@ -301,11 +323,10 @@ async fn main() -> anyhow::Result<()> { Ok(()) } -pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { - let embedding_model = TextEmbedding::try_new( - InitOptions::new(EMBEDDING_MODEL).with_show_download_progress(true), - )?; - +pub async fn bootstrap( + neo4j: &neo4rs::Graph, + embedding_model: &TextEmbedding, +) -> anyhow::Result<()> { let triples = vec![ // Value types Triple::new(system_ids::CHECKBOX, system_ids::NAME_ATTRIBUTE, "Checkbox"), @@ -401,6 +422,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { // Create properties create_property( neo4j, + &embedding_model, "Properties", system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::ATTRIBUTE), @@ -410,6 +432,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_property( neo4j, + &embedding_model, "Types", system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::SCHEMA_TYPE), @@ -419,6 +442,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_property( neo4j, + &embedding_model, "Value Type", system_ids::RELATION_SCHEMA_TYPE, None::<&str>, @@ -428,6 +452,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_property( neo4j, + &embedding_model, "Relation type attribute", system_ids::RELATION_SCHEMA_TYPE, None::<&str>, @@ -437,6 +462,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_property( neo4j, + &embedding_model, "Relation index", system_ids::TEXT, None::<&str>, @@ -446,6 +472,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_property( neo4j, + &embedding_model, "Relation value type", system_ids::RELATION_SCHEMA_TYPE, Some(system_ids::SCHEMA_TYPE), @@ -455,6 +482,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_property( neo4j, + &embedding_model, "Name", system_ids::TEXT, None::<&str>, @@ -464,6 +492,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_property( neo4j, + &embedding_model, "Description", system_ids::TEXT, None::<&str>, @@ -474,6 +503,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { // Create types create_type( neo4j, + &embedding_model, "Type", [system_ids::SCHEMA_TYPE], [ @@ -488,6 +518,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_type( neo4j, + &embedding_model, "Relation schema type", [system_ids::RELATION_SCHEMA_TYPE], [system_ids::RELATION_VALUE_RELATIONSHIP_TYPE], @@ -497,6 +528,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_type( neo4j, + &embedding_model, "Attribute", [system_ids::SCHEMA_TYPE], [ @@ -510,6 +542,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { create_type( neo4j, + &embedding_model, "Relation instance type", [system_ids::RELATION_TYPE], [ @@ -525,6 +558,7 @@ pub async fn bootstrap(neo4j: &neo4rs::Graph) -> anyhow::Result<()> { pub async fn create_entity( neo4j: &neo4rs::Graph, + embedding_model: &TextEmbedding, name: impl Into, description: Option<&str>, types: impl IntoIterator, @@ -538,10 +572,18 @@ pub async fn create_entity( // Set: Entity.name triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_ID, DEFAULT_VERSION) - .triples(vec![Triple::new( + .triples(vec![Triple::with_embedding( &entity_id, system_ids::NAME_ATTRIBUTE, - name, + name.clone(), + embedding_model + .embed(vec![name], Some(1)) + .unwrap_or(vec![Vec::::new()]) + .get(0) + .unwrap_or(&Vec::::new()) + .iter() + .map(|&x| x as f64) + .collect(), )]) .send() .await?; @@ -596,6 +638,7 @@ pub async fn create_entity( /// Creates a type with the given name, types, and properties. pub async fn create_type( neo4j: &neo4rs::Graph, + embedding_model: &TextEmbedding, name: impl Into, types: impl IntoIterator, properties: impl IntoIterator, @@ -612,10 +655,18 @@ pub async fn create_type( // Set: Type.name triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_TYPE, DEFAULT_VERSION) - .triples(vec![Triple::new( + .triples(vec![Triple::with_embedding( &type_id, system_ids::NAME_ATTRIBUTE, - name, + name.clone(), + embedding_model + .embed(vec![name], Some(1)) + .unwrap_or(vec![Vec::::new()]) + .get(0) + .unwrap_or(&Vec::::new()) + .iter() + .map(|&x| x as f64) + .collect(), )]) .send() .await?; @@ -650,6 +701,7 @@ pub async fn create_type( /// Note: if that is the case, then `value_type` should be the system_ids::RELATION_SCHEMA_TYPE type). pub async fn create_property( neo4j: &neo4rs::Graph, + embedding_model: &TextEmbedding, name: impl Into, value_type: impl Into, relation_value_type: Option>, @@ -658,13 +710,22 @@ pub async fn create_property( let block = BlockMetadata::default(); let property_id = id.map(Into::into).unwrap_or_else(|| ids::create_geo_id()); + let string_name = name.into(); // Set: Property.name triple::insert_many(neo4j, &block, system_ids::ROOT_SPACE_ID, DEFAULT_VERSION) - .triples(vec![Triple::new( + .triples(vec![Triple::with_embedding( &property_id, system_ids::NAME_ATTRIBUTE, - name.into(), + string_name.clone(), + embedding_model + .embed(vec![string_name], Some(1)) + .unwrap_or(vec![Vec::::new()]) + .get(0) + .unwrap_or(&Vec::::new()) + .iter() + .map(|&x| x as f64) + .collect(), )]) .send() .await?;