feat(api): Semantic search on name and description (#47)

* feat(api): Semantic search on name and description

* fix: clippy + bump docker rust version

Author: cvauclair

Cargo.lock

@@ -23,15 +23,15 @@ docker run \
 ### 2. Compile and run the indexer
 In a separate terminal, run the following commands:
 ```bash
-cargo run --bin sink -- \
+CFLAGS='-std=gnu17' cargo run --bin sink -- \
     --reset-db \
     --neo4j-uri neo4j://localhost:7687 \
     --neo4j-user neo4j \
     --neo4j-pass neo4j
 ```
 
 ```bash
-cargo run --bin api -- \
+CFLAGS='-std=gnu17' cargo run --bin api -- \
     --neo4j-uri neo4j://localhost:7687 \
     --neo4j-user neo4j \
     --neo4j-pass neo4j

README.md

@@ -22,6 +22,7 @@ grc20-core = { version = "0.1.0", path = "../grc20-core" }
 grc20-sdk = { version = "0.1.0", path = "../grc20-sdk" }
 cache = { version = "0.1.0", path = "../cache" }
 chrono = "0.4.39"
+fastembed = "4.8.0"
 
 [dev-dependencies]
 serde_path_to_error = "0.1.16"

api/Cargo.toml

@@ -1,17 +1,30 @@
 use cache::KgCache;
+use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
 use grc20_core::neo4rs;
 use std::sync::Arc;
 
+const EMBEDDING_MODEL: EmbeddingModel = EmbeddingModel::AllMiniLML6V2;
+
 #[derive(Clone)]
 pub struct KnowledgeGraph {
     pub neo4j: Arc<neo4rs::Graph>,
     pub cache: Option<Arc<KgCache>>,
+    pub embedding_model: Arc<TextEmbedding>,
 }
 
 impl juniper::Context for KnowledgeGraph {}
 
 impl KnowledgeGraph {
     pub fn new(neo4j: Arc<neo4rs::Graph>, cache: Option<Arc<KgCache>>) -> Self {
-        Self { neo4j, cache }
+        Self {
+            neo4j,
+            cache,
+            embedding_model: Arc::new(
+                TextEmbedding::try_new(
+                    InitOptions::new(EMBEDDING_MODEL).with_show_download_progress(true),
+                )
+                .expect("Failed to initialize embedding model"),
+            ),
+        }
     }
 }

api/src/context.rs

@@ -317,4 +317,33 @@ impl RootQuery {
         .await?
         .map(|triple| Triple::new(triple, space_id, version_index)))
     }
+
+    async fn search<'a, S: ScalarValue>(
+        &'a self,
+        executor: &'a Executor<'_, '_, KnowledgeGraph, S>,
+        query: String,
+        #[graphql(default = 100)] first: i32,
+        // #[graphql(default = 0)] skip: i32,
+    ) -> FieldResult<Vec<Triple>> {
+        let embedding = executor
+            .context()
+            .embedding_model
+            .embed(vec![&query], None)
+            .expect("Failed to get embedding")
+            .pop()
+            .expect("Embedding is empty")
+            .into_iter()
+            .map(|v| v as f64)
+            .collect::<Vec<_>>();
+
+        let query = mapping::triple::semantic_search(&executor.context().neo4j, embedding)
+            .limit(first as usize);
+
+        Ok(query
+            .send()
+            .await?
+            .map_ok(|search_result| Triple::new(search_result.triple, search_result.space_id, None))
+            .try_collect::<Vec<_>>()
+            .await?)
+    }
 }

api/src/schema/query.rs

@@ -1,6 +1,6 @@
 use std::fmt::Display;
 
-use juniper::{graphql_object, Executor, GraphQLEnum, GraphQLObject, ScalarValue};
+use juniper::{graphql_object, Executor, FieldResult, GraphQLEnum, GraphQLObject, ScalarValue};
 
 use grc20_core::{
     mapping::{self, query_utils::Query, triple},
@@ -9,8 +9,11 @@ use grc20_core::{
 
 use crate::context::KnowledgeGraph;
 
+use super::Entity;
+
 #[derive(Debug)]
 pub struct Triple {
+    entity_id: String,
     pub attribute: String,
     pub value: String,
     pub value_type: ValueType,
@@ -23,6 +26,7 @@ pub struct Triple {
 impl Triple {
     pub fn new(triple: mapping::Triple, space_id: String, space_version: Option<String>) -> Self {
         Self {
+            entity_id: triple.entity,
             attribute: triple.attribute,
             value: triple.value.value,
             value_type: triple.value.value_type.into(),
@@ -79,6 +83,20 @@ impl Triple {
         .expect("Failed to find triple name attribute")
         .map(|triple| triple.value.value)
     }
+
+    async fn entity<'a, S: ScalarValue>(
+        &'a self,
+        executor: &'a Executor<'_, '_, KnowledgeGraph, S>,
+    ) -> FieldResult<Option<Entity>> {
+        Entity::load(
+            &executor.context().neo4j,
+            self.entity_id.clone(),
+            self.space_id.clone(),
+            self.space_version.clone(),
+            false,
+        )
+        .await
+    }
 }
 
 impl From<mapping::ValueType> for ValueType {

api/src/schema/triple.rs

@@ -1,10 +1,10 @@
-FROM rust:1.81.0 AS builder
+FROM rust:1.87.0 AS builder
 
 WORKDIR /kg-node
 COPY . .
 RUN apt-get update && apt-get upgrade -y
 RUN apt-get install libssl-dev protobuf-compiler -y
-RUN cargo build --release --bin sink --bin api
+RUN CFLAGS='-std=gnu17' cargo build --release --bin sink --bin api
 
 # Run image
 FROM debian:bookworm-slim AS run

docker/Dockerfile

@@ -5,3 +5,8 @@ pub mod network_ids;
 pub mod system_ids;
 
 pub use id::*;
+
+pub fn indexed(id: &str) -> bool {
+    // Add other ids to this list as needed
+    id == system_ids::DESCRIPTION_ATTRIBUTE || id == system_ids::NAME_ATTRIBUTE
+}

grc20-core/src/ids/mod.rs

@@ -837,15 +837,11 @@ mod tests {
         .await
         .expect("Insert failed");
 
-        Triple {
-            entity: "abc".to_string(),
-            attribute: "bar".to_string(),
-            value: 456u64.into(),
-        }
-        .insert(&neo4j, &BlockMetadata::default(), "space_id", "1")
-        .send()
-        .await
-        .expect("Failed to insert triple");
+        Triple::new("abc", "bar", 456u64)
+            .insert(&neo4j, &BlockMetadata::default(), "space_id", "1")
+            .send()
+            .await
+            .expect("Failed to insert triple");
 
         let foo_v2 = entity::find_one::<Entity<Foo>>(&neo4j, "abc")
             .space_id("space_id")

grc20-core/src/mapping/attributes.rs

@@ -168,11 +168,7 @@ mod tests {
         // Setup a local Neo 4J container for testing. NOTE: docker service must be running.
         let (_container, neo4j) = crate::test_utils::setup_neo4j().await;
 
-        let triple = Triple {
-            entity: "abc".to_string(),
-            attribute: "name".to_string(),
-            value: "Alice".into(),
-        };
+        let triple = Triple::new("abc", "name", "Alice");
 
         triple
             .insert(&neo4j, &BlockMetadata::default(), "space_id", "0")