From 812b744eb103a3ea863b1782946b2219f7b56148 Mon Sep 17 00:00:00 2001 From: Francisco Diaz Date: Wed, 15 Nov 2023 12:07:53 +0100 Subject: [PATCH] Reorganize PostgreSQL code (#541) * Reorganize Postgres configuration * Create Hikari data source as resource * Remove duplicated code * Remove PostgresVectorStoreConfig class * Remove unnecessary println --- .../config/PostgreSQLVectorStoreConfig.kt | 15 ++++ .../store/migrations/DatabaseMigrations.kt | 49 +++++++++++++ .../store/migrations/PsqlVectorStoreConfig.kt | 51 ------------- .../com/xebia/functional/xef/server/Server.kt | 63 ++++++++-------- .../xef/server/db/VectorStoreConfig.kt | 8 -- .../server/db/local/LocalVectorStoreConfig.kt | 16 ---- .../functional/xef/server/db/psql/Migrate.kt | 27 ------- .../server/db/psql/PSQLVectorStoreConfig.kt | 73 ------------------- .../xef/server/db/psql/XefDatabaseConfig.kt | 17 ++--- .../services/PostgresVectorStoreService.kt | 33 +++------ .../xef/server/services/RepositoryService.kt | 27 ++++--- .../xef/server/services/VectorStoreService.kt | 73 ++++++++++++------- server/src/main/resources/database.conf | 27 ++----- 13 files changed, 183 insertions(+), 296 deletions(-) create mode 100644 integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/config/PostgreSQLVectorStoreConfig.kt create mode 100644 integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/DatabaseMigrations.kt delete mode 100644 integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt delete mode 100644 server/src/main/kotlin/com/xebia/functional/xef/server/db/VectorStoreConfig.kt delete mode 100644 server/src/main/kotlin/com/xebia/functional/xef/server/db/local/LocalVectorStoreConfig.kt delete mode 100644 server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/Migrate.kt delete mode 100644 server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt diff --git a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/config/PostgreSQLVectorStoreConfig.kt b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/config/PostgreSQLVectorStoreConfig.kt new file mode 100644 index 000000000..6ade0d8ee --- /dev/null +++ b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/config/PostgreSQLVectorStoreConfig.kt @@ -0,0 +1,15 @@ +package com.xebia.functional.xef.store.config + +import kotlinx.serialization.Serializable + +@Serializable +data class PostgreSQLVectorStoreConfig( + val url: String, + val driver: String, + val user: String, + val password: String, + val collectionName: String, + val vectorSize: Int, + val migrationsTable: String = "migrations", + val migrationsLocations: List = listOf("vectorStore/migrations") +) diff --git a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/DatabaseMigrations.kt b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/DatabaseMigrations.kt new file mode 100644 index 000000000..c420a69d9 --- /dev/null +++ b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/DatabaseMigrations.kt @@ -0,0 +1,49 @@ +package com.xebia.functional.xef.store.migrations + +import com.xebia.functional.xef.store.config.PostgreSQLVectorStoreConfig +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.withContext +import org.flywaydb.core.Flyway +import org.flywaydb.core.api.configuration.FluentConfiguration +import org.flywaydb.core.api.output.MigrateResult +import javax.sql.DataSource + +suspend fun runDatabaseMigrations( + dataSource: DataSource, + migrationsTable: String, + migrationsLocations: List +): MigrateResult = + withContext(Dispatchers.IO) { + val migration: FluentConfiguration = Flyway.configure() + .dataSource(dataSource) + .table(migrationsTable) + .locations(*migrationsLocations.toTypedArray()) + .loggers("slf4j") + val isValid = migration.ignoreMigrationPatterns("*:pending").load().validateWithResult() + if (!isValid.validationSuccessful) { + throw IllegalStateException("Migration validation failed: ${isValid.errorDetails}") + } + migration.load().migrate() + } + +suspend fun runDatabaseMigrations( + config: PostgreSQLVectorStoreConfig +): MigrateResult = + withContext(Dispatchers.IO) { + with(config) { + val migration: FluentConfiguration = Flyway.configure() + .dataSource( + url, + user, + password + ) + .table(migrationsTable) + .locations(*migrationsLocations.toTypedArray()) + .loggers("slf4j") + val isValid = migration.ignoreMigrationPatterns("*:pending").load().validateWithResult() + if (!isValid.validationSuccessful) { + throw IllegalStateException("Migration validation failed: ${isValid.errorDetails}") + } + migration.load().migrate() + } + } diff --git a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt deleted file mode 100644 index 0f93526f9..000000000 --- a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt +++ /dev/null @@ -1,51 +0,0 @@ -package com.xebia.functional.xef.store.migrations - -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.withContext -import org.flywaydb.core.Flyway -import org.flywaydb.core.api.configuration.FluentConfiguration -import org.flywaydb.core.api.output.MigrateResult - -class PsqlVectorStoreConfig( - val url: String, - val driver: String, - val user: String, - val password: String, - val migrationsTable: String, - val migrationsLocations: List = listOf("vectorStore/migrations") -) { - suspend fun migrate(): MigrateResult = - withContext(Dispatchers.IO) { - val migration: FluentConfiguration = Flyway.configure() - .dataSource( - url, - user, - password - ) - .table(migrationsTable) - .locations(*migrationsLocations.toTypedArray()) - .loggers("slf4j") - val isValid = migration.ignoreMigrationPatterns("*:pending").load().validateWithResult() - if (!isValid.validationSuccessful) { - throw IllegalStateException("Migration validation failed: ${isValid.errorDetails}") - } - migration.load().migrate() - } - - companion object { - operator fun invoke( - host: String, - port: Int, - database: String, - driver: String, - user: String, - password: String, - migrationsTable: String, - migrationsLocations: List = listOf("vectorStore/migrations") - ): PsqlVectorStoreConfig { - val uri = "jdbc:postgresql://${host}:${port}/${database}" - return PsqlVectorStoreConfig(uri, driver, user, password, migrationsTable, migrationsLocations) - } - } - -} diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/Server.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/Server.kt index 1d809c42c..1d301a9cd 100644 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/Server.kt +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/Server.kt @@ -4,27 +4,30 @@ import arrow.continuations.SuspendApp import arrow.continuations.ktor.server import arrow.fx.coroutines.resourceScope import com.typesafe.config.ConfigFactory -import com.xebia.functional.xef.server.db.psql.Migrate import com.xebia.functional.xef.server.db.psql.XefDatabaseConfig import com.xebia.functional.xef.server.exceptions.exceptionsHandler -import com.xebia.functional.xef.server.http.routes.* -import com.xebia.functional.xef.server.services.PostgresVectorStoreService -import com.xebia.functional.xef.server.services.RepositoryService -import com.xebia.functional.xef.server.services.VectorStoreService -import io.ktor.client.* -import io.ktor.client.engine.cio.* -import io.ktor.client.plugins.auth.* +import com.xebia.functional.xef.server.http.routes.aiRoutes +import com.xebia.functional.xef.server.http.routes.xefRoutes +import com.xebia.functional.xef.server.services.hikariDataSource +import com.xebia.functional.xef.server.services.vectorStoreService +import com.xebia.functional.xef.store.migrations.runDatabaseMigrations +import io.ktor.client.HttpClient +import io.ktor.client.engine.cio.CIO +import io.ktor.client.plugins.auth.Auth import io.ktor.client.plugins.contentnegotiation.ContentNegotiation as ClientContentNegotiation -import io.ktor.client.plugins.logging.* -import io.ktor.http.* -import io.ktor.serialization.kotlinx.json.* -import io.ktor.server.application.* -import io.ktor.server.auth.* -import io.ktor.server.netty.* -import io.ktor.server.plugins.contentnegotiation.* -import io.ktor.server.plugins.cors.routing.* -import io.ktor.server.resources.* -import io.ktor.server.routing.* +import io.ktor.client.plugins.logging.LogLevel +import io.ktor.client.plugins.logging.Logging +import io.ktor.http.HttpMethod +import io.ktor.serialization.kotlinx.json.json +import io.ktor.server.application.install +import io.ktor.server.auth.Authentication +import io.ktor.server.auth.UserIdPrincipal +import io.ktor.server.auth.bearer +import io.ktor.server.netty.Netty +import io.ktor.server.plugins.contentnegotiation.ContentNegotiation +import io.ktor.server.plugins.cors.routing.CORS +import io.ktor.server.resources.Resources +import io.ktor.server.routing.routing import kotlinx.coroutines.awaitCancellation import org.jetbrains.exposed.sql.Database import org.slf4j.LoggerFactory @@ -33,24 +36,22 @@ object Server { @JvmStatic fun main(args: Array) = SuspendApp { resourceScope { + val logger = LoggerFactory.getLogger("xef-server") + val config = ConfigFactory.load("database.conf").resolve() - val xefDBConfig = XefDatabaseConfig.load("xef", config) - Migrate.migrate(xefDBConfig) + val xefDBConfig = XefDatabaseConfig.load("xef-database", config) - val logger = LoggerFactory.getLogger("xef-server") + val xefDatasource = hikariDataSource(xefDBConfig.url, xefDBConfig.user, xefDBConfig.password) - val hikariDataSourceXefDB = - RepositoryService.getHikariDataSource( - xefDBConfig.getUrl(), - xefDBConfig.user, - xefDBConfig.password - ) - Database.connect(hikariDataSourceXefDB) + runDatabaseMigrations( + xefDatasource, + xefDBConfig.migrationsTable, + xefDBConfig.migrationsLocations + ) - val vectorStoreService = - VectorStoreService.load("xef-vector-store", config).getVectorStoreService(logger) + Database.connect(xefDatasource) - (vectorStoreService as? PostgresVectorStoreService)?.addCollection() + vectorStoreService("xef-vector-store", config, logger) val ktorClient = HttpClient(CIO) { diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/db/VectorStoreConfig.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/VectorStoreConfig.kt deleted file mode 100644 index f370ee2e1..000000000 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/db/VectorStoreConfig.kt +++ /dev/null @@ -1,8 +0,0 @@ -package com.xebia.functional.xef.server.db - -import com.xebia.functional.xef.server.services.VectorStoreService -import org.slf4j.Logger - -interface VectorStoreConfig { - suspend fun getVectorStoreService(logger: Logger): VectorStoreService -} diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/db/local/LocalVectorStoreConfig.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/local/LocalVectorStoreConfig.kt deleted file mode 100644 index 7540183a5..000000000 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/db/local/LocalVectorStoreConfig.kt +++ /dev/null @@ -1,16 +0,0 @@ -package com.xebia.functional.xef.server.db.local - -import com.xebia.functional.xef.server.db.VectorStoreConfig -import com.xebia.functional.xef.server.services.LocalVectorStoreService -import kotlinx.serialization.Serializable -import org.slf4j.Logger - -@Serializable -class LocalVectorStoreConfig() : VectorStoreConfig { - override suspend fun getVectorStoreService(logger: Logger): LocalVectorStoreService = - LocalVectorStoreService() - - companion object { - fun load(): LocalVectorStoreConfig = LocalVectorStoreConfig() - } -} diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/Migrate.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/Migrate.kt deleted file mode 100644 index 728adf227..000000000 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/Migrate.kt +++ /dev/null @@ -1,27 +0,0 @@ -package com.xebia.functional.xef.server.db.psql - -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.withContext -import org.flywaydb.core.Flyway -import org.flywaydb.core.api.configuration.FluentConfiguration -import org.flywaydb.core.api.output.MigrateResult - -object Migrate { - suspend fun migrate( - config: XefDatabaseConfig, - ): MigrateResult = - withContext(Dispatchers.IO) { - val url = config.getUrl() - val migration: FluentConfiguration = - Flyway.configure() - .dataSource(url, config.user, config.password) - .table(config.migrationsTable) - .locations(*config.migrationsLocations.toTypedArray()) - .loggers("slf4j") - val isValid = migration.ignoreMigrationPatterns("*:pending").load().validateWithResult() - if (!isValid.validationSuccessful) { - throw IllegalStateException("Migration validation failed: ${isValid.errorDetails}") - } - migration.load().migrate() - } -} diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt deleted file mode 100644 index fd216bb5f..000000000 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt +++ /dev/null @@ -1,73 +0,0 @@ -package com.xebia.functional.xef.server.db.psql - -import com.typesafe.config.Config -import com.typesafe.config.ConfigFactory -import com.xebia.functional.xef.server.db.VectorStoreConfig -import com.xebia.functional.xef.server.services.PostgreSQLXef -import com.xebia.functional.xef.server.services.PostgresVectorStoreService -import com.xebia.functional.xef.server.services.RepositoryService -import com.xebia.functional.xef.store.migrations.PsqlVectorStoreConfig -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.withContext -import kotlinx.serialization.ExperimentalSerializationApi -import kotlinx.serialization.Serializable -import kotlinx.serialization.hocon.Hocon -import org.slf4j.Logger - -@Serializable -class PSQLVectorStoreConfig( - val url: String, - val driver: String, - val user: String, - val password: String, - val collectionName: String, - val vectorSize: Int -) : VectorStoreConfig { - - override suspend fun getVectorStoreService(logger: Logger): PostgresVectorStoreService { - val vectorStoreHikariDataSource = RepositoryService.getHikariDataSource(url, user, password) - return PostgresVectorStoreService(toPGVectorStoreConfig(), logger, vectorStoreHikariDataSource) - } - - private fun toPGVectorStoreConfig() = - PostgreSQLXef.PGVectorStoreConfig( - dbConfig = PostgreSQLXef.DBConfig(url = url, user = user, password = password), - collectionName = collectionName, - vectorSize = vectorSize - ) - - companion object { - operator fun invoke( - host: String, - port: Int, - database: String, - driver: String, - user: String, - password: String, - collectionName: String, - vectorSize: Int - ): PSQLVectorStoreConfig { - val url = "jdbc:postgresql://${host}:${port}/${database}" - return PSQLVectorStoreConfig(url, driver, user, password, collectionName, vectorSize) - } - - @OptIn(ExperimentalSerializationApi::class) - suspend fun load(configNamespace: String, config: Config?): PSQLVectorStoreConfig = - withContext(Dispatchers.IO) { - val rawConfig = config ?: ConfigFactory.load().resolve() - val jdbcConfig = rawConfig.getConfig(configNamespace) - val psqlConfig = Hocon.decodeFromConfig(serializer(), jdbcConfig) - psqlConfig.toPSQLConfig().migrate() - psqlConfig - } - - private fun PSQLVectorStoreConfig.toPSQLConfig(): PsqlVectorStoreConfig = - PsqlVectorStoreConfig( - url = this.url, - driver = this.driver, - user = this.user, - password = this.password, - migrationsTable = "migration" - ) - } -} diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/XefDatabaseConfig.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/XefDatabaseConfig.kt index c86dd19eb..ed6bc9564 100644 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/XefDatabaseConfig.kt +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/XefDatabaseConfig.kt @@ -1,7 +1,6 @@ package com.xebia.functional.xef.server.db.psql import com.typesafe.config.Config -import com.typesafe.config.ConfigFactory import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.withContext import kotlinx.serialization.ExperimentalSerializationApi @@ -9,25 +8,19 @@ import kotlinx.serialization.Serializable import kotlinx.serialization.hocon.Hocon @Serializable -class XefDatabaseConfig( - val host: String, - val port: Int, - val database: String, +data class XefDatabaseConfig( + val url: String, val user: String, val password: String, val migrationsTable: String, val migrationsLocations: List ) { - - fun getUrl(): String = "jdbc:postgresql://$host:$port/$database" - companion object { @OptIn(ExperimentalSerializationApi::class) - suspend fun load(configNamespace: String, config: Config? = null): XefDatabaseConfig = + suspend fun load(configNamespace: String, config: Config): XefDatabaseConfig = withContext(Dispatchers.IO) { - val rawConfig = config ?: ConfigFactory.load().resolve() - val jdbcConfig = rawConfig.getConfig(configNamespace) - Hocon.decodeFromConfig(serializer(), jdbcConfig) + val databaseConfig = config.getConfig(configNamespace) + Hocon.decodeFromConfig(serializer(), databaseConfig) } } } diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt index b661af17f..65441d554 100644 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt @@ -8,27 +8,18 @@ import com.xebia.functional.xef.store.VectorStore import com.xebia.functional.xef.store.postgresql.PGDistanceStrategy import com.xebia.functional.xef.store.postgresql.addNewCollection import com.xebia.functional.xef.store.postgresql.connection -import com.zaxxer.hikari.HikariDataSource +import javax.sql.DataSource import kotlinx.uuid.UUID import kotlinx.uuid.generateUUID import org.slf4j.Logger -object PostgreSQLXef { - data class DBConfig(val url: String, val user: String, val password: String) - - data class PGVectorStoreConfig( - val dbConfig: DBConfig, - val vectorSize: Int = 1536, // OpenAI default - val collectionName: String = "xef_collection", - val preDeleteCollection: Boolean = false, - val chunkSize: Int? = null, - ) -} - class PostgresVectorStoreService( - private val config: PostgreSQLXef.PGVectorStoreConfig, private val logger: Logger, - private val dataSource: HikariDataSource + private val dataSource: DataSource, + private val collectionName: String, + private val vectorSize: Int, + private val preDeleteCollection: Boolean = false, + private val chunkSize: Int? = null, ) : VectorStoreService() { fun addCollection() { @@ -37,9 +28,9 @@ class PostgresVectorStoreService( val uuid = UUID.generateUUID() update(addNewCollection) { bind(uuid.toString()) - bind(config.collectionName) + bind(collectionName) } - .also { logger.info("Created collection ${config.collectionName}") } + .also { logger.info("Created collection $collectionName") } } } @@ -53,14 +44,14 @@ class PostgresVectorStoreService( } return PGVectorStore( - vectorSize = config.vectorSize, + vectorSize = vectorSize, dataSource = dataSource, embeddings = embeddings, - collectionName = config.collectionName, + collectionName = collectionName, distanceStrategy = PGDistanceStrategy.Euclidean, - preDeleteCollection = config.preDeleteCollection, + preDeleteCollection = preDeleteCollection, requestConfig = RequestConfig(user = RequestConfig.Companion.User("user")), - chunkSize = config.chunkSize + chunkSize = chunkSize ) } } diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/services/RepositoryService.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/services/RepositoryService.kt index 032506fcd..1b7430ed7 100644 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/services/RepositoryService.kt +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/services/RepositoryService.kt @@ -1,25 +1,28 @@ package com.xebia.functional.xef.server.services +import arrow.fx.coroutines.ResourceScope +import arrow.fx.coroutines.closeable import com.xebia.functional.xef.server.db.tables.User import com.xebia.functional.xef.server.db.tables.UsersTable import com.xebia.functional.xef.server.models.Token import com.xebia.functional.xef.server.models.exceptions.XefExceptions import com.zaxxer.hikari.HikariConfig import com.zaxxer.hikari.HikariDataSource -import org.jetbrains.exposed.sql.SqlExpressionBuilder.eq -object RepositoryService { - fun getHikariDataSource(url: String, usr: String, passwd: String): HikariDataSource { - val hikariConfig = - HikariConfig().apply { - jdbcUrl = url - username = usr - password = passwd - driverClassName = "org.postgresql.Driver" - } +suspend fun ResourceScope.hikariDataSource( + url: String, + usr: String, + passwd: String +): HikariDataSource { + val hikariConfig = + HikariConfig().apply { + jdbcUrl = url + username = usr + password = passwd + driverClassName = "org.postgresql.Driver" + } - return HikariDataSource(hikariConfig) - } + return closeable { HikariDataSource(hikariConfig) } } fun Token.getUser() = diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/services/VectorStoreService.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/services/VectorStoreService.kt index fad7548d8..5c8eafab6 100644 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/services/VectorStoreService.kt +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/services/VectorStoreService.kt @@ -1,45 +1,68 @@ package com.xebia.functional.xef.server.services +import arrow.fx.coroutines.ResourceScope import com.typesafe.config.Config -import com.typesafe.config.ConfigFactory -import com.xebia.functional.xef.server.db.VectorStoreConfig -import com.xebia.functional.xef.server.db.local.LocalVectorStoreConfig -import com.xebia.functional.xef.server.db.psql.PSQLVectorStoreConfig import com.xebia.functional.xef.server.http.routes.Provider import com.xebia.functional.xef.store.VectorStore +import com.xebia.functional.xef.store.config.PostgreSQLVectorStoreConfig +import com.xebia.functional.xef.store.migrations.runDatabaseMigrations import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.withContext import kotlinx.serialization.ExperimentalSerializationApi -import kotlinx.serialization.Serializable import kotlinx.serialization.hocon.Hocon +import org.slf4j.Logger enum class XefVectorStoreType { PSQL, - LOCAL + LOCAL; + + companion object { + fun loadFromConfiguration(config: Config): XefVectorStoreType { + val name = config.getString("type") + return entries.firstOrNull { it.name == name } ?: PSQL + } + } } +@OptIn(ExperimentalSerializationApi::class) +suspend fun ResourceScope.vectorStoreService( + configNamespace: String, + config: Config, + logger: Logger +): VectorStoreService = + withContext(Dispatchers.IO) { + val vectorStoreConfig = config.getConfig(configNamespace) + + when (XefVectorStoreType.loadFromConfiguration(vectorStoreConfig)) { + XefVectorStoreType.LOCAL -> LocalVectorStoreService() + XefVectorStoreType.PSQL -> { + val postgresVectorStoreConfig = + Hocon.decodeFromConfig(PostgreSQLVectorStoreConfig.serializer(), vectorStoreConfig) + val dataSource = + hikariDataSource( + postgresVectorStoreConfig.url, + postgresVectorStoreConfig.user, + postgresVectorStoreConfig.password + ) + runDatabaseMigrations( + dataSource, + postgresVectorStoreConfig.migrationsTable, + postgresVectorStoreConfig.migrationsLocations + ) + PostgresVectorStoreService( + logger, + dataSource, + postgresVectorStoreConfig.collectionName, + postgresVectorStoreConfig.vectorSize + ) + .also { it.addCollection() } + } + } + } + abstract class VectorStoreService { abstract fun getVectorStore( provider: Provider = Provider.OPENAI, token: String? = null ): VectorStore - - companion object { - @OptIn(ExperimentalSerializationApi::class) - suspend fun load(configNamespace: String, config: Config?): VectorStoreConfig = - withContext(Dispatchers.IO) { - val rawConfig = config ?: ConfigFactory.load().resolve() - val jdbcConfig = rawConfig.getConfig(configNamespace) - val typeConfig = Hocon.decodeFromConfig(VectorStoreTypeConfig.serializer(), jdbcConfig) - when (typeConfig.type) { - XefVectorStoreType.PSQL -> PSQLVectorStoreConfig.load(configNamespace, rawConfig) - XefVectorStoreType.LOCAL -> LocalVectorStoreConfig.load() - } - } - } } - -@Serializable -class VectorStoreTypeConfig( - val type: XefVectorStoreType, -) diff --git a/server/src/main/resources/database.conf b/server/src/main/resources/database.conf index 449f1fa0d..ed0bb69b6 100644 --- a/server/src/main/resources/database.conf +++ b/server/src/main/resources/database.conf @@ -5,14 +5,8 @@ xef-vector-store { driver = "org.postgresql.Driver" - host = "localhost" - host = ${?XEF_DB_VECTOR_STORE_HOST} - - port = 5432 - port = ${?XEF_DB_VECTOR_STORE_PORT} - - database = "xef-vector-store" - database = ${?XEF_DB_VECTOR_STORE_NAME} + url = "jdbc:postgresql://localhost:5432/xef_vector_store" + url = ${?XEF_DB_VECTOR_STORE_URL} user = "postgres" user = ${?XEF_DB_VECTOR_STORE_USER} @@ -27,15 +21,9 @@ xef-vector-store { vectorSize = ${?XEF_DB_VECTOR_SIZE} } -xef { - host = "localhost" - host = ${?XEF_DB_HOST} - - port = 5433 - port = ${?XEF_DB_PORT} - - database = "xefdb" - database = ${?XEF_DB_NAME} +xef-database { + url = "jdbc:postgresql://localhost:5433/xef_database" + url = ${?XEF_DB_URL} user = "postgres" user = ${?XEF_DB_USER} @@ -44,7 +32,6 @@ xef { password = ${?XEF_DB_PASSWORD} migrationsTable = "migrations" - migrationsLocations = [ - "classpath:db/migrations/psql" - ] + + migrationsLocations = [ "classpath:db/migrations/psql" ] }