Quellcode durchsuchen

feat: client library — similaritySearch + vector-aware createCollection

fszontagh vor 4 Monaten
Ursprung
Commit
8d78a21efc
2 geänderte Dateien mit 72 neuen und 4 gelöschten Zeilen
  1. 24 1
      client/include/smartbotic/database/client.hpp
  2. 48 3
      client/src/client.cpp

+ 24 - 1
client/include/smartbotic/database/client.hpp

@@ -196,10 +196,33 @@ public:
     [[nodiscard]] std::vector<std::string> setMembers(const std::string& collection, const std::string& setId);
     [[nodiscard]] bool setIsMember(const std::string& collection, const std::string& setId, const std::string& member);
 
+    // ===== Vector Search =====
+
+    struct SimilarityResult {
+        std::string id;
+        float score;
+        nlohmann::json data;
+    };
+
+    /**
+     * Perform a vector similarity search.
+     * @param collection Collection name (must have vector_dimension set)
+     * @param queryVector Query vector (must match collection's vector_dimension)
+     * @param topK Maximum number of results to return
+     * @param minScore Minimum similarity score threshold (0.0 = return all)
+     * @return Ordered list of results by descending similarity score
+     */
+    [[nodiscard]] std::vector<SimilarityResult> similaritySearch(
+        const std::string& collection,
+        const std::vector<float>& queryVector,
+        uint32_t topK,
+        float minScore = 0.0f);
+
     // ===== Collection Management =====
 
     bool createCollection(const std::string& name, uint32_t defaultTtlSeconds = 0,
-                          bool encrypted = false, uint32_t maxVersions = 0);
+                          bool encrypted = false, uint32_t maxVersions = 0,
+                          uint32_t vectorDimension = 0);
     bool dropCollection(const std::string& name);
     [[nodiscard]] std::vector<std::string> listCollections();
 

+ 48 - 3
client/src/client.cpp

@@ -467,8 +467,44 @@ public:
 
     // ===== Collection Management =====
 
+    std::vector<Client::SimilarityResult> similaritySearch(
+        const std::string& collection, const std::vector<float>& queryVector,
+        uint32_t topK, float minScore) {
+        smartbotic::databasepb::SimilaritySearchRequest request;
+        request.set_collection(collection);
+        for (float v : queryVector) {
+            request.add_query_vector(v);
+        }
+        request.set_top_k(topK);
+        request.set_min_score(minScore);
+
+        smartbotic::databasepb::SimilaritySearchResponse response;
+        grpc::ClientContext context;
+        setDeadline(context);
+
+        auto status = stub_->SimilaritySearch(&context, request, &response);
+        if (!status.ok()) {
+            spdlog::error("Client::similaritySearch failed: {}", status.error_message());
+            throw std::runtime_error(status.error_message());
+        }
+
+        std::vector<Client::SimilarityResult> results;
+        results.reserve(response.results_size());
+        for (const auto& r : response.results()) {
+            Client::SimilarityResult entry;
+            entry.id = r.id();
+            entry.score = r.score();
+            if (!r.data().empty()) {
+                entry.data = nlohmann::json::parse(r.data());
+            }
+            results.push_back(std::move(entry));
+        }
+
+        return results;
+    }
+
     bool createCollection(const std::string& name, uint32_t defaultTtlSeconds, bool encrypted,
-                          uint32_t maxVersions) {
+                          uint32_t maxVersions, uint32_t vectorDimension) {
         smartbotic::databasepb::CreateCollectionRequest request;
         request.set_name(name);
         auto* options = request.mutable_options();
@@ -479,6 +515,9 @@ public:
         if (maxVersions > 0) {
             options->set_max_versions(maxVersions);
         }
+        if (vectorDimension > 0) {
+            options->set_vector_dimension(vectorDimension);
+        }
 
         smartbotic::databasepb::CreateCollectionResponse response;
         grpc::ClientContext context;
@@ -972,9 +1011,15 @@ bool Client::setIsMember(const std::string& collection, const std::string& setId
     return impl_->setIsMember(collection, setId, member);
 }
 
+std::vector<Client::SimilarityResult> Client::similaritySearch(
+    const std::string& collection, const std::vector<float>& queryVector,
+    uint32_t topK, float minScore) {
+    return impl_->similaritySearch(collection, queryVector, topK, minScore);
+}
+
 bool Client::createCollection(const std::string& name, uint32_t defaultTtlSeconds,
-                              bool encrypted, uint32_t maxVersions) {
-    return impl_->createCollection(name, defaultTtlSeconds, encrypted, maxVersions);
+                              bool encrypted, uint32_t maxVersions, uint32_t vectorDimension) {
+    return impl_->createCollection(name, defaultTtlSeconds, encrypted, maxVersions, vectorDimension);
 }
 
 bool Client::dropCollection(const std::string& name) {