Агент на Kotlin без фреймворков

Единственная часть, которую было сложно расширять и поддерживать в прошлой статье — сам агент. Тут мы напишем такое решение, чтобы агент собирался, как конструктор, и чтобы любую часть можно было легко вынести и переиспользовать в других агентах.

Если нет времени читать, можете глянуть PR с реализацией агента на графах и PR с добавлением RAG.

С 2015, когда я начал карьеру, постоянно появляются библиотеки-решения для организации UI-архитектуры на основе паттернов: MVC, MVP, MVVM, TEA, VIPER, Flux/Redux. Я пробовал все паттерны из перечисленных и смело могу сказать, что особой разницы нет, пока вся команда придерживается одного подхода. Но каждый раз использование чьей-то библиотеки приводило к страданиям. Потому что находился кейс, который библиотека упускала и не давала решить легко. Были баги, которые не починить без форка. Код библиотек запутан и сложен. Всегда проще было самому написать основу и жить с ней.

То же самое и с фреймворками по написанию ИИ-агентов. Лучше разберитесь с основами и напишите под себя легковесное решение, которое вы будете понимать и контролировать.

Я встречал java-бэкенд разработчиков, которые несколько лет работали с Postgres, но не знают, как взять лок. Разгадка — фреймворк Hibernate. Такой разработчик может ходить в базу в цикле или внутри транзакций к базе выполнять HTTP-запросы. Конечно, есть исключения, но при прочих равных, фреймворк лишает понимания, задерживает развитие, способствует написанию неоптимального кода и даже негативно влияет на экологию (сколько энергии было потрачено на разработку и компиляцию этого фреймворка? SW).

То же самое и с фреймворками по написанию ИИ-агентов. Под капотом происходит всего лишь вызов нескольких ручек HTTP, а ваш фреймворк падает с out of memory.

import com.fasterxml.jackson.annotation.JsonProperty
import java.util.*

object GigaResponse {

    data class Token(
        @JsonProperty("access_token") val accessToken: String,
        @JsonProperty("expires_at") val expiresAt: Date
    )

    sealed interface Chat {
        data class Ok(val choices: List<Choice>, val created: Long, val model: String) : Chat
        data class Error(val status: Int, val message: String) : Chat
    }

    data class Choice(
        val message: Message,
        val index: Int,
        @JsonProperty("finish_reason")
        val finishReason: String
    )

    data class Message(
        val content: String,
        val role: GigaMessageRole,
        @JsonProperty("function_call")
        val functionCall: FunctionCall? = null,
        @JsonProperty("functions_state_id")
        val functionsStateId: String?
    )

    data class FunctionCall(
        val name: String,
        val arguments: Map<String, Any>
    )
}

object GigaRequest {
    data class Chat(
        val model: String = "GigaChat-Max",
        val messages: List<Message>,
        @JsonProperty("function_call")
        val functionCall: String = "auto",
        val functions: List<Function>? = null,
    )

    data class Message(
        val role: GigaMessageRole,
        val content: String, // Could be String or FunctionCall object
        @JsonProperty("functions_state_id")
        val functionsStateId: String? = null
    )

    data class Function(
        val name: String,
        val description: String,
        val parameters: Parameters
    )

    data class Parameters(
        val type: String,
        val properties: Map<String, Property>
    )

    data class Property(
        val type: String,
        val description: String? = null
    )
}

@Suppress("EnumEntryName")
enum class GigaMessageRole { system, user, assistant, function }

const val MAX_TOKENS = 8192

enum class GigaModel(val alias: String, val maxTokens: Int) {
    Lite("GigaChat-2", MAX_TOKENS),
    Pro("GigaChat-Pro", MAX_TOKENS),
    Max("GigaChat-Max", MAX_TOKENS),
}

fun String.toSystemPromptMessage() = GigaRequest.Message(
    role = GigaMessageRole.system,
    content = this
)

import io.ktor.client.*
import io.ktor.client.call.*
import io.ktor.client.engine.cio.*
import io.ktor.client.request.*
import io.ktor.client.request.forms.*
import io.ktor.http.*

