diff --git a/langstream-agents/langstream-ai-agents/pom.xml b/langstream-agents/langstream-ai-agents/pom.xml
index b8d9474ba..0fc60bd9b 100644
--- a/langstream-agents/langstream-ai-agents/pom.xml
+++ b/langstream-agents/langstream-ai-agents/pom.xml
@@ -66,7 +66,17 @@
com.datastax.astra
astra-sdk-devops
- 0.6.9
+ 1.0
+
+
+ com.datastax.astra
+ astra-db-client
+ 1.0
+
+
+ com.datastax.astra
+ astra-sdk
+ 1.0
io.netty.incubator
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/datasource/impl/AstraDataSource.java b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/datasource/impl/CassandraDataSourceProvider.java
similarity index 94%
rename from langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/datasource/impl/AstraDataSource.java
rename to langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/datasource/impl/CassandraDataSourceProvider.java
index 70fcfb0b7..99cbce9ab 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/datasource/impl/AstraDataSource.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/datasource/impl/CassandraDataSourceProvider.java
@@ -20,7 +20,7 @@
import com.datastax.oss.streaming.ai.datasource.QueryStepDataSource;
import java.util.Map;
-public class AstraDataSource implements DataSourceProvider {
+public class CassandraDataSourceProvider implements DataSourceProvider {
@Override
public boolean supports(Map dataSourceConfig) {
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/datasource/CassandraDataSource.java b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/datasource/CassandraDataSource.java
index a3d2808e2..61e667765 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/datasource/CassandraDataSource.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/datasource/CassandraDataSource.java
@@ -31,9 +31,9 @@
import com.datastax.oss.driver.api.core.type.reflect.GenericType;
import com.datastax.oss.driver.internal.core.type.codec.CqlVectorCodec;
import com.datastax.oss.driver.internal.core.type.codec.registry.DefaultCodecRegistry;
-import com.dtsx.astra.sdk.db.AstraDbClient;
-import com.dtsx.astra.sdk.db.DatabaseClient;
-import com.dtsx.astra.sdk.utils.ApiLocator;
+import com.dtsx.astra.sdk.db.AstraDBOpsClient;
+import com.dtsx.astra.sdk.db.DbOpsClient;
+import com.dtsx.astra.sdk.utils.AstraEnvironment;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.io.ByteArrayInputStream;
import java.net.InetSocketAddress;
@@ -241,13 +241,13 @@ private CqlSession buildCqlSession(Map dataSourceConfig) {
log.info(
"Automatically downloading the secure bundle for database name {} from AstraDB",
astraDatabase);
- DatabaseClient databaseClient = this.buildAstraClient();
+ DbOpsClient databaseClient = this.buildAstraClient();
secureBundleDecoded = downloadSecureBundle(databaseClient);
} else if (!astraDatabaseId.isEmpty() && !astraToken.isEmpty()) {
log.info(
"Automatically downloading the secure bundle for database id {} from AstraDB",
astraDatabaseId);
- DatabaseClient databaseClient = this.buildAstraClient();
+ DbOpsClient databaseClient = this.buildAstraClient();
secureBundleDecoded = downloadSecureBundle(databaseClient);
} else {
log.info("No secure bundle provided, using the default CQL driver for Cassandra");
@@ -281,11 +281,11 @@ public CqlSession getSession() {
return session;
}
- public DatabaseClient buildAstraClient() {
+ public DbOpsClient buildAstraClient() {
return buildAstraClient(astraToken, astraDatabase, astraDatabaseId, astraEnvironment);
}
- public static DatabaseClient buildAstraClient(
+ public static DbOpsClient buildAstraClient(
String astraToken,
String astraDatabase,
String astraDatabaseId,
@@ -293,9 +293,8 @@ public static DatabaseClient buildAstraClient(
if (astraToken.isEmpty()) {
throw new IllegalArgumentException("You must configure the AstraDB token");
}
- AstraDbClient astraDbClient =
- new AstraDbClient(
- astraToken, ApiLocator.AstraEnvironment.valueOf(astraEnvironment));
+ AstraDBOpsClient astraDbClient =
+ new AstraDBOpsClient(astraToken, AstraEnvironment.valueOf(astraEnvironment));
if (!astraDatabase.isEmpty()) {
return astraDbClient.databaseByName(astraDatabase);
} else if (!astraDatabaseId.isEmpty()) {
@@ -306,7 +305,7 @@ public static DatabaseClient buildAstraClient(
}
}
- public static byte[] downloadSecureBundle(DatabaseClient databaseClient) {
+ public static byte[] downloadSecureBundle(DbOpsClient databaseClient) {
long start = System.currentTimeMillis();
byte[] secureBundleDecoded = databaseClient.downloadDefaultSecureConnectBundle();
long delta = System.currentTimeMillis() - start;
diff --git a/langstream-agents/langstream-ai-agents/src/main/resources/META-INF/services/ai.langstream.ai.agents.datasource.DataSourceProvider b/langstream-agents/langstream-ai-agents/src/main/resources/META-INF/services/ai.langstream.ai.agents.datasource.DataSourceProvider
index 3c023fedb..233a8ee8a 100644
--- a/langstream-agents/langstream-ai-agents/src/main/resources/META-INF/services/ai.langstream.ai.agents.datasource.DataSourceProvider
+++ b/langstream-agents/langstream-ai-agents/src/main/resources/META-INF/services/ai.langstream.ai.agents.datasource.DataSourceProvider
@@ -1,2 +1,2 @@
-ai.langstream.ai.agents.datasource.impl.AstraDataSource
+ai.langstream.ai.agents.datasource.impl.CassandraDataSourceProvider
ai.langstream.ai.agents.datasource.impl.JdbcDataSourceProvider
\ No newline at end of file
diff --git a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/astra/AstraCollectionsAssetsManagerProvider.java b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/astra/AstraCollectionsAssetsManagerProvider.java
new file mode 100644
index 000000000..c915b4ab3
--- /dev/null
+++ b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/astra/AstraCollectionsAssetsManagerProvider.java
@@ -0,0 +1,117 @@
+/*
+ * 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.
+ */
+package ai.langstream.agents.vector.astra;
+
+import ai.langstream.api.model.AssetDefinition;
+import ai.langstream.api.runner.assets.AssetManager;
+import ai.langstream.api.runner.assets.AssetManagerProvider;
+import ai.langstream.api.util.ConfigurationUtils;
+import io.stargate.sdk.doc.exception.CollectionNotFoundException;
+import java.util.Map;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class AstraCollectionsAssetsManagerProvider implements AssetManagerProvider {
+
+ @Override
+ public boolean supports(String assetType) {
+ return "astra-collection".equals(assetType);
+ }
+
+ @Override
+ public AssetManager createInstance(String assetType) {
+
+ switch (assetType) {
+ case "astra-collection":
+ return new AstraDBCollectionAssetManager();
+ default:
+ throw new IllegalArgumentException();
+ }
+ }
+
+ private abstract static class BaseAstraAssetManager implements AssetManager {
+
+ AstraCollectionsDataSource datasource;
+ AssetDefinition assetDefinition;
+
+ @Override
+ public void initialize(AssetDefinition assetDefinition) {
+ this.datasource = buildDataSource(assetDefinition);
+ this.assetDefinition = assetDefinition;
+ }
+
+ @Override
+ public void close() throws Exception {
+ if (datasource != null) {
+ datasource.close();
+ }
+ }
+ }
+
+ private static class AstraDBCollectionAssetManager extends BaseAstraAssetManager {
+
+ @Override
+ public boolean assetExists() throws Exception {
+ String collection = getCollection();
+ log.info("Checking if collection {} exists", collection);
+ return datasource.getAstraDB().isCollectionExists(collection);
+ }
+
+ @Override
+ public void deployAsset() throws Exception {
+ int vectorDimension = getVectorDimension();
+
+ String collection = getCollection();
+ log.info("Create collection {} with vector dimension {}", collection, vectorDimension);
+ datasource.getAstraDB().createCollection(collection, vectorDimension);
+ }
+
+ private String getCollection() {
+ return ConfigurationUtils.getString("collection", null, assetDefinition.getConfig());
+ }
+
+ private int getVectorDimension() {
+ return ConfigurationUtils.getInt("vector-dimension", 1536, assetDefinition.getConfig());
+ }
+
+ @Override
+ public boolean deleteAssetIfExists() throws Exception {
+ String collection = getCollection();
+
+ log.info("Deleting collection {}", collection);
+
+ try {
+ datasource.getAstraDB().deleteCollection(collection);
+ return true;
+ } catch (CollectionNotFoundException e) {
+ log.info(
+ "collection does not exist, maybe it was deleted by another agent ({})",
+ e.toString());
+ return false;
+ }
+ }
+ }
+
+ private static AstraCollectionsDataSource buildDataSource(AssetDefinition assetDefinition) {
+ AstraCollectionsDataSource dataSource = new AstraCollectionsDataSource();
+ Map datasourceDefinition =
+ ConfigurationUtils.getMap("datasource", Map.of(), assetDefinition.getConfig());
+ Map configuration =
+ ConfigurationUtils.getMap("configuration", Map.of(), datasourceDefinition);
+ dataSource.initialize(configuration);
+ return dataSource;
+ }
+}
diff --git a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/astra/AstraCollectionsDataSource.java b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/astra/AstraCollectionsDataSource.java
new file mode 100644
index 000000000..1ed0caee2
--- /dev/null
+++ b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/astra/AstraCollectionsDataSource.java
@@ -0,0 +1,76 @@
+/*
+ * 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.
+ */
+package ai.langstream.agents.vector.astra;
+
+import ai.langstream.api.util.ConfigurationUtils;
+import com.datastax.oss.streaming.ai.datasource.QueryStepDataSource;
+import com.dtsx.astra.sdk.AstraDB;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class AstraCollectionsDataSource implements QueryStepDataSource {
+
+ AstraDB astraDB;
+
+ @Override
+ public void initialize(Map dataSourceConfig) {
+ log.info(
+ "Initializing CassandraDataSource with config {}",
+ ConfigurationUtils.redactSecrets(dataSourceConfig));
+ String astraToken = ConfigurationUtils.getString("token", "", dataSourceConfig);
+ String astraEndpoint = ConfigurationUtils.getString("endpoint", "", dataSourceConfig);
+ this.astraDB = new AstraDB(astraToken, astraEndpoint);
+ }
+
+ @Override
+ public void close() {}
+
+ @Override
+ public List