diff --git a/Cargo.lock b/Cargo.lock index 43f3095..1bafabf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -99,16 +99,16 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.95" +version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04" +checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" [[package]] name = "api" version = "0.1.0" dependencies = [ "anyhow", - "axum", + "axum 0.7.9", "cache", "chrono", "clap", @@ -417,7 +417,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" dependencies = [ "async-trait", - "axum-core", + "axum-core 0.4.5", "bytes", "futures-util", "http 1.2.0", @@ -426,7 +426,41 @@ dependencies = [ "hyper 1.6.0", "hyper-util", "itoa", - "matchit", + "matchit 0.7.3", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "rustversion", + "serde", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper 1.0.2", + "tokio", + "tower 0.5.2", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "021e862c184ae977658b36c4500f7feac3221ca5da43e3f25bd04ab6c79a29b5" +dependencies = [ + "axum-core 0.5.2", + "bytes", + "form_urlencoded", + "futures-util", + "http 1.2.0", + "http-body 1.0.1", + "http-body-util", + "hyper 1.6.0", + "hyper-util", + "itoa", + "matchit 0.8.4", "memchr", "mime", "percent-encoding", @@ -465,6 +499,26 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum-core" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68464cd0412f486726fb3373129ef5d2993f90c34bc2bc1c1e9943b2f4fc7ca6" +dependencies = [ + "bytes", + "futures-core", + "http 1.2.0", + "http-body 1.0.1", + "http-body-util", + "mime", + "pin-project-lite", + "rustversion", + "sync_wrapper 1.0.2", + "tower-layer", + "tower-service", + "tracing", +] + [[package]] name = "backoff" version = "0.4.0" @@ -475,7 +529,7 @@ dependencies = [ "getrandom 0.2.15", "instant", "pin-project-lite", - "rand", + "rand 0.8.5", "tokio", ] @@ -756,9 +810,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.27" +version = "4.5.39" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "769b0145982b4b48713e01ec42d61614425f27b7058bda7180a3a41f30104796" +checksum = "fd60e63e9be68e5fb56422e397cf9baddded06dae1d2e523401542383bc72a9f" dependencies = [ "clap_builder", "clap_derive", @@ -766,9 +820,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.27" +version = "4.5.39" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b26884eb4b57140e4d2d93652abfa49498b938b3c9179f9fc487b0acc3edad7" +checksum = "89cc6392a1f72bbeb820d71f32108f61fdaf18bc526e1d23954168a67759ef51" dependencies = [ "anstream", "anstyle", @@ -778,9 +832,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.24" +version = "4.5.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54b755194d6389280185988721fffba69495eed5ee9feeee9a599b53db80318c" +checksum = "09176aae279615badda0765c0c0b3f6ed53f4709118af73cf4655d85d1530cd7" dependencies = [ "heck", "proc-macro2", @@ -1170,6 +1224,12 @@ dependencies = [ "serde_json", ] +[[package]] +name = "dyn-clone" +version = "1.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c7a8fb8a9fbf66c1f703fe16184d10ca0ee9d23be5b4436400408ba54a95005" + [[package]] name = "either" version = "1.13.0" @@ -1597,7 +1657,7 @@ dependencies = [ "neo4rs", "pretty_assertions", "prost", - "rand", + "rand 0.8.5", "serde", "serde_json", "serde_with", @@ -1737,7 +1797,7 @@ dependencies = [ "libc", "log", "native-tls", - "rand", + "rand 0.8.5", "reqwest 0.12.12", "serde", "serde_json", @@ -2397,7 +2457,7 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed3b21b9af313a2967572c8d4b8875c53fc8062e10768470de4748c16ce7b992" dependencies = [ - "axum", + "axum 0.7.9", "bytes", "juniper", "juniper_graphql_ws", @@ -2626,6 +2686,12 @@ version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + [[package]] name = "matrixmultiply" version = "0.3.9" @@ -2646,6 +2712,27 @@ dependencies = [ "rayon", ] +[[package]] +name = "mcp-server" +version = "0.1.0" +dependencies = [ + "anyhow", + "axum 0.8.4", + "clap", + "fastembed", + "futures", + "grc20-core", + "grc20-sdk", + "rmcp", + "schemars", + "serde", + "serde_json", + "tokio", + "tokio-util", + "tracing", + "tracing-subscriber", +] + [[package]] name = "md-5" version = "0.10.6" @@ -2666,7 +2753,7 @@ dependencies = [ "enum_dispatch", "openssl", "r2d2", - "rand", + "rand 0.8.5", "url", ] @@ -3141,7 +3228,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" dependencies = [ "phf_shared", - "rand", + "rand 0.8.5", ] [[package]] @@ -3421,8 +3508,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.3", ] [[package]] @@ -3432,7 +3529,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.3", ] [[package]] @@ -3444,6 +3551,15 @@ dependencies = [ "getrandom 0.2.15", ] +[[package]] +name = "rand_core" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" +dependencies = [ + "getrandom 0.3.1", +] + [[package]] name = "rav1e" version = "0.7.1" @@ -3470,8 +3586,8 @@ dependencies = [ "once_cell", "paste", "profiling", - "rand", - "rand_chacha", + "rand 0.8.5", + "rand_chacha 0.3.1", "simd_helpers", "system-deps", "thiserror 1.0.69", @@ -3717,6 +3833,41 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rmcp" +version = "0.1.5" +source = "git+https://github.com/modelcontextprotocol/rust-sdk?branch=main#d5a72e43c17d688086738030387af1cd39a9ce38" +dependencies = [ + "axum 0.8.4", + "base64 0.22.1", + "chrono", + "futures", + "paste", + "pin-project-lite", + "rand 0.9.1", + "rmcp-macros", + "schemars", + "serde", + "serde_json", + "thiserror 2.0.11", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "uuid", +] + +[[package]] +name = "rmcp-macros" +version = "0.1.5" +source = "git+https://github.com/modelcontextprotocol/rust-sdk?branch=main#d5a72e43c17d688086738030387af1cd39a9ce38" +dependencies = [ + "proc-macro2", + "quote", + "serde_json", + "syn 2.0.96", +] + [[package]] name = "rustc-demangle" version = "0.1.24" @@ -3850,6 +4001,31 @@ dependencies = [ "parking_lot", ] +[[package]] +name = "schemars" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615" +dependencies = [ + "chrono", + "dyn-clone", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 2.0.96", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -3912,6 +4088,17 @@ dependencies = [ "syn 2.0.96", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.96", +] + [[package]] name = "serde_json" version = "1.0.140" @@ -4090,7 +4277,7 @@ name = "sink" version = "0.1.0" dependencies = [ "anyhow", - "axum", + "axum 0.7.9", "cache", "chrono", "clap", @@ -4612,7 +4799,7 @@ dependencies = [ "monostate", "onig", "paste", - "rand", + "rand 0.8.5", "rayon", "rayon-cond", "regex", @@ -4628,9 +4815,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.44.1" +version = "1.45.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f382da615b842244d4b8738c82ed1275e6c5dd90c459a30941cd07080b06c91a" +checksum = "75ef51a33ef1da925cea3e4eb122833cb377c61439ca401b770f54902b806779" dependencies = [ "backtrace", "bytes", @@ -4672,7 +4859,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f57eb36ecbe0fc510036adff84824dd3c24bb781e21bfa67b69d556aa85214f" dependencies = [ "pin-project", - "rand", + "rand 0.8.5", "tokio", ] @@ -4715,9 +4902,9 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.13" +version = "0.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7fcaa8d55a2bdd6b83ace262b016eca0d79ee02818c5c1bcdf0305114081078" +checksum = "66a539a9ad6d5d281510d5bd368c973d636c02dbf8a67300bfb6b950696ad7df" dependencies = [ "bytes", "futures-core", @@ -4768,7 +4955,7 @@ checksum = "877c5b330756d856ffcc4553ab34a5684481ade925ecc54bcd1bf02b1d0d4d52" dependencies = [ "async-stream", "async-trait", - "axum", + "axum 0.7.9", "base64 0.22.1", "bytes", "flate2", @@ -4805,7 +4992,7 @@ dependencies = [ "indexmap 1.9.3", "pin-project", "pin-project-lite", - "rand", + "rand 0.8.5", "slab", "tokio", "tokio-util", diff --git a/Cargo.toml b/Cargo.toml index c34c1b8..29e4a9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,5 +8,5 @@ members = [ "web3-utils", "grc20-core", "grc20-macros", - "grc20-sdk", + "grc20-sdk", "mcp-server", ] diff --git a/README.md b/README.md index 977910a..83b3218 100644 --- a/README.md +++ b/README.md @@ -43,5 +43,13 @@ Schema introspection npx get-graphql-schema http://127.0.0.1:8080/graphql > api/schema.graphql ``` +## MCP Server +```bash +CFLAGS='-std=gnu17' cargo run --bin mcp-server -- \ + --neo4j-uri neo4j://localhost:7687 \ + --neo4j-user neo4j \ + --neo4j-pass neo4j +``` + ## GRC20 CLI Coming soon™️ \ No newline at end of file diff --git a/api/src/schema/query.rs b/api/src/schema/query.rs index 791ebeb..0e49f5f 100644 --- a/api/src/schema/query.rs +++ b/api/src/schema/query.rs @@ -318,12 +318,12 @@ impl RootQuery { .map(|triple| Triple::new(triple, space_id, version_index))) } - async fn search<'a, S: ScalarValue>( + async fn search_triples<'a, S: ScalarValue>( &'a self, executor: &'a Executor<'_, '_, KnowledgeGraph, S>, query: String, #[graphql(default = 100)] first: i32, - // #[graphql(default = 0)] skip: i32, + #[graphql(default = 0)] skip: i32, ) -> FieldResult> { let embedding = executor .context() @@ -336,8 +336,9 @@ impl RootQuery { .map(|v| v as f64) .collect::>(); - let query = mapping::triple::semantic_search(&executor.context().neo4j, embedding) - .limit(first as usize); + let query = mapping::triple::search(&executor.context().neo4j, embedding) + .limit(first as usize) + .skip(skip as usize); Ok(query .send() diff --git a/grc20-core/src/mapping/entity/find_many.rs b/grc20-core/src/mapping/entity/find_many.rs index 1a2f620..ff3704d 100644 --- a/grc20-core/src/mapping/entity/find_many.rs +++ b/grc20-core/src/mapping/entity/find_many.rs @@ -157,6 +157,7 @@ impl QueryStream> for FindManyQuery> { "e", "attrs", "types", + None, "RETURN e{.*, attrs: attrs, types: types}", ), ); diff --git a/grc20-core/src/mapping/entity/find_one.rs b/grc20-core/src/mapping/entity/find_one.rs index 451b286..f0ebe9e 100644 --- a/grc20-core/src/mapping/entity/find_one.rs +++ b/grc20-core/src/mapping/entity/find_one.rs @@ -88,6 +88,7 @@ impl Query>> for FindOneQuery> { "e", "attrs", "types", + None, "RETURN e{.*, attrs: attrs, types: types}", ), ) diff --git a/grc20-core/src/mapping/entity/mod.rs b/grc20-core/src/mapping/entity/mod.rs index be732f9..96f48c4 100644 --- a/grc20-core/src/mapping/entity/mod.rs +++ b/grc20-core/src/mapping/entity/mod.rs @@ -5,6 +5,7 @@ pub mod find_one; pub mod insert_many; pub mod insert_one; pub mod models; +pub mod semantic_search; pub mod utils; pub use delete_one::DeleteOneQuery; @@ -12,6 +13,7 @@ pub use find_many::FindManyQuery; pub use find_one::FindOneQuery; pub use insert_one::InsertOneQuery; pub use models::{Entity, EntityNode, EntityNodeRef, SystemProperties}; +pub use semantic_search::SemanticSearchQuery; pub use utils::{EntityFilter, EntityRelationFilter}; use crate::block::BlockMetadata; @@ -40,6 +42,10 @@ pub fn find_many(neo4j: &neo4rs::Graph) -> FindManyQuery { FindManyQuery::new(neo4j) } +pub fn search(neo4j: &neo4rs::Graph, vector: Vec) -> SemanticSearchQuery { + SemanticSearchQuery::new(neo4j, vector) +} + pub fn insert_one( neo4j: &neo4rs::Graph, block: &BlockMetadata, diff --git a/grc20-core/src/mapping/entity/semantic_search.rs b/grc20-core/src/mapping/entity/semantic_search.rs new file mode 100644 index 0000000..f8607f9 --- /dev/null +++ b/grc20-core/src/mapping/entity/semantic_search.rs @@ -0,0 +1,218 @@ +use futures::{Stream, StreamExt, TryStreamExt}; + +use crate::{ + entity::utils::MatchEntity, + error::DatabaseError, + mapping::{ + query_utils::VersionFilter, AttributeNode, FromAttributes, PropFilter, QueryBuilder, + QueryStream, Subquery, + }, +}; + +use super::{Entity, EntityFilter, EntityNode}; + +pub struct SemanticSearchQuery { + neo4j: neo4rs::Graph, + vector: Vec, + filter: EntityFilter, + space_id: Option>, + version: VersionFilter, + limit: usize, + skip: Option, + + _marker: std::marker::PhantomData, +} + +impl SemanticSearchQuery { + pub fn new(neo4j: &neo4rs::Graph, vector: Vec) -> Self { + Self { + neo4j: neo4j.clone(), + vector, + filter: EntityFilter::default(), + space_id: None, + version: VersionFilter::default(), + limit: 100, + skip: None, + + _marker: std::marker::PhantomData, + } + } + + pub fn filter(mut self, filter: EntityFilter) -> Self { + self.filter = 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 + } + + 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 + MATCH (e:Entity) -[r:ATTRIBUTE]-> (n) + "#; + + // Exact neighbor search using vector index (very expensive but allows prefiltering) + // const QUERY: &str = const_format::formatcp!( + // r#" + // MATCH (e:Entity) -[r:ATTRIBUTE]-> (a:Attribute:Indexed) + // WHERE r.max_version IS null + // AND a.embedding IS NOT NULL + // WITH e, a, r, vector.similarity.cosine(a.embedding, $vector) AS score + // ORDER BY score DESC + // WHERE score IS NOT null + // LIMIT $limit + // RETURN a{{.*, entity: e.id, space_version: r.min_version, space_id: r.space_id, score: score}} + // "#, + // ); + + QueryBuilder::default() + .subquery(QUERY) + .subquery(self.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) + } +} + +#[derive(Clone, Debug, PartialEq)] +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, + ) -> Result< + impl Stream, DatabaseError>>, + DatabaseError, + > { + let query = self.subquery().r#return("DISTINCT e, score"); + + if cfg!(debug_assertions) || cfg!(test) { + tracing::info!( + "entity_node::FindManyQuery:::\n{}\nparams:{:?}", + query.compile(), + query.params() + ); + }; + + #[derive(Debug, serde::Deserialize)] + struct RowResult { + e: EntityNode, + score: f64, + } + + Ok(self + .neo4j + .execute(query.build()) + .await? + .into_stream_as::() + .map_err(DatabaseError::from) + .and_then(|row| async move { + Ok(SemanticSearchResult { + entity: row.e, + score: row.score, + }) + })) + } +} + +impl QueryStream>> + for SemanticSearchQuery> +{ + 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(), "score".to_string()], + match_entity.chain( + "e", + "attrs", + "types", + Some(vec!["score".to_string()]), + "RETURN e{.*, attrs: attrs, types: types, score: score}", + ), + ); + + if cfg!(debug_assertions) || cfg!(test) { + tracing::info!( + "entity_node::FindManyQuery::>:\n{}\nparams:{:?}", + query.compile(), + query.params + ); + }; + + #[derive(Debug, serde::Deserialize)] + struct RowResult { + #[serde(flatten)] + node: EntityNode, + attrs: Vec, + types: Vec, + score: f64, + } + + 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| SemanticSearchResult { + entity: Entity { + node: row.node, + attributes: data, + types: row.types.into_iter().map(|t| t.id).collect(), + }, + score: row.score, + }) + .map_err(DatabaseError::from) + }) + }); + + Ok(stream) + } +} diff --git a/grc20-core/src/mapping/entity/utils.rs b/grc20-core/src/mapping/entity/utils.rs index fc05811..3fd6117 100644 --- a/grc20-core/src/mapping/entity/utils.rs +++ b/grc20-core/src/mapping/entity/utils.rs @@ -16,6 +16,8 @@ pub struct EntityFilter { pub(crate) id: Option>, pub(crate) attributes: Vec, pub(crate) relations: 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>, } @@ -284,6 +286,7 @@ impl<'a> MatchEntity<'a> { node_var: impl Into, attributes_node_var: impl Into, types_node_var: impl Into, + extra_vars: Option>, next: impl Subquery, ) -> QueryBuilder { let node_var = node_var.into(); @@ -294,21 +297,23 @@ impl<'a> MatchEntity<'a> { // let attributes_node_var = format!("{entity_node_var}_attributes"); // let types_node_var = format!("{entity_node_var}_types"); + let with_vars = vec![ + node_var.clone(), + format!("COLLECT(DISTINCT {attributes_node_var}{{.*}}) AS {attributes_node_var}"), + format!("COLLECT(DISTINCT {types_node_var}{{.*}}) AS {types_node_var}"), + ]; + let with_vars = if let Some(extra_vars) = extra_vars { + with_vars.into_iter().chain(extra_vars).collect() + } else { + with_vars + }; + QueryBuilder::default() .subquery( self.match_attributes .subquery(&node_var, &attributes_node_var), ) .subquery(self.match_types.subquery(&node_var, &types_node_var)) - .with( - vec![ - node_var, - format!( - "COLLECT(DISTINCT {attributes_node_var}{{.*}}) AS {attributes_node_var}" - ), - format!("COLLECT(DISTINCT {types_node_var}{{.*}}) AS {types_node_var}"), - ], - next, - ) + .with(with_vars, next) } } diff --git a/grc20-core/src/mapping/query_utils/attributes_filter.rs b/grc20-core/src/mapping/query_utils/attributes_filter.rs index f23be78..7860342 100644 --- a/grc20-core/src/mapping/query_utils/attributes_filter.rs +++ b/grc20-core/src/mapping/query_utils/attributes_filter.rs @@ -1,5 +1,19 @@ use super::{prop_filter::PropFilter, query_builder::MatchQuery, version_filter::VersionFilter}; +/// Struct representing an attribute filter subquery for an entity's attributes. +/// +/// IMPORTANT: This filter subquery is designed to be used to filter an entity by its +/// attributes (and not filter a list of attributes!) +/// +/// The struct follows the builder pattern to set the filter parameters. +/// ```rust +/// let filter = AttributeFilter::new(system_ids::NAME_ATTRIBUTE) +/// .value(["Bob", "Alice"]) +/// .value_type("TEXT") +/// .space_id("25omwWh6HYgeRQKCaSpVpa"); +/// +/// let subquery = filter.subquery("e"); +/// ``` #[derive(Clone, Debug)] pub struct AttributeFilter { attribute: String, @@ -10,6 +24,9 @@ pub struct AttributeFilter { } impl AttributeFilter { + /// Create a new filter subquery for the provided `attribute`. By default, if no other + /// parameters are set, this filter subquery will filter entities for which the `attribute` + /// exists in the current version of the knowledge graph. pub fn new(attribute: &str) -> Self { Self { attribute: attribute.to_owned(), @@ -20,23 +37,23 @@ impl AttributeFilter { } } - pub fn space_id(mut self, space_id: PropFilter) -> Self { - self.space_id = Some(space_id); + pub fn space_id(mut self, space_id: impl Into>) -> Self { + self.space_id = Some(space_id.into()); self } - pub fn space_id_mut(&mut self, space_id: PropFilter) -> &mut Self { - self.space_id = Some(space_id); + pub fn space_id_mut(&mut self, space_id: impl Into>) -> &mut Self { + self.space_id = Some(space_id.into()); self } - pub fn value(mut self, value: PropFilter) -> Self { - self.value = Some(value); + pub fn value(mut self, value: impl Into>) -> Self { + self.value = Some(value.into()); self } - pub fn value_type(mut self, value_type: PropFilter) -> Self { - self.value_type = Some(value_type); + pub fn value_type(mut self, value_type: impl Into>) -> Self { + self.value_type = Some(value_type.into()); self } @@ -55,6 +72,35 @@ impl AttributeFilter { self } + /// Compiles the attribute filter into a Neo4j subquery that will filter the nodes + /// identified by `node_var` according to the provided parameters. + /// + /// The subquery will have the following form: + /// ```cypher + /// MATCH ({node_var}) -[r_{node_var}_attribute:ATTRIBUTE]-> ({node_var}_attribute:Attribute {id: $attribute}) + /// WHERE {VERSION_FILTER} + /// AND {SPACE_ID_FITLER} + /// AND {VALUE_FILTER} + /// AND {VALUE_TYPE_FILTER} + /// ``` + /// + /// For example, if: + /// - the attribute to filter on is `LuBWqZAu6pz54eiJS5mLv8` + /// - the nodes to filter are bound to the variable `e` + /// - the version filter is set to filter the current version + /// - the value filter is set to `["foo", "bar"]` + /// - the value type filter set to `TEXT` + /// - the space id filter set to `25omwWh6HYgeRQKCaSpVpa` + /// + /// the subquery will be: + /// ```cypher + /// MATCH (e) -[r_e_attribute:ATTRIBUTE]-> (e_attribute:Attribute {id: $attribute}) + /// WHERE r_e_attribute.max_version IS NULL + /// AND r_e_attribute.space_id = "25omwWh6HYgeRQKCaSpVpa" + /// AND e_attribute.value IN ["foo", "bar"] + /// AND e_attribute.value_type = "TEXT" + /// ``` + /// Note: the `$attribute` query parameter will contain the value `LuBWqZAu6pz54eiJS5mLv8` 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); diff --git a/grc20-core/src/mapping/query_utils/prop_filter.rs b/grc20-core/src/mapping/query_utils/prop_filter.rs index 15d2e0a..9eb816f 100644 --- a/grc20-core/src/mapping/query_utils/prop_filter.rs +++ b/grc20-core/src/mapping/query_utils/prop_filter.rs @@ -158,13 +158,41 @@ impl PropFilter { } impl> PropFilter { - /// Converts the filter into a query part. - /// The `node_var` is the variable name of the node in the query. - /// The `key` is the property key of the node. - /// The `expr` is an optional expression to use instead of the property key. - /// If `expr` is `None`, the node_var and key will be used as the expression to - /// filter, e.g. `{node_var}.{key} = $value` - pub(crate) fn subquery(&self, node_var: &str, key: &str, expr: Option<&str>) -> WhereClause { + /// Compiles the attribute filter into a [WhereClause] Neo4j subquery that will apply + /// a filter on the `key` field of the `node_var` node(s) (i.e.: `{node_var}.{key}`). + /// + /// If `expr` is set, then it will used as the filter target instead of the above. + /// + /// For example, given the following [PropFilter] (which creates a property filter) + /// ```rust + /// # fn main() { + /// # use std::collections::HashMap; + /// # use grc20_core::mapping::{PropFilter, Subquery}; + /// let prop_filter = PropFilter::default() + /// .value_not("Bob") + /// .value_lt("Gary"); + /// + /// let query = prop_filter.subquery("e", "name", None); + /// assert_eq!( + /// query.compile(), + /// "WHERE e.name <> $e_name_value_not\nAND e.name < $e_name_value_lt" + /// ) + /// assert_eq!( + /// query.params(), + /// HashMap::from([ + /// ("e_name_value_not", "Bob"), + /// ("e_name_value_lt", "Gary") + /// ]) + /// ) + /// + /// let query = prop_filter.subquery("e", "name", Some("my_expr")); + /// assert_eq!( + /// query.compile(), + /// "WHERE my_expr <> $e_name_value_not\nAND my_expr < $e_name_value_lt" + /// ) + /// # } + /// ``` + pub fn subquery(&self, node_var: &str, key: &str, expr: Option<&str>) -> WhereClause { let mut where_clause = WhereClause::default(); let expr = expr @@ -175,56 +203,56 @@ impl> PropFilter { let param_key = format!("{node_var}_{key}_value"); where_clause = where_clause .clause(format!("{expr} = ${param_key}")) - .params(param_key, value.clone()); + .set_param(param_key, value.clone()); } if let Some(value_gt) = &self.value_gt { let param_key = format!("{node_var}_{key}_value_gt"); where_clause = where_clause .clause(format!("{expr} > ${param_key}")) - .params(param_key, value_gt.clone()); + .set_param(param_key, value_gt.clone()); } if let Some(value_gte) = &self.value_gte { let param_key = format!("{node_var}_{key}_value_gte"); where_clause = where_clause .clause(format!("{expr} >= ${param_key}")) - .params(param_key, value_gte.clone()); + .set_param(param_key, value_gte.clone()); } if let Some(value_lt) = &self.value_lt { let param_key = format!("{node_var}_{key}_value_lt"); where_clause = where_clause .clause(format!("{expr} < ${param_key}")) - .params(param_key, value_lt.clone()); + .set_param(param_key, value_lt.clone()); } if let Some(value_lte) = &self.value_lte { let param_key = format!("{node_var}_{key}_value_lte"); where_clause = where_clause .clause(format!("{expr} <= ${param_key}")) - .params(param_key, value_lte.clone()); + .set_param(param_key, value_lte.clone()); } if let Some(value_not) = &self.value_not { let param_key = format!("{node_var}_{key}_value_not"); where_clause = where_clause .clause(format!("{expr} <> ${param_key}")) - .params(param_key, value_not.clone()); + .set_param(param_key, value_not.clone()); } if let Some(value_in) = &self.value_in { let param_key = format!("{node_var}_{key}_value_in"); where_clause = where_clause .clause(format!("{expr} IN ${param_key}")) - .params(param_key, value_in.clone()); + .set_param(param_key, value_in.clone()); } if let Some(value_not_in) = &self.value_not_in { let param_key = format!("{node_var}_{key}_value_not_in"); where_clause = where_clause .clause(format!("{expr} NOT IN ${param_key}")) - .params(param_key, value_not_in.clone()); + .set_param(param_key, value_not_in.clone()); } where_clause diff --git a/grc20-core/src/mapping/query_utils/query_builder.rs b/grc20-core/src/mapping/query_utils/query_builder.rs index 3562a32..16d8564 100644 --- a/grc20-core/src/mapping/query_utils/query_builder.rs +++ b/grc20-core/src/mapping/query_utils/query_builder.rs @@ -236,7 +236,7 @@ impl WhereClause { self } - pub fn params(mut self, key: impl Into, value: impl Into) -> Self { + pub fn set_param(mut self, key: impl Into, value: impl Into) -> Self { self.params.insert(key.into(), value.into()); self } diff --git a/grc20-core/src/mapping/query_utils/version_filter.rs b/grc20-core/src/mapping/query_utils/version_filter.rs index 62e0351..615c066 100644 --- a/grc20-core/src/mapping/query_utils/version_filter.rs +++ b/grc20-core/src/mapping/query_utils/version_filter.rs @@ -28,7 +28,7 @@ impl VersionFilter { let param_key = format!("{}_version", var); WhereClause::new(format!("{var}.min_version <= ${param_key} AND ({var}.max_version IS NULL OR {var}.max_version > ${param_key})")) - .params(param_key, version.clone()) + .set_param(param_key, version.clone()) } else { WhereClause::new(format!("{var}.max_version IS NULL")) } diff --git a/grc20-core/src/mapping/relation/find_many_to.rs b/grc20-core/src/mapping/relation/find_many_to.rs index 5c73963..d5b6086 100644 --- a/grc20-core/src/mapping/relation/find_many_to.rs +++ b/grc20-core/src/mapping/relation/find_many_to.rs @@ -149,6 +149,7 @@ impl QueryStream> for FindManyToQuery> { "to", "attrs", "types", + None, "RETURN to{.*, attrs: attrs, types: types}", ), ); diff --git a/grc20-core/src/mapping/relation/find_one_to.rs b/grc20-core/src/mapping/relation/find_one_to.rs index 6c88c2d..bf9ec1a 100644 --- a/grc20-core/src/mapping/relation/find_one_to.rs +++ b/grc20-core/src/mapping/relation/find_one_to.rs @@ -84,6 +84,7 @@ impl Query>> for FindOneToQuery> { "to", "attrs", "types", + None, "RETURN to{.*, attrs: attrs, types: types}", ), ); diff --git a/grc20-core/src/mapping/triple.rs b/grc20-core/src/mapping/triple.rs index e5c1e0c..61d504b 100644 --- a/grc20-core/src/mapping/triple.rs +++ b/grc20-core/src/mapping/triple.rs @@ -136,7 +136,7 @@ pub fn find_many(neo4j: &neo4rs::Graph) -> FindManyQuery { FindManyQuery::new(neo4j) } -pub fn semantic_search(neo4j: &neo4rs::Graph, vector: Vec) -> SemanticSearchQuery { +pub fn search(neo4j: &neo4rs::Graph, vector: Vec) -> SemanticSearchQuery { SemanticSearchQuery::new(neo4j, vector) } @@ -621,7 +621,7 @@ pub struct SemanticSearchQuery { // space_id: Option>, // space_version: VersionFilter, limit: usize, - // skip: Option, + skip: Option, } impl SemanticSearchQuery { @@ -632,7 +632,7 @@ impl SemanticSearchQuery { // space_id: None, // space_version: VersionFilter::default(), limit: 100, - // skip: None, + skip: None, } } @@ -658,15 +658,15 @@ impl SemanticSearchQuery { self } - // pub fn skip(mut self, skip: usize) -> Self { - // self.skip = Some(skip); - // 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 skip_opt(mut self, skip: Option) -> Self { + self.skip = skip; + self + } } #[derive(Clone, Debug, Default, Deserialize, PartialEq)] @@ -678,34 +678,40 @@ 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, ) -> Result>, DatabaseError> { - // const QUERY: &str = const_format::formatcp!( - // r#" - // CALL db.index.vector.queryNodes('vector_index', $limit, $vector) - // YIELD node AS n, score AS score - // MATCH (e:Entity) -[r:ATTRIBUTE]-> (n) - // RETURN n{{.*, entity: e.id, space_version: r.min_version, space_id: r.space_id, score: score}} - // "# - // ); const QUERY: &str = const_format::formatcp!( r#" - MATCH (e:Entity) -[r:ATTRIBUTE]-> (a:Attribute:Indexed) - WHERE r.max_version IS null - WITH e, a, r, vector.similarity.cosine(a.embedding, $vector) AS score + CALL db.index.vector.queryNodes('vector_index', $limit * $effective_search_ratio, $vector) + YIELD node AS n, score AS score ORDER BY score DESC - WHERE score IS NOT null LIMIT $limit - RETURN a{{.*, entity: e.id, space_version: r.min_version, space_id: r.space_id, score: score}} - "#, + MATCH (e:Entity) -[r:ATTRIBUTE]-> (n) + RETURN n{{.*, entity: e.id, space_version: r.min_version, space_id: r.space_id, score: score}} + "# ); + // const QUERY: &str = const_format::formatcp!( + // r#" + // MATCH (e:Entity) -[r:ATTRIBUTE]-> (a:Attribute:Indexed) + // WHERE r.max_version IS null + // AND a.embedding IS NOT NULL + // WITH e, a, r, vector.similarity.cosine(a.embedding, $vector) AS score + // ORDER BY score DESC + // WHERE score IS NOT null + // LIMIT $limit + // RETURN a{{.*, entity: e.id, space_version: r.min_version, space_id: r.space_id, score: score}} + // "#, + // ); let query = neo4rs::query(QUERY) .param("vector", self.vector) - .param("limit", self.limit as i64); + .param("limit", self.limit as i64) + .param("effective_search_ratio", EFFECTIVE_SEARCH_RATIO); Ok(self .neo4j diff --git a/grc20-sdk/src/models/base_entity.rs b/grc20-sdk/src/models/base_entity.rs index 6924f51..76983c7 100644 --- a/grc20-sdk/src/models/base_entity.rs +++ b/grc20-sdk/src/models/base_entity.rs @@ -11,10 +11,10 @@ use grc20_core::{ #[grc20_core::entity] pub struct BaseEntity { #[grc20(attribute = system_ids::NAME_ATTRIBUTE)] - name: Option, + pub name: Option, #[grc20(attribute = system_ids::DESCRIPTION_ATTRIBUTE)] - description: Option, + pub description: Option, } pub async fn blocks( diff --git a/mcp-server/Cargo.toml b/mcp-server/Cargo.toml new file mode 100644 index 0000000..388ddbc --- /dev/null +++ b/mcp-server/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "mcp-server" +version = "0.1.0" +edition = "2024" + +[dependencies] +anyhow = "1.0.98" +axum = "0.8.4" +clap = { version = "4.5.39", features = ["derive", "env"] } +fastembed = "4.8.0" +futures = "0.3.31" +grc20-core = { version = "0.1.0", path = "../grc20-core" } +grc20-sdk = { version = "0.1.0", path = "../grc20-sdk" } +rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "main", features = ["server", "transport-sse-server"] } +schemars = "0.8.22" +serde = { version = "1.0.219", features = ["derive"] } +serde_json = "1.0.140" +tokio = { version = "1.45.1", features = ["rt-multi-thread", "macros", "signal"] } +tokio-util = "0.7.15" +tracing = "0.1.41" +tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } diff --git a/mcp-server/resources/get_properties_description.md b/mcp-server/resources/get_properties_description.md new file mode 100644 index 0000000..e69de29 diff --git a/mcp-server/resources/instructions.md b/mcp-server/resources/instructions.md new file mode 100644 index 0000000..e251cd8 --- /dev/null +++ b/mcp-server/resources/instructions.md @@ -0,0 +1,13 @@ +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. + +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! + +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. + +Here is an example: +User> What are the properties of the Person type? + +ToolCall> search_type("person") +ToolResult> +``` +``` \ No newline at end of file diff --git a/mcp-server/resources/search_entity_description.md b/mcp-server/resources/search_entity_description.md new file mode 100644 index 0000000..e69de29 diff --git a/mcp-server/resources/search_relation_type_description.md b/mcp-server/resources/search_relation_type_description.md new file mode 100644 index 0000000..e69de29 diff --git a/mcp-server/resources/search_type_description.md b/mcp-server/resources/search_type_description.md new file mode 100644 index 0000000..e69de29 diff --git a/mcp-server/src/main.rs b/mcp-server/src/main.rs new file mode 100644 index 0000000..caf58cf --- /dev/null +++ b/mcp-server/src/main.rs @@ -0,0 +1,380 @@ +use clap::{Args, Parser}; +use fastembed::{EmbeddingModel, InitOptions, TextEmbedding}; +use futures::TryStreamExt; +use grc20_core::{ + entity::{self, Entity, EntityRelationFilter}, + mapping::{Query, QueryStream, query_utils::TypesFilter}, + neo4rs, system_ids, +}; +use grc20_sdk::models::BaseEntity; +use rmcp::{ + Error as McpError, RoleServer, ServerHandler, + model::*, + service::RequestContext, + tool, + transport::sse_server::{SseServer, SseServerConfig}, +}; +use serde_json::json; +use std::sync::Arc; +use tracing_subscriber::{ + layer::SubscriberExt, + util::SubscriberInitExt, + {self}, +}; + +const BIND_ADDRESS: &str = "127.0.0.1:8000"; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_subscriber::registry() + .with( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "debug".to_string().into()), + ) + .with(tracing_subscriber::fmt::layer()) + .init(); + + let args = AppArgs::parse(); + + let neo4j = neo4rs::Graph::new( + &args.neo4j_args.neo4j_uri, + &args.neo4j_args.neo4j_user, + &args.neo4j_args.neo4j_pass, + ) + .await?; + + let config = SseServerConfig { + bind: BIND_ADDRESS.parse()?, + sse_path: "/sse".to_string(), + post_path: "/message".to_string(), + ct: tokio_util::sync::CancellationToken::new(), + sse_keep_alive: None, + }; + + let (sse_server, router) = SseServer::new(config); + + // Do something with the router, e.g., add routes or middleware + + let listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?; + + let ct = sse_server.config.ct.child_token(); + + let server = axum::serve(listener, router).with_graceful_shutdown(async move { + ct.cancelled().await; + tracing::info!("sse server cancelled"); + }); + + tokio::spawn(async move { + if let Err(e) = server.await { + tracing::error!(error = %e, "sse server shutdown with error"); + } + }); + + let ct = sse_server.with_service(move || KnowledgeGraph::new(neo4j.clone())); + + tokio::signal::ctrl_c().await?; + ct.cancel(); + Ok(()) +} + +#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] +pub struct StructRequest { + pub a: i32, + pub b: i32, +} + +const EMBEDDING_MODEL: EmbeddingModel = EmbeddingModel::AllMiniLML6V2; + +#[derive(Clone)] +pub struct KnowledgeGraph { + neo4j: neo4rs::Graph, + pub embedding_model: Arc, +} + +#[tool(tool_box)] +impl KnowledgeGraph { + #[allow(dead_code)] + pub fn new(neo4j: neo4rs::Graph) -> Self { + Self { + neo4j, + embedding_model: Arc::new( + TextEmbedding::try_new( + InitOptions::new(EMBEDDING_MODEL).with_show_download_progress(true), + ) + .expect("Failed to initialize embedding model"), + ), + } + } + + fn _create_resource_text(&self, uri: &str, name: &str) -> Resource { + RawResource::new(uri, name.to_string()).no_annotation() + } + + #[tool(description = "Search Types")] + async fn search_types( + &self, + #[tool(param)] + #[schemars(description = "The query string to search for types")] + query: String, + ) -> Result { + let embedding = self + .embedding_model + .embed(vec![&query], None) + .expect("Failed to get embedding") + .pop() + .expect("Embedding is empty") + .into_iter() + .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(8) + .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() })), + ) + })?; + + 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") + }) + .collect(), + )) + } + + #[tool(description = "Search Relation Types")] + async fn search_relation_types( + &self, + #[tool(param)] + #[schemars(description = "The query string to search for relation types")] + query: String, + ) -> Result { + let embedding = self + .embedding_model + .embed(vec![&query], None) + .expect("Failed to get embedding") + .pop() + .expect("Embedding is empty") + .into_iter() + .map(|v| v as f64) + .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), + ), + ) + .limit(8) + .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() })), + ) + })?; + + 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") + }) + .collect(), + )) + } + + #[tool(description = "Search Properties")] + async fn search_properties( + &self, + #[tool(param)] + #[schemars(description = "The query string to search for properties")] + query: String, + ) -> Result { + let embedding = self + .embedding_model + .embed(vec![&query], None) + .expect("Failed to get embedding") + .pop() + .expect("Embedding is empty") + .into_iter() + .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)), + ) + .limit(8) + .send() + .await + .map_err(|e| { + McpError::internal_error( + "search_properties", + Some(json!({ "error": e.to_string() })), + ) + })? + .try_collect::>() + .await + .map_err(|e| { + McpError::internal_error( + "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") + }) + .collect(), + )) + } + + // #[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( + &self, + #[tool(param)] + #[schemars( + description = "Return an entity by its ID along with its attributes (name, description, etc.) and types" + )] + id: String, + ) -> Result { + let entity = entity::find_one::>(&self.neo4j, &id) + .send() + .await + .map_err(|e| { + McpError::internal_error("get_entity", Some(json!({ "error": e.to_string() }))) + })? + .ok_or_else(|| { + McpError::internal_error("entity_not_found", Some(json!({ "id": id }))) + })?; + + tracing::info!("Found entity with ID '{}'", id); + + Ok(CallToolResult::success(vec![ + Content::json(json!({ + "id": entity.id(), + "name": entity.attributes.name, + "description": entity.attributes.description, + "types": entity.types, + })) + .expect("Failed to create JSON content"), + ])) + } +} + +#[tool(tool_box)] +impl ServerHandler for KnowledgeGraph { + fn get_info(&self) -> ServerInfo { + ServerInfo { + protocol_version: ProtocolVersion::V_2024_11_05, + capabilities: ServerCapabilities::builder() + .enable_prompts() + .enable_resources() + .enable_tools() + .build(), + server_info: Implementation::from_build_env(), + instructions: Some(include_str!("../resources/instructions.md").to_string()), + } + } + + async fn initialize( + &self, + _request: InitializeRequestParam, + context: RequestContext, + ) -> Result { + if let Some(http_request_part) = context.extensions.get::() { + let initialize_headers = &http_request_part.headers; + let initialize_uri = &http_request_part.uri; + tracing::info!(?initialize_headers, %initialize_uri, "initialize from http server"); + } + Ok(self.get_info()) + } +} + +#[derive(Debug, Parser)] +#[command(name = "stdout", version, about, arg_required_else_help = true)] +struct AppArgs { + #[clap(flatten)] + neo4j_args: Neo4jArgs, +} + +#[derive(Debug, Args)] +struct Neo4jArgs { + /// Neo4j database host + #[arg(long)] + neo4j_uri: String, + + /// Neo4j database user name + #[arg(long)] + neo4j_user: String, + + /// Neo4j database user password + #[arg(long)] + neo4j_pass: String, +} diff --git a/sink/src/bootstrap/boostrap_indexer.rs b/sink/src/bootstrap/boostrap_indexer.rs index 257424c..a2a10aa 100644 --- a/sink/src/bootstrap/boostrap_indexer.rs +++ b/sink/src/bootstrap/boostrap_indexer.rs @@ -54,8 +54,6 @@ pub fn triples() -> Vec { system_ids::NAME_ATTRIBUTE, "Space Kind", ), - // Triple::new(indexer_ids::SPACE_VERSION_COUNTER, system_ids::NAME_ATTRIBUTE, "Space Version Counter"), - // Member and Editor relations Triple::new( indexer_ids::MEMBER_RELATION,