object GigaAuth {
    suspend fun requestToken(apiKey: String): String {
        val client = HttpClient(CIO) {
            gigaDefaults()
        }
        val response = client.submitForm(
            url = "https://ngw.devices.sberbank.ru:9443/api/v2/oauth",
            formParameters = Parameters.build {
                append("scope", "GIGACHAT_API_PERS")
            }
        ) {
            header("Content-Type", "application/x-www-form-urlencoded")
            header("Authorization", "Basic $apiKey")
        }.body<GigaResponse.Token>()

        client.close()
        return response.accessToken
    }
}

import io.ktor.client.*
import io.ktor.client.call.*
import io.ktor.client.engine.cio.*
import io.ktor.client.plugins.auth.*
import io.ktor.client.plugins.auth.providers.*
import io.ktor.client.plugins.logging.LogLevel
import io.ktor.client.plugins.logging.Logging
import io.ktor.client.request.*
import io.ktor.http.*

class GigaChatAPI(private val auth: GigaAuth) {
    private val client = HttpClient(CIO) {
        var token = "" // get form env, or cache, or db
        val gigaKey = System.getenv("GIGA_KEY")
        gigaDefaults()
        install(Auth) {
            bearer {
                loadTokens {
                    BearerTokens(token, "")
                }
                refreshTokens {
                    token = auth.requestToken(gigaKey)
                    BearerTokens(token, "")
                }
            }
        }
        install(Logging) {
            val envLevel = LogLevel.INFO
            level = envLevel
            sanitizeHeader { it.equals(HttpHeaders.Authorization, true) }
        }
    }

    suspend fun message(body: GigaRequest.Chat): GigaResponse.Chat {
        val response = client.post("https://gigachat.devices.sberbank.ru/api/v1/chat/completions") {
            setBody(body)
        }
        return when {
            response.status.isSuccess() -> response.body<GigaResponse.Chat.Ok>()
            else -> response.body<GigaResponse.Chat.Error>()
        }
    }

    fun clear() = client.close()
}

import com.fasterxml.jackson.databind.DeserializationFeature
import io.ktor.client.*
import io.ktor.client.engine.cio.*
import io.ktor.client.plugins.*
import io.ktor.client.plugins.contentnegotiation.*
import io.ktor.client.request.*
import io.ktor.http.*
import io.ktor.serialization.jackson.*
import java.security.cert.X509Certificate
import java.util.*
import javax.net.ssl.X509TrustManager

fun HttpClientConfig<CIOEngineConfig>.gigaDefaults() {
    this.defaultRequest {
        header(HttpHeaders.ContentType, "application/json")
        header(HttpHeaders.Accept, "application/json")
        header("RqUID", UUID.randomUUID().toString())
    }
    install(HttpTimeout) {
        requestTimeoutMillis = 40000
    }
    install(ContentNegotiation) {
        jackson { this.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES) }
    }
    engine {
        https {
            trustManager = object : X509TrustManager {
                override fun checkClientTrusted(chain: Array<out X509Certificate>?, authType: String?) {}
                override fun checkServerTrusted(chain: Array<out X509Certificate>?, authType: String?) {}
                override fun getAcceptedIssuers(): Array<X509Certificate> = arrayOf()
            }
        }
    }
}

@Target(AnnotationTarget.PROPERTY)
@Retention(AnnotationRetention.RUNTIME)
annotation class InputParamDescription(val value: String)

interface ToolSetup<Input> {

    val name: String
    val description: String

    operator fun invoke(input: Input): String
}

class BadInputException(msg: String) : Exception(msg)

Пример реализации тула:

object ToolRunBashCommand : ToolSetup<ToolRunBashCommand.Input> {
    override val name = "RunBashCommand"
    override val description = "Executes a bash command and returns its output"

    override fun invoke(input: Input): String {
        val process = ProcessBuilder("bash", "-c", input.command)
            .redirectErrorStream(true)
            .start()
        val output = process.inputStream.bufferedReader().use(BufferedReader::readText)
        val exitCode = process.waitFor()
        if (exitCode != 0) throw RuntimeException("Command failed with exit code $exitCode")
        return output.trim()
    }

