From 72a945b787c89f128ac8274eafed093a7deade7d Mon Sep 17 00:00:00 2001 From: Jan Martiska Date: Fri, 8 Dec 2023 09:41:39 +0100 Subject: [PATCH] Allow injecting non-default datasource for PgVector store, produce an EmbeddingStoreBuildItem --- .../quarkus-langchain4j-pgvector.adoc | 17 +++++++++++++ .../Langchain4jPgvectorProcessor.java | 24 +++++++++++++++---- ...PgVectorEmbeddingStoreBuildTimeConfig.java | 20 ++++++++++++++++ .../PgVectorEmbeddingStoreRecorder.java | 10 +++++--- 4 files changed, 64 insertions(+), 7 deletions(-) create mode 100644 pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/PgVectorEmbeddingStoreBuildTimeConfig.java diff --git a/docs/modules/ROOT/pages/includes/quarkus-langchain4j-pgvector.adoc b/docs/modules/ROOT/pages/includes/quarkus-langchain4j-pgvector.adoc index 1e5cfcce8..bcfd5c448 100644 --- a/docs/modules/ROOT/pages/includes/quarkus-langchain4j-pgvector.adoc +++ b/docs/modules/ROOT/pages/includes/quarkus-langchain4j-pgvector.adoc @@ -10,6 +10,23 @@ h|[[quarkus-langchain4j-pgvector_configuration]]link:#quarkus-langchain4j-pgvect h|Type h|Default +a|icon:lock[title=Fixed at build time] [[quarkus-langchain4j-pgvector_quarkus.langchain4j.pgvector.datasource]]`link:#quarkus-langchain4j-pgvector_quarkus.langchain4j.pgvector.datasource[quarkus.langchain4j.pgvector.datasource]` + + +[.description] +-- +The name of the configured Postgres datasource to use for this store. If not set, the default datasource from the Agroal extension will be used. + +ifdef::add-copy-button-to-env-var[] +Environment variable: env_var_with_copy_button:+++QUARKUS_LANGCHAIN4J_PGVECTOR_DATASOURCE+++[] +endif::add-copy-button-to-env-var[] +ifndef::add-copy-button-to-env-var[] +Environment variable: `+++QUARKUS_LANGCHAIN4J_PGVECTOR_DATASOURCE+++` +endif::add-copy-button-to-env-var[] +--|string +| + + a| [[quarkus-langchain4j-pgvector_quarkus.langchain4j.pgvector.table]]`link:#quarkus-langchain4j-pgvector_quarkus.langchain4j.pgvector.table[quarkus.langchain4j.pgvector.table]` diff --git a/pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/Langchain4jPgvectorProcessor.java b/pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/Langchain4jPgvectorProcessor.java index 59f579d23..2e328e43b 100644 --- a/pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/Langchain4jPgvectorProcessor.java +++ b/pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/Langchain4jPgvectorProcessor.java @@ -1,7 +1,9 @@ package io.quarkiverse.langchain4j.pgvector.deployment; import jakarta.enterprise.context.ApplicationScoped; +import jakarta.enterprise.inject.Default; +import org.jboss.jandex.AnnotationInstance; import org.jboss.jandex.ClassType; import org.jboss.jandex.DotName; import org.jboss.jandex.ParameterizedType; @@ -11,9 +13,11 @@ import dev.langchain4j.data.segment.TextSegment; import dev.langchain4j.store.embedding.EmbeddingStore; import io.agroal.api.AgroalDataSource; +import io.quarkiverse.langchain4j.deployment.EmbeddingStoreBuildItem; import io.quarkiverse.langchain4j.pgvector.PgVectorEmbeddingStore; import io.quarkiverse.langchain4j.pgvector.runtime.PgVectorEmbeddingStoreConfig; import io.quarkiverse.langchain4j.pgvector.runtime.PgVectorEmbeddingStoreRecorder; +import io.quarkus.agroal.DataSource; import io.quarkus.arc.deployment.SyntheticBeanBuildItem; import io.quarkus.deployment.annotations.BuildProducer; import io.quarkus.deployment.annotations.BuildStep; @@ -38,7 +42,19 @@ FeatureBuildItem feature() { public void createBean( BuildProducer beanProducer, PgVectorEmbeddingStoreRecorder recorder, - PgVectorEmbeddingStoreConfig config) { + PgVectorEmbeddingStoreConfig config, + PgVectorEmbeddingStoreBuildTimeConfig buildTimeConfig, + BuildProducer embeddingStoreProducer) { + String datasourceName = buildTimeConfig.datasource().orElse(null); + AnnotationInstance datasourceQualifier; + if (datasourceName == null) { + datasourceQualifier = AnnotationInstance.builder(Default.class).build(); + + } else { + datasourceQualifier = AnnotationInstance.builder(DataSource.class) + .add("value", datasourceName) + .build(); + } beanProducer.produce(SyntheticBeanBuildItem .configure(PGVECTOR_EMBEDDING_STORE) .types(ClassType.create(EmbeddingStore.class), @@ -46,10 +62,10 @@ public void createBean( .setRuntimeInit() .defaultBean() .scope(ApplicationScoped.class) - .addInjectionPoint(ClassType.create(DotName.createSimple(AgroalDataSource.class))) - .createWith(recorder.embeddingStoreFunction(config)) + .createWith(recorder.embeddingStoreFunction(config, buildTimeConfig.datasource().orElse(null))) + .addInjectionPoint(ClassType.create(DotName.createSimple(AgroalDataSource.class)), datasourceQualifier) .done()); - + embeddingStoreProducer.produce(new EmbeddingStoreBuildItem()); } @BuildStep diff --git a/pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/PgVectorEmbeddingStoreBuildTimeConfig.java b/pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/PgVectorEmbeddingStoreBuildTimeConfig.java new file mode 100644 index 000000000..7d7185d6c --- /dev/null +++ b/pgvector/deployment/src/main/java/io/quarkiverse/langchain4j/pgvector/deployment/PgVectorEmbeddingStoreBuildTimeConfig.java @@ -0,0 +1,20 @@ +package io.quarkiverse.langchain4j.pgvector.deployment; + +import static io.quarkus.runtime.annotations.ConfigPhase.BUILD_TIME; + +import java.util.Optional; + +import io.quarkus.runtime.annotations.ConfigRoot; +import io.smallrye.config.ConfigMapping; + +@ConfigRoot(phase = BUILD_TIME) +@ConfigMapping(prefix = "quarkus.langchain4j.pgvector") +public interface PgVectorEmbeddingStoreBuildTimeConfig { + + /** + * The name of the configured Postgres datasource to use for this store. If not set, + * the default datasource from the Agroal extension will be used. + */ + Optional datasource(); + +} diff --git a/pgvector/runtime/src/main/java/io/quarkiverse/langchain4j/pgvector/runtime/PgVectorEmbeddingStoreRecorder.java b/pgvector/runtime/src/main/java/io/quarkiverse/langchain4j/pgvector/runtime/PgVectorEmbeddingStoreRecorder.java index 42ec4fada..b4a7ffcb4 100644 --- a/pgvector/runtime/src/main/java/io/quarkiverse/langchain4j/pgvector/runtime/PgVectorEmbeddingStoreRecorder.java +++ b/pgvector/runtime/src/main/java/io/quarkiverse/langchain4j/pgvector/runtime/PgVectorEmbeddingStoreRecorder.java @@ -6,6 +6,7 @@ import io.agroal.api.AgroalDataSource; import io.quarkiverse.langchain4j.pgvector.PgVectorEmbeddingStore; +import io.quarkus.agroal.DataSource.DataSourceLiteral; import io.quarkus.arc.SyntheticCreationalContext; import io.quarkus.runtime.annotations.Recorder; @@ -13,13 +14,16 @@ public class PgVectorEmbeddingStoreRecorder { public Function, PgVectorEmbeddingStore> embeddingStoreFunction( - PgVectorEmbeddingStoreConfig config) { + PgVectorEmbeddingStoreConfig config, String datasourceName) { return new Function<>() { @Override public PgVectorEmbeddingStore apply(SyntheticCreationalContext context) { AgroalDataSource dataSource; - //TODO handle named datasources - dataSource = context.getInjectedReference(AgroalDataSource.class, new Default.Literal()); + if (datasourceName == null) { + dataSource = context.getInjectedReference(AgroalDataSource.class, new Default.Literal()); + } else { + dataSource = context.getInjectedReference(AgroalDataSource.class, new DataSourceLiteral(datasourceName)); + } return new PgVectorEmbeddingStore(dataSource, config.table(), config.dimension(), config.useIndex(), config.indexListSize(), config.createTable(), config.dropTableFirst()); }