ChromaDB.kt

@file:Suppress("UNCHECKED_CAST", "NOTHING_TO_INLINE")

package kore.chromaDB

import kore.vo.VO
import kore.vo.field.value.string
import java.util.concurrent.ConcurrentHashMap

class ChromaDB:VO(){
    //    companion object{
//        fun option(ids:LFS?, keyword:String?, limit:Long?):HashMap<String, Any>
//        = hashMapOf<String, Any>().also{m->
//            ids?.let{m["ids"] = it}
//            keyword?.let{m["keyword"] = it}
//            limit?.let{m["limit"] = it}
//        }
//    }
    var url by string

    @PublishedApi internal val tenants = ConcurrentHashMap<String, String>()
    @PublishedApi internal val dbs = ConcurrentHashMap<String, String>()
    @PublishedApi internal val tables = ConcurrentHashMap<String, String>()

    inline fun unsafe():Unsafe = Unsafe(this)
    suspend inline fun tenant(name:String):TENANT{
        val u = unsafe()
        if(!u.hasTenant(name)){
            u.addTenant(name) ?: throw Exception("Failed to add tenant $name")
        }
        return TENANT(Struct(this, name))
    }
    suspend inline fun <T> tenant(name:String, block:TENANT.()->T):T = tenant(name).let(block)
}

COLLECTION.kt

@file:Suppress("UNCHECKED_CAST", "NOTHING_TO_INLINE")

package kore.chromaDB

import kore.coroutine.awaitAll
import kore.net.RestAPI
import kore.net.getWebClient
import kore.savable.JSON
import kore.util.arrayListFillOf
import kore.util.uuid
import kore.vectorDB.VectorDB
import kore.vo.VO
import kore.vo.field.LF
import kore.vo.field.LFD
import kore.vo.field.LFS
import kore.vo.field.MF
import kore.vo.field.MFA
import kore.vo.field.MFI
import kore.vo.field.MFS
import kore.vo.field.list.doubleListList
import kore.vo.field.list.stringListList
import kore.vo.field.map.stringMap
import kore.vo.field.value.double
import kore.vo.field.value.long
import kore.vo.field.value.string
import kore.vo.save.fromJSON
import kore.vo.save.toJSON
import kotlinx.coroutines.Dispatchers
import org.springframework.http.MediaType
import org.springframework.web.reactive.function.client.awaitBody

