Couchbase-parameters (#46)

Updated couchbase to allow for scope and collection names to be passed in order to write to serveral different collections for experiements. *Trying the vscode PR extension Co-authored-by: benfrank241 <[email protected]>
vectorize-io · Jun 3, 2024 · 2ccb779 · 2ccb779
1 parent cde8cea
commit 2ccb779
Show file tree

Hide file tree

Showing 9 changed files with 63 additions and 28 deletions.
diff --git a/conf/cli.yaml b/conf/cli.yaml
@@ -1,19 +1,3 @@
-#
-# Copyright DataStax, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
 ---
 webServiceUrl: "http://localhost:8090"
 apiGatewayUrl: "ws://localhost:8091"

diff --git a/examples/applications/query-couchbase/README.md b/examples/applications/query-couchbase/README.md
@@ -38,16 +38,22 @@ couchbases://cb.shnnjztaidekg6i.cloud.couchbase.com
 
 The above is an example of a connection string.
 
+## Create a scope
+
+Navigate to Databases > Data tools tab.
+On the right hand side hover over the bucket of your choice and click the ellipsis that appears when you hover over it and create a scope.
+Do the same but with the scope you just created to create a collection.
+
 ## Create a search index in the database in the Couchbase Capella UI
 
 1. **Select Your Cluster:**
    - From the dashboard, select the cluster where you want to create the search index.
 
-2. **Open the "Indexes" Tab:**
-   - In the cluster's overview page, go to the "Indexes" tab which can be found in the sidebar.
+2. **Open the "Search" Tab:**
+   - In the cluster's overview page, go to the "Search" tab which can be found in the sidebar.
 
 3. **Create a New Search Index:**
-   - Click on the "Add Index" button.
+   - Click on the "Create Search Index" button.
    - Choose "Full Text Search" from the index type options.
 
 4. **Configure the Index:**
@@ -79,7 +85,7 @@ Using the docker image:
 bin/langstream gateway produce test write-topic -v "{\"id\":\"myid\",\"document\":\"Kafkaesque: extremely unpleasant, frightening, and confusing, and similar to situations described in the novels of Franz Kafka.\"}" -p sessionId=$(uuidgen)
 
 ```
-You can view the uploaded document in the _default scope and _default collection of the bucket you selected.
+You can view the uploaded document in the example scope and default collection of the bucket you selected.
 
 ## Start a chat using the gateway to query the document
 

diff --git a/examples/applications/query-couchbase/configuration.yaml b/examples/applications/query-couchbase/configuration.yaml
@@ -31,3 +31,5 @@ configuration:
         username: "${secrets.couchbase.username}"
         password: "${secrets.couchbase.password}"
         bucket-name: "${secrets.couchbase.bucket-name}"
+        scope-name: "example"
+        collection-name: "default"
diff --git a/examples/applications/query-couchbase/write.yaml b/examples/applications/query-couchbase/write.yaml
@@ -38,5 +38,8 @@ pipeline:
       username: "${secrets.couchbase.username}"
       password: "${secrets.couchbase.password}"
       connection-string: "${secrets.couchbase.connection-string}"
+      scope-name: "example"
+      collection-name: "default"
       vector.id: "value.id"
       vector.vector: "value.embeddings"
+
diff --git a/...ector-agents/src/main/java/ai/langstream/agents/vector/couchbase/CouchbaseDataSource.java b/...ector-agents/src/main/java/ai/langstream/agents/vector/couchbase/CouchbaseDataSource.java
@@ -59,6 +59,12 @@ public static final class CouchbaseConfig {
 
         @JsonProperty(value = "password", required = true)
         private String password;
+
+        @JsonProperty(value = "scope-name", required = true)
+        private String scopeName;
+
+        @JsonProperty(value = "collection-name", required = true)
+        private String collectionName;
     }
 
     @Override
@@ -100,17 +106,23 @@ public List<Map<String, Object>> fetchData(String query, List<Object> params) {
 
                 float[] vector = JstlFunctions.toArrayOfFloat(queryMap.remove("vector"));
                 Integer topK = (Integer) queryMap.remove("topK");
+                // scope namen comes from querymap
 
                 SearchRequest request =
                         SearchRequest.create(
                                 VectorSearch.create(
                                         VectorQuery.create("embeddings", vector)
                                                 .numCandidates(topK)));
+                log.debug("SearchRequest created: {}", request);
 
                 SearchResult result =
                         cluster.search(
-                                "" + clientConfig.bucketName + "._default.vector-search",
-                                request); //
+                                ""
+                                        + clientConfig.bucketName
+                                        + "."
+                                        + clientConfig.scopeName
+                                        + ".vector-search",
+                                request);
 
                 return result.rows().stream()
                         .map(

diff --git a/...am-vector-agents/src/main/java/ai/langstream/agents/vector/couchbase/CouchbaseWriter.java b/...am-vector-agents/src/main/java/ai/langstream/agents/vector/couchbase/CouchbaseWriter.java
@@ -61,10 +61,10 @@ public CouchbaseDatabaseWriter(Map<String, Object> datasourceConfig) {
             String username = (String) datasourceConfig.get("username");
             String password = (String) datasourceConfig.get("password");
             String bucketName = (String) datasourceConfig.get("bucket-name");
-            String scopeName = (String) datasourceConfig.getOrDefault("scopeName", "_default");
+            String scopeName = (String) datasourceConfig.getOrDefault("scope-name", "_default");
             String connectionString = (String) datasourceConfig.get("connection-string");
             String collectionName =
-                    (String) datasourceConfig.getOrDefault("collectionName", "_default");
+                    (String) datasourceConfig.getOrDefault("collection-name", "_default");
 
             // Create a cluster with the WAN profile
             ClusterOptions clusterOptions =

diff --git a/...core/src/main/java/ai/langstream/impl/resources/datasource/CouchbaseDatasourceConfig.java b/...core/src/main/java/ai/langstream/impl/resources/datasource/CouchbaseDatasourceConfig.java
@@ -82,4 +82,20 @@ public void validate(Resource resource) {
                                     """)
     @JsonProperty("bucket-name")
     private String bucketName;
+
+    @ConfigProperty(
+            description =
+                    """
+                            Scope to use in Couchbase.
+                                    """)
+    @JsonProperty("scope-name")
+    private String scopeName;
+
+    @ConfigProperty(
+            description =
+                    """
+                            Collection to use in Couchbase.
+                                    """)
+    @JsonProperty("collection-name")
+    private String collectionName;
 }
diff --git a/...va/ai/langstream/runtime/impl/k8s/agents/vectors/CouchbaseVectorDatabaseWriterConfig.java b/...va/ai/langstream/runtime/impl/k8s/agents/vectors/CouchbaseVectorDatabaseWriterConfig.java
@@ -62,8 +62,18 @@ public boolean isAgentConfigModelAllowUnknownProperties() {
     private String username;
 
     @ConfigProperty(
-            description = "The password to connect to the Couchbase cluster.",
-            required = true)
-    @JsonProperty("password")
-    private String password;
+            description =
+                    """
+                            Scope to use in Couchbase.
+                                    """)
+    @JsonProperty("scope-name")
+    private String scopeName;
+
+    @ConfigProperty(
+            description =
+                    """
+                            Collection to use in Couchbase.
+                                    """)
+    @JsonProperty("collection-name")
+    private String collectionName;
 }
diff --git a/...e/src/test/java/ai/langstream/runtime/impl/k8s/agents/QueryVectorDBAgentProviderTest.java b/...e/src/test/java/ai/langstream/runtime/impl/k8s/agents/QueryVectorDBAgentProviderTest.java
@@ -138,6 +138,8 @@ public void testWriteCouchbase() {
                                 bucket-name: "vectorize"
                                 username: ""
                                 password: ""
+                                scope-name: "example"
+                                collection-name: "default"
                                 connection-string: "couchbases://"
                                 vector.id: "value.id"
                                 vector.vector: "value.embeddings"