Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Interface for VectorStoreService #510

Merged
merged 4 commits into from
Oct 30, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions server/src/main/kotlin/com/xebia/functional/xef/server/Server.kt
Original file line number Diff line number Diff line change
Expand Up @@ -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.*
Expand Down Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
@@ -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
}
Original file line number Diff line number Diff line change
@@ -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()
}
}
Original file line number Diff line number Diff line change
@@ -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"
)
}
}

This file was deleted.

Original file line number Diff line number Diff line change
@@ -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)
}
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ class PostgresVectorStoreService(
private val dataSource: HikariDataSource
) : VectorStoreService() {

override fun addCollection() {
fun addCollection() {
dataSource.connection {
// Create collection
val uuid = UUID.generateUUID()
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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,
)
6 changes: 6 additions & 0 deletions server/src/main/resources/database.conf
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down