    data class Input(
        @InputParamDescription("The bash command to run, e.g., 'ls', 'echo Hello', './gradlew tasks'")
        val command: String
    )
}

Маппинг на модели гигачата:

import com.dumch.tool.InputParamDescription
import com.dumch.tool.ToolSetup
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import kotlin.reflect.KCallable
import kotlin.reflect.full.declaredMembers
import kotlin.reflect.full.findAnnotation


interface GigaToolSetup {
    val fn: GigaRequest.Function
    operator fun invoke(functionCall: GigaResponse.FunctionCall): GigaRequest.Message
}

val gigaJsonMapper = jacksonObjectMapper()

inline fun <reified Input> ToolSetup<Input>.toGiga(): GigaToolSetup {
    val toolSetup = this
    return object : GigaToolSetup {
        override val fn: GigaRequest.Function = GigaRequest.Function(
            name = toolSetup.name,
            description = toolSetup.description,
            parameters = GigaRequest.Parameters(
                "object",
                properties = HashMap<String, GigaRequest.Property>().apply {
                    val clazz = Input::class
                    for (kProperty: KCallable<*> in clazz.declaredMembers) {
                        val annotation = kProperty.findAnnotation<InputParamDescription>() ?: continue
                        val description = annotation.value
                        val type = kProperty.returnType.toString().substringAfterLast(".").lowercase()
                        val gigaProperty = GigaRequest.Property(type, description)
                        put(kProperty.name, gigaProperty)
                    }
                }
            )
        )

        override fun invoke(
            functionCall: GigaResponse.FunctionCall,
        ): GigaRequest.Message {
            return try {
                val input: Input = gigaJsonMapper.convertValue(functionCall.arguments, Input::class.java)
                val toolResult = toolSetup.invoke(input)
                val gigaResult = gigaJsonMapper.writeValueAsString(
                    mapOf("result" to toolResult)
                )
                GigaRequest.Message(
                    role = GigaMessageRole.function,
                    content = gigaResult,
                )
            } catch (e: Exception) {
                e.toGigaToolMessage()
            }
        }
    }
}

fun Exception.toGigaToolMessage(): GigaRequest.Message {
    return GigaRequest.Message(
        role = GigaMessageRole.function,
        content = """{"result": "${message ?: toString()}"}""",
    )
}

interface Node<IN, OUT> {
    val name: String
    suspend fun execute(ctx: AgentContext<IN>, runtime: GraphRuntime): AgentContext<OUT>
}

/**
 * Create new [Node] implementation based on [op]
 */
fun <IN, OUT> Node(
    name: String,
    op: suspend (AgentContext<IN>) -> AgentContext<OUT>,
): Node<IN, OUT> = object : Node<IN, OUT> {
    override val name: String = "Node $name; ${Integer.toHexString(hashCode())}"
    override suspend fun execute(ctx: AgentContext<IN>, runtime: GraphRuntime) = op(ctx)
}

data class AgentContext<I>(
    val input: I,
    val settings: AgentSettings,
    val history: List<GigaRequest.Message>,
    val tools: List<GigaRequest.Function>,
    val systemPrompt: String,
) {
    inline fun <reified O> map(
        settings: AgentSettings = this.settings,
        history: List<GigaRequest.Message> = this.history,
        activeTools: List<GigaRequest.Function> = this.tools,
        systemPrompt: String = this.systemPrompt,
        transform: (I) -> O = { it as O },
    ): AgentContext<O> = AgentContext(input = transform(input), settings, history, activeTools, systemPrompt)
}

data class AgentSettings(
    val model: String,
    val temperature: Float,
    val tools: Map<String, GigaToolSetup>
)

