feat: add vectorindex proto (#185)

momentohq · Aug 8, 2023 · 43ec538 · 43ec538
1 parent dd54989
commit 43ec538
Show file tree

Hide file tree

Showing 7 changed files with 69 additions and 4 deletions.
diff --git a/csharp/Momento.Protos.csproj b/csharp/Momento.Protos.csproj
@@ -19,6 +19,7 @@
     <Protobuf Include="../proto/cacheclient.proto" GrpcServices="Client" />
     <Protobuf Include="../proto/controlclient.proto" GrpcServices="Client" />
     <Protobuf Include="../proto/cacheping.proto" GrpcServices="Client" />
+    <Protobuf Include="../proto/vectorindex.proto" GrpcServices="Client" />
   </ItemGroup>
 
   <ItemGroup>

diff --git a/javascript-web/generate_protos.sh b/javascript-web/generate_protos.sh
@@ -65,7 +65,7 @@ mkdir $out
 # So we do a terrible hack to comment out the package declaration before generating the JS types,
 # but add them back before generating the GRPC web bindings
 
-proto_file_list=" extensions.proto cacheclient.proto controlclient.proto auth.proto cacheping.proto cachepubsub.proto "
+proto_file_list=" extensions.proto cacheclient.proto controlclient.proto auth.proto cacheping.proto cachepubsub.proto vectorindex.proto "
 
 echo "Backing up protos dir"
 cp -r ../proto ../proto.bak

diff --git a/javascript/README.md b/javascript/README.md
@@ -2,5 +2,5 @@
 ```bash
 npm install
 mkdir dist
-PATH=node_modules/protoc-gen-ts/bin/:$PATH protoc -I=../proto --ts_out=dist cacheclient.proto controlclient.proto cachepubsub.proto
+PATH=node_modules/protoc-gen-ts/bin/:$PATH protoc -I=../proto --ts_out=dist cacheclient.proto controlclient.proto cachepubsub.proto vectorindex.proto
 ```
diff --git a/javascript/generate_protos.sh b/javascript/generate_protos.sh
@@ -2,4 +2,4 @@
 set -e
 set -x
 
-PATH=node_modules/protoc-gen-ts/bin/:$PATH protoc -I=../proto -I=/usr/local/include --ts_out=src cacheclient.proto controlclient.proto cachepubsub.proto auth.proto cacheping.proto
+PATH=node_modules/protoc-gen-ts/bin/:$PATH protoc -I=../proto -I=/usr/local/include --ts_out=src cacheclient.proto controlclient.proto cachepubsub.proto auth.proto cacheping.proto vectorindex.proto
diff --git a/proto/vectorindex.proto b/proto/vectorindex.proto
@@ -0,0 +1,63 @@
+syntax = "proto3";
+
+package vectorindex;
+
+service VectorIndex {
+    rpc UpsertItemBatch(UpsertItemBatchRequest) returns (UpsertItemBatchResponse) {}
+    rpc Search(SearchRequest) returns (SearchResponse) {}
+}
+
+message Item {
+    string id = 1;
+    Vector vector = 2;
+    repeated Metadata metadata = 3;
+}
+
+message UpsertItemBatchRequest {
+    string index_name = 1;
+    repeated Item items = 2;
+}
+
+message UpsertItemBatchResponse {
+    repeated uint32 error_indices = 1;
+}
+
+message Vector {
+    repeated float elements = 1;
+}
+
+message Metadata {
+    string field = 1;
+    oneof value {
+        // Eventually can support ints, dates, etc
+        string string_value = 2;
+    }
+}
+
+message MetadataRequest {
+    message Some {
+        repeated string fields = 1;
+    }
+    message All {}
+
+    oneof kind {
+        Some some = 2;
+    }
+}
+
+message SearchRequest {
+    string index_name = 1;
+    uint32 top_k = 2;
+    Vector query_vector = 3;
+    MetadataRequest metadata_fields = 4;
+}
+
+message SearchHit {
+    string id = 1;
+    float distance = 2;
+    repeated Metadata metadata = 3;
+}
+
+message SearchResponse {
+    repeated SearchHit hits = 1;
+}
diff --git a/python/run-protoc.sh b/python/run-protoc.sh
@@ -23,7 +23,7 @@ do
     if [[ $python_protobuf_version == "protobuf>4" ]]; then
         pyi_out="--pyi_out=$src_path"
     fi
-    poetry run python -m grpc_tools.protoc -I../proto --python_out=$src_path $pyi_out --grpc_python_out=$src_path extensions.proto cacheclient.proto controlclient.proto auth.proto cachepubsub.proto
+    poetry run python -m grpc_tools.protoc -I../proto --python_out=$src_path $pyi_out --grpc_python_out=$src_path extensions.proto cacheclient.proto controlclient.proto auth.proto cachepubsub.proto vectorindex.proto
 
     # A shortcoming of the generated code is in the grpc generated code,
     # the protobuf imports are absolute instead of relative.

diff --git a/rust/proto_generator/build.rs b/rust/proto_generator/build.rs
@@ -19,6 +19,7 @@ fn main() {
                 format!("{proto_dir}/cacheclient.proto"),
                 format!("{proto_dir}/cachepubsub.proto"),
                 format!("{proto_dir}/controlclient.proto"),
+                format!("{proto_dir}/vectorindex.proto"),
             ],
             &[proto_dir],
         )