@JvmInline value class COLLECTION @PublishedApi internal constructor(val v:Struct){
    class DocsList{
        val list = arrayListOf<Any>()
        fun and(block:(DocsList)->Unit){list.add(hashMapOf($$"$and" to DocsList().also{block(it)}.list))}
        fun or(block:(DocsList)->Unit){list.add(hashMapOf($$"$or" to DocsList().also{block(it)}.list))}
        fun contains(keyword:String){list.add(hashMapOf($$"$contains" to keyword))}
        fun notContains(keyword:String){list.add(hashMapOf($$"$not_contains" to keyword))}
    }
    class DocsMap{
        val map = hashMapOf<String, Any>()
        fun and(block:(DocsList)->Unit){map[$$"$and"] = DocsList().also{block(it)}.list}
        fun or(block:(DocsList)->Unit){map[$$"$or"] = DocsList().also{block(it)}.list}
        fun contains(keyword:String){map[$$"$contains"] = keyword}
        fun notContains(keyword:String){map[$$"$not_contains"] = keyword}
    }
    @JvmInline value class AddDSL @PublishedApi internal constructor(val v:MFA){
        inline fun add(embedded:LFD, documents:String, id:String = uuid(), meta:MFS? = null, uris:String = ""){
            v["embeddings"] as LF<LFD> += embedded
            v["documents"] as LFS += documents
            v["uris"] as LFS += uris
            v["ids"] as LFS += id
            v["metadatas"] as LF<MFS?> += meta
        }
    }
    @JvmInline value class DelDSL @PublishedApi internal constructor(val v:MFA){
        inline fun addId(id:String){v.getOrPut("ids"){arrayListOf<String>()} as LFS += id}
        inline fun addMeta(k:String, value:Any){(v.getOrPut("where"){hashMapOf<String, Any>()} as MFA)[k] = value}
        inline fun document(block:(DocsMap)->Unit){v["where_document"] = DocsMap().also{block(it)}.map}
    }
    @JvmInline value class GetDSL @PublishedApi internal constructor(val v:MFA){
        inline fun addId(id:String){v.getOrPut("ids"){arrayListOf<String>()} as LFS += id}
        inline fun addMeta(k:String, value:Any){(v.getOrPut("where"){hashMapOf<String, Any>()} as MFA)[k] = value}
        inline fun document(block:(DocsMap)->Unit){v["where_document"] = DocsMap().also{block(it)}.map}
        inline fun limit(limit:Long){v["limit"] = limit}
        inline fun offset(offset:Long){v["offset"] = offset}
        @PublishedApi internal fun removeInclude(k:String) {
            (v.getOrPut("include") { arrayListOf("ids", "documents", "metadatas", "embeddings", "distances") } as LFS).remove(k)
        }
        inline fun excludeDocument(){removeInclude("documents")}
        inline fun excludeEmbedding(){removeInclude("embeddings")}
        inline fun excludeMeta(){removeInclude("metadatas")}
        inline fun excludeId(){removeInclude("ids")}
        inline fun excludeDistances(){removeInclude("distances")}
    }
    class QueryDSL @PublishedApi internal constructor(){
        val v:MFA = hashMapOf()
        val keys: HashMap<Int, String> = hashMapOf()
        inline fun addId(id:String){v.getOrPut("ids"){arrayListOf<String>()} as LFS += id}
        inline fun addMeta(k:String, value:Any){(v.getOrPut("where"){hashMapOf<String, Any>()} as MFA)[k] = value}
        inline fun document(block:(DocsMap)->Unit){v["where_document"] = DocsMap().also{block(it)}.map}
        inline fun limit(limit:Long){v["n_results"] = limit}
        inline fun addEmbedding(key:String, embed:LFD){
            (v.getOrPut("query_embeddings") { arrayListOf<LFD>() } as LF<LFD>).let{
                it.add(embed)
                keys[it.lastIndex] = key
            }
        }
        @PublishedApi internal fun removeInclude(k:String) {
            (v.getOrPut("include") { arrayListOf("ids", "documents", "metadatas", "embeddings", "distances") } as LFS).remove(k)
        }
        inline fun excludeDocument(){removeInclude("documents")}
        inline fun excludeEmbedding(){removeInclude("embeddings")}
        inline fun excludeMeta(){removeInclude("metadatas")}
        inline fun excludeId(){removeInclude("ids")}
        inline fun excludeDistances(){removeInclude("distances")}
    }
    suspend inline fun add(isNorm:Boolean, block:AddDSL.()->Unit):Boolean{
        val request:MFA = hashMapOf()
        AddDSL(request).block()
        if(isNorm) request["embeddings"] = (request["embeddings"] as LF<LFD>).awaitAll(Dispatchers.IO){
            VectorDB.l2Norm(it)
        }
        return try {
            getWebClient(v.cr.url).post()
                .uri("/api/v2/tenants/${v.tenant}/databases/${v.db}/collections/${v.collection}/add")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(JSON.save(request))
                .retrieve()
                .awaitBody<String>()
            true
        }catch(e:Exception){
            false
        }
    }
    suspend fun remove(block:(DelDSL)->Unit):Boolean{
        val request:MFA = hashMapOf()
        block(DelDSL(request))
        return try {
            getWebClient(v.cr.url).post()
                .uri("/api/v2/tenants/${v.tenant}/databases/${v.db}/collections/${v.collection}/delete")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(JSON.save(request))
                .retrieve()
                .awaitBody<String>()
            true
        }catch(e:Exception){
            false
        }
    }
    suspend fun get(block:(GetDSL)->Unit = {}):LF<Record>?{
        val request:MFA = hashMapOf()
        block(GetDSL(request))
        return try {
            val res = getWebClient(v.cr.url).post()
                .uri("/api/v2/tenants/${v.tenant}/databases/${v.db}/collections/${v.collection}/get")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(JSON.save(request))
                .retrieve()
                .awaitBody<String>()
            val j:Map<String, Any?> = JSON.restore(res)
            val ids = j["ids"] as? LFS
            val embeddings = j["embeddings"] as? LF<LFD>
            val documents = j["documents"] as? LFS
            val metadatas = j["metadatas"] as? LF<MFS?>
            val uris = j["uris"] as? LFS
            val distance = j["distances"] as? LF<Double?>
            return (ids ?: documents ?: embeddings ?: metadatas ?: uris)?.foldIndexed(arrayListOf()) {index, acc, doc->
                acc.add(Record(
                    ids?.get(index) ?: "",
                    embeddings?.get(index),
                    documents?.get(index) ?: "",
                    metadatas?.get(index),
                    uris?.get(index) ?: "",
                    distance?.get(index) ?: 0.0
                ))
                acc
            } ?: null
        }catch(e:Exception){
            null
        }
    }
    suspend fun query(block:QueryDSL.()->Unit):MF<LF<Record>>?{
        val request = QueryDSL()
        return try {
            val res = getWebClient(v.cr.url).post()
                .uri("/api/v2/tenants/${v.tenant}/databases/${v.db}/collections/${v.collection}/query")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(JSON.save(request.v))
                .retrieve()
                .awaitBody<String>()
            val j:Map<String, Any?> = JSON.restore(res)
            val ids = j["ids"] as LF<LFS>
            val embeddings = j["embeddings"] as? LF<LF<LFD>>
            val documents = j["documents"] as? LF<LFS>
            val metadatas = j["metadatas"] as? LF<LF<MFS?>>
            val uris = j["uris"] as? LF<LFS>
            val distance = j["distances"] as? LF<LFD>
            return ids.foldIndexed(hashMapOf()) {index, acc, ids->
                request.keys[index]?.let{
                    val v = arrayListOf<Record>()
                    acc[it] = v
                    ids.forEachIndexed {i, id->
                        v.add(Record(
                            id,
                            embeddings?.get(index)?.get(i),
                            documents?.get(index)?.get(i) ?: "",
                            metadatas?.get(index)?.get(i),
                            uris?.get(index)?.get(i) ?: "",
                            distance?.get(index)?.get(i) ?: 0.0
                        ))
                    }
                }
                acc
            }
        }catch(e:Exception){
            null
        }
    }
}