class Graph<IN, OUT> internal constructor(
    label: String,
    private val enter: Node<IN, *>,
    private val exit: Node<OUT, OUT>,
    private val retryPolicy: RetryPolicy,
    private val definition: GraphDefinition,
) : Node<IN, OUT> {

    private val runner = GraphRunner()

    override val name: String = "$label::graph"

    @Suppress("UNCHECKED_CAST")
    override suspend fun execute(ctx: AgentContext<IN>, runtime: GraphRuntime): AgentContext<OUT> {
        val result = runner.run(
            start = enter as Node<Any?, Any?>,
            seed = ctx as AgentContext<Any?>,
            runtime = runtime,
            definition = definition, // ребра передадим в Runner
            stopPredicate = { node, _ -> node === exit }
        )
        return result as AgentContext<OUT>
    }

    suspend fun start(
        seed: AgentContext<IN>,
        maxSteps: Int = 1000,
        onStep: ((step: StepInfo, node: Node<Any?, Any?>, ctx: AgentContext<Any?>) -> Unit)? = null,
    ): AgentContext<OUT> {
        val runtime = GraphRuntime(
            retryPolicy = retryPolicy,
            maxSteps = maxSteps,
            onStep = onStep,
        )
        return execute(seed, runtime)
    }
}

class GraphBuilder<IN, OUT> internal constructor(
    private val graphName: String,
    private val retryPolicy: RetryPolicy,
) {
    val nodeInput: Node<IN, IN> = Node("$graphName::enter") { it }
    val nodeFinish: Node<OUT, OUT> = Node("$graphName::exit") { it }

    private val transitions: MutableMap<Node<*, *>, MutableList<Transition<*>>> = mutableMapOf()

    fun <IN, OUT, OUT2> Node<IN, OUT>.edgeTo(target: Node<OUT, OUT2>): Node<OUT, OUT2> {
        registerTransition(this, Transition.Static(target))
        return target
    }

    fun <IN, OUT> Node<IN, OUT>.edgeTo(router: suspend (AgentContext<OUT>) -> Node<OUT, *>): Unit {
        registerTransition(this, Transition.Dynamic(router))
    }

    private fun <OUT> registerTransition(from: Node<*, OUT>, transition: Transition<OUT>) {
        val bucket = transitions.getOrPut(from) { mutableListOf() }
        bucket += transition
    }

    internal fun build(): Graph<IN, OUT> = Graph(
        graphName,
        nodeInput,
        nodeFinish,
        retryPolicy,
        GraphDefinition(transitions.mapValues { it.value.toList() }),
    )
}

// Вынесем еще и абстракцию для хранения ребер
internal class GraphDefinition(
    private val transitions: Map<Node<*, *>, List<Transition<*>>>,
) {
    @Suppress("UNCHECKED_CAST")
    suspend fun nextNodes(node: Node<Any?, Any?>, ctx: AgentContext<Any?>): List<Node<Any?, *>> {
        val registered = transitions[node] as? List<Transition<Any?>> ?: emptyList()
        if (registered.isEmpty()) return emptyList()

        val next = ArrayList<Node<Any?, *>>(registered.size)
        for (transition in registered) {
            when (transition) {
                is Transition.Static -> next.addOrWarn(transition.target as Node<Any?, *>)
                is Transition.Dynamic -> next.addOrWarn(transition.router(ctx) as Node<Any?, *>)
            }
        }
        return next
    }

    private fun MutableCollection<Node<Any?, *>>.addOrWarn(node: Node<Any?, *>) {
        if (contains(node)) {
            add(node)
        }
    }
}

internal sealed interface Transition<OUT> {
    class Static<OUT>(val target: Node<OUT, *>) : Transition<OUT>
    class Dynamic<OUT>(val router: suspend (AgentContext<OUT>) -> Node<OUT, *>) : Transition<OUT>
}

// Ниже всего лишь бойлерплейт для делегатов (by).

fun <I, O> buildGraph(
    name: String = "Graph",
    retryPolicy: RetryPolicy = RetryPolicy(),
    configure: GraphBuilder<I, O>.() -> Unit
): Graph<I, O> {
    val builder = GraphBuilder<I, O>(name, retryPolicy)
    builder.configure()
    return builder.build()
}

fun <I, O> graph(
    name: String? = null,
    retryPolicy: RetryPolicy = RetryPolicy(),
    configure: GraphBuilder<I, O>.() -> Unit
): ReadOnlyProperty<Any?, Graph<I, O>> = GraphDelegate(name, retryPolicy, configure)

