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 c2a49d7bd..1d809c42c 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 @@ -6,11 +6,11 @@ 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.db.psql.XefVectorStoreConfig -import com.xebia.functional.xef.server.db.psql.XefVectorStoreConfig.Companion.getVectorStoreService 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.* @@ -46,9 +46,11 @@ object Server { xefDBConfig.password ) Database.connect(hikariDataSourceXefDB) - val vectorStoreConfig = XefVectorStoreConfig.load("xef-vector-store", config) - val vectorStoreService = vectorStoreConfig.getVectorStoreService(config, logger) - vectorStoreService.addCollection() + + val vectorStoreService = + VectorStoreService.load("xef-vector-store", config).getVectorStoreService(logger) + + (vectorStoreService as? PostgresVectorStoreService)?.addCollection() 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 new file mode 100644 index 000000000..f370ee2e1 --- /dev/null +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/db/VectorStoreConfig.kt @@ -0,0 +1,8 @@ +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 new file mode 100644 index 000000000..7540183a5 --- /dev/null +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/db/local/LocalVectorStoreConfig.kt @@ -0,0 +1,16 @@ +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/PSQLVectorStoreConfig.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt new file mode 100644 index 000000000..6147d4a60 --- /dev/null +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt @@ -0,0 +1,73 @@ +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 host: String, + val port: Int, + val database: String, + val driver: String, + val user: String, + val password: String, + val collectionName: String, + val vectorSize: Int +) : VectorStoreConfig { + + fun getUrl(): String = "jdbc:postgresql://$host:$port/$database" + + override suspend fun getVectorStoreService(logger: Logger): PostgresVectorStoreService { + val vectorStoreHikariDataSource = + RepositoryService.getHikariDataSource(getUrl(), user, password) + return PostgresVectorStoreService(toPGVectorStoreConfig(), logger, vectorStoreHikariDataSource) + } + + private fun toPGVectorStoreConfig() = + PostgreSQLXef.PGVectorStoreConfig( + dbConfig = + PostgreSQLXef.DBConfig( + host = host, + port = port, + database = database, + user = user, + password = password + ), + collectionName = collectionName, + vectorSize = vectorSize + ) + + companion object { + @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( + host = this.host, + port = this.port, + database = this.database, + 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/XefVectorStoreConfig.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/XefVectorStoreConfig.kt deleted file mode 100644 index e6bf2b1ae..000000000 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/XefVectorStoreConfig.kt +++ /dev/null @@ -1,99 +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.services.PostgreSQLXef -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 com.xebia.functional.xef.store.migrations.PsqlVectorStoreConfig -import com.zaxxer.hikari.HikariDataSource -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 -} - -@Serializable -class XefVectorStoreConfig( - val type: XefVectorStoreType, - val host: String, - val port: Int, - val database: String, - val driver: String, - val user: String, - val password: String -) { - - fun getUrl(): String = "jdbc:postgresql://$host:$port/$database" - - companion object { - @OptIn(ExperimentalSerializationApi::class) - suspend fun load(configNamespace: String, config: Config? = null): XefVectorStoreConfig = - withContext(Dispatchers.IO) { - val rawConfig = config ?: ConfigFactory.load().resolve() - val jdbcConfig = rawConfig.getConfig(configNamespace) - val config = Hocon.decodeFromConfig(serializer(), jdbcConfig) - config.migrate() - config - } - - suspend fun XefVectorStoreConfig.getVectorStoreService( - config: Config, - logger: Logger - ): VectorStoreService { - when (this.type) { - XefVectorStoreType.PSQL -> { - val vectorStoreHikariDataSource = - RepositoryService.getHikariDataSource(getUrl(), user, password) - return getPsqlVectorStoreService(config, vectorStoreHikariDataSource, logger) - } - } - } - - private suspend fun getPsqlVectorStoreService( - config: Config, - dataSource: HikariDataSource, - logger: Logger - ): VectorStoreService { - val vectorStoreConfig = XefVectorStoreConfig.load("xef-vector-store", config) - val pgVectorStoreConfig = - PostgreSQLXef.PGVectorStoreConfig( - dbConfig = - PostgreSQLXef.DBConfig( - host = vectorStoreConfig.host, - port = vectorStoreConfig.port, - database = vectorStoreConfig.database, - user = vectorStoreConfig.user, - password = vectorStoreConfig.password - ) - ) - return PostgresVectorStoreService(pgVectorStoreConfig, logger, dataSource) - } - } - - private suspend fun migrate() { - when (this.type) { - XefVectorStoreType.PSQL -> { - val psqlConfig = this.toPSQLConfig() - psqlConfig.migrate() - } - } - } - - private fun XefVectorStoreConfig.toPSQLConfig(): PsqlVectorStoreConfig = - PsqlVectorStoreConfig( - host = this.host, - port = this.port, - database = this.database, - driver = this.driver, - user = this.user, - password = this.password, - migrationsTable = "migration" - ) -} diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/services/LocalVectorStoreService.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/services/LocalVectorStoreService.kt new file mode 100644 index 000000000..06364727b --- /dev/null +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/services/LocalVectorStoreService.kt @@ -0,0 +1,11 @@ +package com.xebia.functional.xef.server.services + +import com.xebia.functional.xef.conversation.llm.openai.OpenAI +import com.xebia.functional.xef.server.http.routes.Provider +import com.xebia.functional.xef.store.LocalVectorStore +import com.xebia.functional.xef.store.VectorStore + +class LocalVectorStoreService : VectorStoreService() { + override fun getVectorStore(provider: Provider, token: String?): VectorStore = + LocalVectorStore(OpenAI().DEFAULT_EMBEDDING) +} 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 4b6799114..665464834 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 @@ -37,7 +37,7 @@ class PostgresVectorStoreService( private val dataSource: HikariDataSource ) : VectorStoreService() { - override fun addCollection() { + fun addCollection() { dataSource.connection { // Create collection val uuid = UUID.generateUUID() @@ -49,7 +49,7 @@ class PostgresVectorStoreService( } } - override fun getVectorStore(provider: Provider, token: String): VectorStore { + override fun getVectorStore(provider: Provider, token: String?): VectorStore { val embeddings = when (provider) { Provider.OPENAI -> OpenAI(token).DEFAULT_EMBEDDING diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/services/TokenRepositoryService.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/services/TokenRepositoryService.kt index 7a4369857..d3c8d23d1 100644 --- a/server/src/main/kotlin/com/xebia/functional/xef/server/services/TokenRepositoryService.kt +++ b/server/src/main/kotlin/com/xebia/functional/xef/server/services/TokenRepositoryService.kt @@ -1,6 +1,10 @@ package com.xebia.functional.xef.server.services import com.xebia.functional.xef.server.db.tables.* +import com.xebia.functional.xef.server.db.tables.Project +import com.xebia.functional.xef.server.db.tables.ProjectsTable +import com.xebia.functional.xef.server.db.tables.XefTokens +import com.xebia.functional.xef.server.db.tables.XefTokensTable import com.xebia.functional.xef.server.models.* import com.xebia.functional.xef.server.models.exceptions.XefExceptions.* import java.util.UUID 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 7a0a30ec0..fad7548d8 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,11 +1,45 @@ package com.xebia.functional.xef.server.services +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 kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.withContext +import kotlinx.serialization.ExperimentalSerializationApi +import kotlinx.serialization.Serializable +import kotlinx.serialization.hocon.Hocon -abstract class VectorStoreService { +enum class XefVectorStoreType { + PSQL, + LOCAL +} - abstract fun addCollection(): Unit +abstract class VectorStoreService { + abstract fun getVectorStore( + provider: Provider = Provider.OPENAI, + token: String? = null + ): VectorStore - abstract fun getVectorStore(provider: Provider = Provider.OPENAI, token: String): 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 f3667c898..449f1fa0d 100644 --- a/server/src/main/resources/database.conf +++ b/server/src/main/resources/database.conf @@ -19,6 +19,12 @@ xef-vector-store { password = "postgres" password = ${?XEF_DB_VECTOR_STORE_PASSWORD} + + collectionName = "xef_collection" + collectionName = ${?XEF_DB_COLLECTION_NAME} + + vectorSize = 1536 + vectorSize = ${?XEF_DB_VECTOR_SIZE} } xef {