Nidx shards API & tests (#2609)

nuclia · Nov 12, 2024 · c29bd5f · c29bd5f
1 parent 8d71b6f
commit c29bd5f
Show file tree

Hide file tree

Showing 19 changed files with 474 additions and 249 deletions.
diff --git a/nidx/.sqlx/query-ef25d1b091b8edce55dcd7f0a2f528537cd9329ab7297b81ace0aba8c0e1bb0e.json b/nidx/.sqlx/query-ef25d1b091b8edce55dcd7f0a2f528537cd9329ab7297b81ace0aba8c0e1bb0e.json
diff --git a/...e16de56ea704d4b1c8784dafcce996983342.json → ...b475f6af02768c22f0e583c74062e2e821ce.json b/...e16de56ea704d4b1c8784dafcce996983342.json → ...b475f6af02768c22f0e583c74062e2e821ce.json
diff --git a/nidx/Cargo.lock b/nidx/Cargo.lock
diff --git a/nidx/Cargo.toml b/nidx/Cargo.toml
@@ -31,6 +31,9 @@ nidx_text = { version = "0.1.0", path = "nidx_text" }
 nidx_paragraph = { version = "0.1.0", path = "nidx_paragraph" }
 nidx_relation = { version = "0.1.0", path = "nidx_relation" }
 config = { version = "0.14.1", default-features = false }
+thiserror = "2.0.0"
+tower = "0.5.1"
+http = "1.1.0"
 
 [dev-dependencies]
 nidx_tests = { path = "nidx_tests" }

diff --git a/nidx/nidx_protos/build.rs b/nidx/nidx_protos/build.rs
@@ -22,13 +22,15 @@ use std::io::Result;
 
 fn main() -> Result<()> {
     println!("cargo:rerun-if-changed=../../nucliadb_protos");
+    println!("cargo:rerun-if-changed=src");
     tonic_build::configure().emit_rerun_if_changed(false).compile_protos(
         &[
+            "src/nidx.proto",
             "../../nucliadb_protos/nodereader.proto",
             "../../nucliadb_protos/noderesources.proto",
             "../../nucliadb_protos/nodewriter.proto",
         ],
-        &["../../"],
+        &["../../", "src"],
     )?;
     Ok(())
 }
diff --git a/nidx/nidx_protos/src/lib.rs b/nidx/nidx_protos/src/lib.rs
@@ -34,6 +34,10 @@ pub mod nodewriter {
     tonic::include_proto!("nodewriter");
 }
 
+pub mod nidx {
+    tonic::include_proto!("nidx");
+}
+
 pub use nodereader::*;
 pub use noderesources::*;
 pub use nodewriter::*;

diff --git a/nidx/nidx_protos/src/nidx.proto b/nidx/nidx_protos/src/nidx.proto
@@ -0,0 +1,49 @@
+syntax = "proto3";
+
+package nidx;
+
+import "nucliadb_protos/nodereader.proto";
+import "nucliadb_protos/nodewriter.proto";
+
+service NidxApi {
+  rpc GetShard(nodereader.GetShardRequest) returns (noderesources.Shard) {}
+
+  // rpc GetShardFiles(GetShardFilesRequest) returns (ShardFileList) {};
+  // rpc DownloadShardFile(DownloadShardFileRequest) returns (stream
+  // ShardFileChunk) {};
+
+  rpc NewShard(nodewriter.NewShardRequest) returns (noderesources.ShardCreated) {}
+  rpc DeleteShard(noderesources.ShardId) returns (noderesources.ShardId) {}
+  rpc ListShards(noderesources.EmptyQuery) returns (noderesources.ShardIds) {}
+  // rpc GC(noderesources.ShardId) returns (GarbageCollectorResponse) {}
+  // rpc Merge(noderesources.ShardId) returns (MergeResponse) {}
+
+  // rpc SetResource(noderesources.Resource) returns (OpStatus) {}
+  // rpc SetResourceFromStorage(IndexMessage) returns (OpStatus) {}
+  // rpc RemoveResource(noderesources.ResourceID) returns (OpStatus) {}
+
+  rpc AddVectorSet(nodewriter.NewVectorSetRequest) returns (nodewriter.OpStatus) {} // New
+  rpc RemoveVectorSet(noderesources.VectorSetID) returns (nodewriter.OpStatus) {}
+  rpc ListVectorSets(noderesources.ShardId) returns (noderesources.VectorSetList) {}
+
+  rpc GetMetadata(noderesources.EmptyQuery) returns (noderesources.NodeMetadata) {}
+}
+
+service NidxSearcher {
+  rpc DocumentSearch(nodereader.DocumentSearchRequest) returns (nodereader.DocumentSearchResponse) {}
+  rpc ParagraphSearch(nodereader.ParagraphSearchRequest) returns (nodereader.ParagraphSearchResponse) {}
+  rpc VectorSearch(nodereader.VectorSearchRequest) returns (nodereader.VectorSearchResponse) {}
+  rpc RelationSearch(nodereader.RelationSearchRequest) returns (nodereader.RelationSearchResponse) {}
+  rpc DocumentIds(noderesources.ShardId) returns (nodereader.IdCollection) {}
+  rpc ParagraphIds(noderesources.ShardId) returns (nodereader.IdCollection) {}
+  rpc VectorIds(noderesources.VectorSetID) returns (nodereader.IdCollection) {}
+  rpc RelationIds(noderesources.ShardId) returns (nodereader.IdCollection) {}
+  rpc RelationEdges(noderesources.ShardId) returns (nodereader.EdgeList) {}
+
+  rpc Search(nodereader.SearchRequest) returns (nodereader.SearchResponse) {}
+  rpc Suggest(nodereader.SuggestRequest) returns (nodereader.SuggestResponse) {}
+
+  // Streams
+  rpc Paragraphs(nodereader.StreamRequest) returns (stream nodereader.ParagraphItem) {}
+  rpc Documents(nodereader.StreamRequest) returns (stream nodereader.DocumentItem) {}
+}