private class GraphDelegate<I, O>(
    private val nameHint: String?,
    private val retryPolicy: RetryPolicy,
    private val configure: GraphBuilder<I, O>.() -> Unit,
) : ReadOnlyProperty<Any?, Graph<I, O>> {
    private var cached: Graph<I, O>? = null

    override fun getValue(thisRef: Any?, property: KProperty<*>): Graph<I, O> {
        return cached ?: build(property.name).also { cached = it }
    }

    private fun build(propertyName: String): Graph<I, O> {
        val name = nameHint ?: propertyName
        val builder = GraphBuilder<I, O>(name, retryPolicy)
        builder.configure()
        return builder.build()
    }
}

internal class GraphRunner {

    suspend fun run(
        start: Node<Any?, Any?>,
        seed: AgentContext<Any?>,
        runtime: GraphRuntime,
        definition: GraphDefinition,
        stopPredicate: ((Node<Any?, Any?>, AgentContext<Any?>) -> Boolean)? = null,
    ): AgentContext<Any?> {
        val queue = ArrayDeque<Frame>().apply { add(Frame(start, seed, 0)) }
        val leaves = mutableListOf<AgentContext<*>>()
        var lastCtx: AgentContext<Any?> = seed

        try {
            while (queue.isNotEmpty() && currentCoroutineContext().isActive) {
                if (runtime.counter.get() >= runtime.maxSteps) {
                    error("Graph maxSteps (${runtime.maxSteps}) reached — potential loop")
                }

                val frame = queue.removeFirst()
                val outCtx = executeWithRetry(frame.node, frame.ctx, runtime)
                val stepInfo = StepInfo(currentGraphIndex = frame.depth, index = runtime.counter.get())
                runtime.onStep?.invoke(stepInfo, frame.node, outCtx)
                lastCtx = outCtx

                if (stopPredicate?.invoke(frame.node, outCtx) == true) return outCtx

                val nextNodes = definition.nextNodes(frame.node, outCtx)
                if (nextNodes.isEmpty()) {
                    leaves += outCtx
                } else {
                    for (child in nextNodes) {
                        @Suppress("UNCHECKED_CAST")
                        queue.add(Frame(child as Node<Any?, Any?>, outCtx, frame.depth + 1))
                    }
                }
                runtime.counter.incrementAndGet()
            }
        } catch (cancel: CancellationException) {
            throw GraphCancellation(lastCtx, cancel)
        }

        @Suppress("UNCHECKED_CAST")
        return leaves.lastOrNull() as? AgentContext<Any?> ?: lastCtx
    }

    private suspend fun executeWithRetry(
        node: Node<Any?, Any?>,
        inCtx: AgentContext<Any?>,
        runtime: GraphRuntime,
    ): AgentContext<Any?> {
        val policy = runtime.retryPolicy
        var attempt = 0
        var lastError: Throwable? = null
        while (attempt < policy.maxAttempts) {
            attempt++
            try {
                return node.execute(inCtx, runtime)
            } catch (t: Throwable) {
                if (t is CancellationException) throw t
                lastError = t
                val shouldRetry = policy.shouldRetry(t, inCtx, node, attempt)
                val attemptsLeft = policy.maxAttempts - attempt
                if (!shouldRetry || attemptsLeft <= 0) break
            }
        }
        throw lastError ?: IllegalStateException("Unknown failure in node ${node.name}")
    }

    private data class Frame(
        val node: Node<Any?, Any?>,
        val ctx: AgentContext<Any?>,
        val depth: Int,
    )
}

class GraphCancellation(
    val lastContext: AgentContext<*>,
    cause: CancellationException? = null
) : CancellationException(cause?.message) {
    init {
        initCause(cause)
    }
}

data class StepInfo(
    /**
     * Sequential index of the executed node within the run (starting from 0).
     */
    val index: Int,
    /**
     * Sequential index of the executed node within the current graph run (starting from 0).
     */
    val currentGraphIndex: Int,
)

class GraphRuntime private constructor(
    val retryPolicy: RetryPolicy,
    val maxSteps: Int,
    val onStep: ((step: StepInfo, node: Node<Any?, Any?>, ctx: AgentContext<Any?>) -> Unit)? = null,
    val counter: AtomicInteger
) {
    constructor(
        retryPolicy: RetryPolicy,
        maxSteps: Int,
        onStep: ((step: StepInfo, node: Node<Any?, Any?>, ctx: AgentContext<Any?>) -> Unit)? = null,
    ): this(retryPolicy, maxSteps, onStep, counter = AtomicInteger())
}

