Swagger vector search (#511)

stargate · Aug 21, 2023 · de9b75d · de9b75d
1 parent 30cfb25
commit de9b75d
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 3 deletions.
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/StargateJsonApi.java b/src/main/java/io/stargate/sgv2/jsonapi/StargateJsonApi.java
@@ -80,6 +80,20 @@
                         }
                       }
                       """),
+              @ExampleObject(
+                  name = "findOneVectorSearch",
+                  summary = "`findOne` command with vector search",
+                  value =
+                      """
+                    {
+                      "findOne": {
+                          "filter": {"location": "London", "race.competitors" : {"$eq" : 100}},
+                          "projection": {"_id":0, "location":1, "race.start_date":1, "tags":1},
+                          "sort" : {"$vector" : [0.25,0.25,0.25,0.25,0.25]},
+                          "options": {"includeSimilarity" : true}
+                      }
+                    }
+                    """),
               @ExampleObject(
                   name = "countDocuments",
                   summary = "`countDocuments` command",
@@ -105,6 +119,20 @@
                         }
                       }
                       """),
+              @ExampleObject(
+                  name = "findVectorSearch",
+                  summary = "`find` command with vector search",
+                  value =
+                      """
+                      {
+                        "find": {
+                           "filter": {"location": "London", "race.competitors" : {"$eq" : 100}},
+                           "projection": {"tags":0},
+                           "sort" : {"$vector" : [0.25,0.25,0.25,0.25,0.25]},
+                           "options": {"limit" : 1000, "includeSimilarity" : true}
+                        }
+                      }
+                    """),
               @ExampleObject(
                   name = "findOneAndUpdate",
                   summary = "`findOneAndUpdate` command",
@@ -234,7 +262,8 @@
                               "competitors": 100,
                               "start_date": "2022-08-15"
                             },
-                            "tags": [ ]
+                            "tags": [ ],
+                            "$vector" : [0.25,0.25,0.25,0.25,0.25]
                           }
                         }
                       }
@@ -254,7 +283,8 @@
                                     "competitors": 100,
                                     "start_date": "2022-08-15"
                                   },
-                                  "tags" : [ "eu" ]
+                                  "tags" : [ "eu" ],
+                                  "$vector" : [0.35,0.35,0.35,0.35,0.35]
                                 },
                                 {
                                   "_id": "2",
@@ -263,7 +293,8 @@
                                     "competitors": 150,
                                     "start_date": "2022-09-15"
                                   },
-                                  "tags": [ "us" ]
+                                  "tags": [ "us" ],
+                                  "$vector" : [0.45,0.45,0.45,0.45,0.45]
                                 }
                             ],
                             "options": {
@@ -332,6 +363,23 @@
                             }
                         }
                       """),
+              @ExampleObject(
+                  name = "createCollectionVectorSearch",
+                  summary = "`CreateCollection` command with vector search",
+                  value =
+                      """
+                        {
+                            "createCollection": {
+                              "name": "events_vector",
+                              "options": {
+                                "vector": {
+                                  "size": 2,
+                                  "function": "cosine"
+                                }
+                              }
+                            }
+                        }
+                     """),
               @ExampleObject(
                   name = "findCollections",
                   summary = "`FindCollections` command",

diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/CollectionResource.java b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/CollectionResource.java
@@ -99,7 +99,9 @@ public CollectionResource(MeteredCommandProcessor meteredCommandProcessor) {
                 @ExampleObject(ref = "deleteOne"),
                 @ExampleObject(ref = "deleteMany"),
                 @ExampleObject(ref = "find"),
+                @ExampleObject(ref = "findVectorSearch"),
                 @ExampleObject(ref = "findOne"),
+                @ExampleObject(ref = "findOneVectorSearch"),
                 @ExampleObject(ref = "findOneAndDelete"),
                 @ExampleObject(ref = "findOneAndReplace"),
                 @ExampleObject(ref = "findOneAndUpdate"),

diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/NamespaceResource.java b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/NamespaceResource.java
@@ -58,6 +58,7 @@ public NamespaceResource(MeteredCommandProcessor meteredCommandProcessor) {
               schema = @Schema(anyOf = {CreateCollectionCommand.class}),
               examples = {
                 @ExampleObject(ref = "createCollection"),
+                @ExampleObject(ref = "createCollectionVectorSearch"),
                 @ExampleObject(ref = "findCollections"),
                 @ExampleObject(ref = "deleteCollection"),
               }))