CollectionOpt.kt

package kore.chromaDB

import kore.vo.VO
import kore.vo.field.map.stringMap
import kore.vo.field.value.string

class CollectionOpt: VO(){
    var name by string
    var configuration by stringMap
}

DB.kt

package kore.chromaDB

@JvmInline value class DB @PublishedApi internal constructor(val v: Struct){
    suspend inline fun colletion(name:String): COLLECTION {
        val u = v.cr.unsafe()
        if(!u.hasCollection(v.tenant, v.db, name)) {
            u.addCollection(v.tenant, v.db, name) ?: throw Exception("Failed to add db: ${v.tenant}, ${v.db}, $name")
        }
        v.collection = v.cr.tables["${v.tenant}@${v.db}@$name"]!!
        return COLLECTION(v)
    }
    suspend inline fun <T> colletion(name:String, block: COLLECTION.()->T):T = colletion(name).let(block)
    suspend fun hasCollection(name:String):Boolean = v.cr.unsafe().hasCollection(v.tenant, v.db, name)
    suspend fun addCollection(name:String):String? = v.cr.unsafe().addCollection(v.tenant, v.collection, name)
    suspend fun removeCollection(name:String):Boolean {
        if(!hasCollection(name)) return false
        return v.cr.unsafe().removeCollection(v.tenant, v.db, name, v.cr.tables["${v.tenant}@${v.db}@$name"]!!)
    }
}

Record.kt

package kore.chromaDB

import kore.vo.field.LFD
import kore.vo.field.MFS

class Record(
    var id:String,
    var embedding:LFD?,
    var document:String,
    var metadata:MFS?,
    var uri:String,
    var distances:Double
)

Struct.kt

package kore.chromaDB

class Struct @PublishedApi internal constructor(val cr:ChromaDB, val tenant:String = "", var db:String = "", var collection:String = "")

TENANT.kt

package kore.chromaDB

@JvmInline value class TENANT @PublishedApi internal constructor(val v: Struct){
    suspend inline fun db(name:String): DB {
        val u = v.cr.unsafe()
        if(!u.hasDB(v.tenant, name)) {
            u.addDB(v.tenant, name) ?: throw Exception("Failed to add db: ${v.tenant}, $name")
        }
        v.db = name
        return DB(v)
    }
    suspend inline fun <T> db(name:String, block: DB.()->T):T = db(name).let(block)
    suspend fun hasDB(name:String):Boolean = v.cr.unsafe().hasDB(v.tenant, name)
    suspend fun addDB(name:String):String? = v.cr.unsafe().addDB(v.tenant, name)
    suspend fun removeDB(name:String):Boolean {
        if(!hasDB(name)) return false
        return v.cr.unsafe().removeDB(v.tenant, name)
    }
}

Unsafe.kt

package kore.chromaDB

import kore.net.getWebClient
import kore.savable.JSON
import kore.vo.save.toJSON
import org.springframework.http.MediaType
import org.springframework.web.reactive.function.client.awaitBody