data class RetryPolicy(
    val maxAttempts: Int = 2,
    val shouldRetry: suspend (
        error: Throwable,
        ctx: AgentContext<*>,
        node: Node<*, *>?,
        attempt: Int
    ) -> Boolean = { _, _, _, _ -> true }
)

object NodesCommon {
    val stringToReq: Node<String, GigaRequest.Chat> = Node("String->Request") { ctx ->
        val usrMsg = GigaRequest.Message(GigaMessageRole.user, ctx.input)
        val history = ArrayList(ctx.history).apply {
            if (isEmpty()) add(ctx.systemPrompt.toSystemPromptMessage())
            add(usrMsg)
        }
        ctx.map(history = history) { ctx.toGigaRequest(history) }
    }

    val respToString: Node<GigaResponse.Chat, String> = Node("Response->String") { ctx ->
        when (val input = ctx.input) {
            is GigaResponse.Chat.Error -> ctx.map { input.message }
            is GigaResponse.Chat.Ok -> ctx.map { input.choices.last().message.content }
        }
    }

    val toolUse: Node<GigaResponse.Chat, GigaRequest.Chat> = Node("toolUse") { ctx ->
        val fnCallMessages = fnCallMessages(ctx)
        val history = ArrayList(ctx.history).apply { addAll(fnCallMessages) }
        ctx.map(history = history) { ctx.toGigaRequest(history) }
    }

    private suspend fun fnCallMessages(ctx: AgentContext<GigaResponse.Chat>): List<GigaRequest.Message> {
        val fnCallMessages = (ctx.input as GigaResponse.Chat.Ok).choices.mapNotNull { choice ->
            val msg = choice.message
            if (msg.functionCall != null && msg.functionsStateId != null) {
                executeTool(ctx.settings, msg.functionCall)
            } else null
        }
        return fnCallMessages
    }

    private suspend fun executeTool(
        settings: AgentSettings,
        functionCall: GigaResponse.FunctionCall,
    ): GigaRequest.Message {
        val tools = settings.tools
        val fn: GigaToolSetup = tools[functionCall.name] ?: return GigaRequest.Message(
            GigaMessageRole.function, """{"result":"no such function ${functionCall.name}"}"""
        )
        return fn.invoke(functionCall)
    }
}

fun <T> AgentContext<T>.toGigaRequest(history: List<GigaRequest.Message>): GigaRequest.Chat {
    val ctx = this
    return GigaRequest.Chat(
        model = ctx.settings.model,
        messages = history,
        functions = ctx.tools,
    )
}

class NodesLLM(llmApi: GigaChatAPI) {

    val chat: Node<GigaRequest.Chat, GigaResponse.Chat> = Node("llmCall") { ctx ->
        val response = withContext(Dispatchers.IO) {
            llmApi.message(ctx.input)
        }
        val history = ArrayList(ctx.history).apply {
            if (response is GigaResponse.Chat.Ok) {
                addAll(response.choices.mapNotNull { it.toMessage() })
            }
        }
        ctx.map(history = history) { response }
    }

    /**
     * Restores the last message, and a system prompt. Other messages are transformed into TLDR
     */
    val summarize: Node<GigaResponse.Chat, GigaResponse.Chat> = Node("llmSummarize") { ctx ->
        val conversation = ArrayList(ctx.history)

        val summaryResponse: GigaResponse.Chat = withContext(Dispatchers.IO) {
            conversation.add(GigaRequest.Message(
                role = GigaMessageRole.user,
                content = "Резюмируй разговор",
            ))
            val request = ctx.toGigaRequest(conversation)
                .copy(functions = emptyList())
            llmApi.message(request)
        }

        val msg: GigaRequest.Message = when (summaryResponse) {
            is GigaResponse.Chat.Error -> {
                throw IOException(summaryResponse.message)
            }
            is GigaResponse.Chat.Ok -> summaryResponse.choices.mapNotNull { it.toMessage() }.last()
        }

        val newHistory = listOf(ctx.systemPrompt.toSystemPromptMessage(), ctx.history.last(), msg)
        ctx.map(history = newHistory) { summaryResponse }
    }

