Interface for VectorStoreService (#510)

* Interface for VectorStoreService

* Build fixed

* Loading vector-store by demand

* Removed shadow var

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) {

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
+}

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()
+  }
+}

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"
+      )
+  }
+}

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)
+}

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

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

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,
+)

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 {