@JvmInline value class Unsafe @PublishedApi internal constructor(val cr:ChromaDB){
    suspend fun hasTenant(name:String):Boolean
    = if(cr.tenants.contains(name)) true else try{
        val res = getWebClient(cr.url).get()
            .uri("/api/v2/tenants/$name")
            .retrieve()
            .awaitBody<String>()
        val result = JSON.restore(res)["name"] == name
        if(result) cr.tenants[name] = name else cr.tenants.remove(name)
        result
    }catch(e:Exception){
        false
    }
    suspend fun addTenant(name:String):String?
    = if(hasTenant(name)) name else try {
        getWebClient(cr.url).post()
            .uri("/api/v2/tenants")
            .contentType(MediaType.APPLICATION_JSON)
            .bodyValue("""{"name":"$name"}""")
            .retrieve()
            .awaitBody<String>()
        cr.tenants[name] = name
        name
    }catch(e:Exception){
        null
    }
    suspend fun removeTenant(name: String): Boolean
    = try {
        getWebClient(cr.url).delete()
            .uri("/api/v2/tenants/$name")
            .retrieve()
            .awaitBody<String>()
        cr.tenants.remove(name)
        true
    }catch(e:Exception){
        false
    }
    suspend fun hasDB(tenant:String, name:String):Boolean
    = if(cr.dbs.contains("$tenant@$name")) true else try{
        val res = getWebClient(cr.url).get()
            .uri("/api/v2/tenants/$tenant/databases")
            .retrieve()
            .awaitBody<String>()
        val info:ArrayList<HashMap<String, Any?>> = JSON.restore("""{"a": $res}""")["a"] as ArrayList<HashMap<String, Any?>>
        if(info.find{it["name"] == name} != null){
            cr.dbs["$tenant@$name"] = name
            true
        }else{
            cr.dbs.remove("$tenant@$name")
            false
        }
    }catch(e:Exception){
        false
    }
    suspend fun addDB(tenant:String, name:String):String?
    = try{
        getWebClient(cr.url).post()
            .uri("/api/v2/tenants/$tenant/databases")
            .contentType(MediaType.APPLICATION_JSON)
            .bodyValue("""{"name":"$name"}""")
            .retrieve()
            .awaitBody<String>()
        cr.dbs["$tenant@$name"] = name
        name
    }catch(e:Exception){
        null
    }
    suspend fun removeDB(tenant:String, name: String):Boolean
    = try {
        getWebClient(cr.url).delete()
            .uri("/api/v2/tenants/$tenant/databases/$name")
            .retrieve()
            .awaitBody<String>()
        cr.dbs.remove("$tenant@$name")
        true
    }catch(e:Exception){
        false
    }
    suspend fun hasCollection(tenant:String, db:String, name:String):Boolean
    = if(cr.tables.contains("$tenant@$db@$name")) true else try{
        val res = getWebClient(cr.url).get()
            .uri("/api/v2/tenants/$tenant/databases/$db/collections")
            .retrieve()
            .awaitBody<String>()
        val info:ArrayList<HashMap<String, Any?>> = JSON.restore("""{"a": $res}""")["a"] as ArrayList<HashMap<String, Any?>>
        val r = info.find{it["name"] == name}
        if(r != null){
            cr.tables["$tenant@$db@$name"] = r["id"] as String
            true
        }else{
            cr.tables.remove("$tenant@$db@$name")
            false
        }
    }catch(e:Exception){
        false
    }
    suspend fun addCollection(tenant:String, db:String, name:String, opt:CollectionOpt? = null):String?
    = try {
        val res = getWebClient(cr.url).post()
            .uri("/api/v2/tenants/$tenant/databases/$db/collections")
            .contentType(MediaType.APPLICATION_JSON)
            .bodyValue(opt?.let{
                it.name = name
                it.toJSON()
            } ?: """{"name":"$name"}""")
            .retrieve()
            .awaitBody<String>()
        val r = JSON.restore(res)["id"] as String
        cr.tables["$tenant@$db@$name"] = r
        r
    }catch(e:Exception){
        null
    }
    suspend fun removeCollection(tenant:String, db:String, name:String, id:String):Boolean
    = try {
        getWebClient(cr.url).delete()
            .uri("/api/v2/tenants/$tenant/databases/$db/collections/$id")
            .retrieve()
            .awaitBody<String>()
        cr.tables.remove("$tenant@$db@$name")
        true
    }catch(e:Exception){
        false
    }
}