    private fun GigaResponse.Choice.toMessage(): GigaRequest.Message? {
        val msg = this.message
        val content: String = when {
            msg.content.isNotBlank() -> msg.content
            msg.functionCall != null -> gigaJsonMapper.writeValueAsString(
                mapOf("name" to msg.functionCall.name, "arguments" to msg.functionCall.arguments)
            )

            else -> return null
        }
        return GigaRequest.Message(
            role = msg.role,
            content = content,
            functionsStateId = msg.functionsStateId
        )
    }
}

object GigaResponse {
    // ... предыдущий код
  
    data class Embeddings(
        val data: List<Embedding>,
        val model: String,
        @JsonProperty("object") val objectType: String,
    )

    data class Embedding(
        val embedding: List<Double>,
        val index: Int,
        @JsonProperty("object") val objectType: String? = null,
    )
}

object GigaRequest {
    // ... предыдущий код
  
    data class Embeddings(
        val model: String = "Embeddings",
        val input: List<String>,
    )
}


class GigaChatAPI(private val auth: GigaAuth) {
    // ... предыдущий код

    suspend fun embeddings(body: GigaRequest.Embeddings): GigaResponse.Embeddings {
        val response = client.post("https://gigachat.devices.sberbank.ru/api/v1/embeddings") {
            setBody(body)
        }
        return when {
            response.status.isSuccess() -> response.body<GigaResponse.Embeddings>()
            response.status == HttpStatusCode.Unauthorized || response.status == HttpStatusCode.Forbidden -> TODO("Auth exception")
            else -> TODO("unexpected error")
        }
    }
}

object VectorDB {
    private const val INDEX_PATH = "build/rag_index"

    init {
        val isInitialized = File(INDEX_PATH).exists() // naive way to check initialization
        if (!isInitialized) {
            val dir = FSDirectory.open(Paths.get(INDEX_PATH))
            IndexWriter(dir, IndexWriterConfig()).use { }
        }
    }

    fun insert(data: List<String>, embeddings: List<List<Double>>) {
        val dir = FSDirectory.open(Paths.get(INDEX_PATH))
        IndexWriter(dir, IndexWriterConfig()).use { writer ->
            data.indices.forEach { idx ->
                val doc = Document()
                doc.add(StoredField("text", data[idx]))
                doc.add(KnnFloatVectorField("embedding", toFloatArray(embeddings[idx])))
                writer.addDocument(doc)
            }
        }
    }

    fun getAllTexts(): List<String> {
        val dir = FSDirectory.open(Paths.get(INDEX_PATH))
        DirectoryReader.open(dir).use { reader ->
            val list = mutableListOf<String>()
            for (i in 0 until reader.maxDoc()) {
                val doc = reader.document(i)
                doc.get("text")?.let { list.add(it) }
            }
            return list
        }
    }

    fun searchSimilar(embedding: List<Double>, limit: Int = 5): List<String> {
        val dir = FSDirectory.open(Paths.get(INDEX_PATH))
        DirectoryReader.open(dir).use { reader ->
            val searcher = IndexSearcher(reader)
            val query = KnnFloatVectorQuery("embedding", toFloatArray(embedding), limit)
            val topDocs = searcher.search(query, limit)
            val texts = mutableListOf<String>()
            topDocs.scoreDocs.forEach { sd ->
                searcher.doc(sd.doc).get("text")?.let { texts.add(it) }
            }
            return texts
        }
    }

    fun clearAllData() {
        val dir = FSDirectory.open(Paths.get(INDEX_PATH))
        IndexWriter(dir, IndexWriterConfig()).use { writer ->
            writer.deleteAll()
        }
    }

    private fun toFloatArray(list: List<Double>): FloatArray {
        val size = min(list.size, MAX_DIM)
        val arr = FloatArray(size)
        for (i in 0 until size) {
            arr[i] = list[i].toFloat()
        }
        return arr
    }

    private const val MAX_DIM = 1024
}