RAG (Retrieval-Augmented Generation), Function / Tools Calling, CAG (Cache-Augmented Generation) за LLM модели

{
  "name": "get_weather",
  "arguments": {
    "location": "New York"
 }
}

{
  "name": "search_docs",
  "arguments": {
    "query": "quantum computing basics",
    "top_k": 5,
    "min_similarity_score": 0.6
  }
}

const val DEFAULT_FUNCTION_CALLING_PROMPT = """
Call tools only when the user explicitly asks to check, verify, look up, find, or search for something. In all other cases, answer directly without calling any tool.

If you decide to invoke a function, output only a JSON object inside a fenced block marked as tool_code without any text before or after it.

Available tools:

search_docs
Description: Searches knowledge and returns the most relevant results as plain text.
Arguments:
query (string) — User string query to search in the vector store.
top_k (integer, default 5) — Number of results to return.
min_similarity_score (float, default 0.6) — Minimum similarity score from 0.0 (no filtering) to 1.0 (exact match).

get_weather
Description: Gets current weather for a location.
Arguments:
location (string) — City or place name.

JSON format inside the fenced block:
tool_code
{
"name": "<tool_name>",
"arguments": {
"<arg_name>": <arg_value>
}
}

Examples:
tool_code
{ "name": "search_docs", "arguments": { "query": "quantum computing basics", "top_k": 5, "min_similarity_score": 0.6 } }

tool_code
{ "name": "get_weather", "arguments": { "location": "New York" } }

Important rules:
Use tools only when explicitly requested by the user to check, verify, look up, find, or search. Otherwise, provide a direct answer.
When calling a tool, the output must be valid JSON with exactly two keys: name and arguments.
Do not include any explanations or extra content outside the tool_code block.

After receiving tool results (the result will be in JSON format):
1. Parse and process the result.
2. Convert it into a simple, clear, and human-readable natural-language response.
3. Always reply to the user with this processed answer as the next message.
4. If no relevant results are found, briefly explain that nothing relevant was found and suggest next steps.
"""

fun main() = runBlocking {

    val messages = mutableListOf(

         // Добавяме промпта в опашката със съобщения с роля system
        ChatMessage(
            role = "system",
            content = DEFAULT_FUNCTION_CALLING_PROMPT
        ),

        // Добавяме в опашката заявката на потребителя; от тази заявка е ясно,
        // че потребителят моли да се провери времето, моделът знае, че му е достъпна
        // функцията за проверка на времето, описана в промпта, и я извиква,
        //  като изпраща в отговор съобщение, започващо с tool_code
        ChatMessage(
            role = "user",
            content = "Check weather in London",
        )
    )
    var iterator = 0
    while (iterator < 3) {
        try {
            chatAnswerWithToolsStream(
                ChatCompletionRequest(
                    model = MODEL_ID,
                    messages = messages,
                    temperature = 0.2
                )
            ).collect { piece ->
                when (piece) {
                    is ChatResult.Debug -> {
                        println("DEBUG MESSAGE:\n${piece.message}")
                    }

                    is ChatResult.Message -> {
                        println("FINAL MESSAGE:\n${piece.message}")
                        iterator = 3
                        return@collect
                    }

                    is ChatResult.SearchDocs -> {
                        println("SEARCH DOCS data:\n$piece")
                        iterator = 3

                        // В този пример ще използвам само проверка на времето,
                        // но добавих още една функция, за да покажа, че може да са много
                        return@collect
                    }

                    is ChatResult.GetWeather -> {
                        println("GET WEATHER result:\n$piece")

                        // Ако моделът е решил да извика функция, добавяме нейното извикване
                        // в опашката със съобщения с роля assistant
                        // Важно е да се добавят всички съобщения в опашката, за да
                        // знае моделът цялата история на диалога
                        messages.add(
                            ChatMessage(
                                role = "assistant",
                                content = piece.message ?: ""
                            )
                        )

                        // правим заявка към API-то за данни за времето
                        val result = getCurrentWeatherByCity(
                            city = piece.location!!,
                        )
                        val jsonResult: String = kotlinxJsonConfig.encodeToString(result)
                        println("GET WEATHER request result:\n$jsonResult")

                        // добавяме в опашката със съобщения с роля user, тъй като този модел
                        // не поддържа специални роли за функционални извиквания
                        messages.add(
                            ChatMessage(
                                role = "user",
                                content = jsonResult
                            )
                        )
                        iterator++
                    }
                }
            }
        } catch (e: Throwable) {
            println(e.message ?: "error")
            break
        }
    }

}

// код за работа с OpenAI API, тук е важно извикването на функцията
// parseToolJson, която ще върне null, ако не е успяла да преобразува отговора във функционално извикване
fun chatAnswerWithToolsStream(req: ChatCompletionRequest): Flow<ChatResult> = flow {
    val response = client.post("$BASE_URL/api/v0/chat/completions") {
        contentType(ContentType.Application.Json)
        setBody(req.copy(stream = true))
    }
    val channel: ByteReadChannel = response.bodyAsChannel()

    val fullBuilder = StringBuilder()

    while (!channel.isClosedForRead) {
        val line: String = channel.readUTF8Line() ?: break
        val payload: String = extractDataLine(line) ?: continue

        val chunk: StreamChatChunk = runCatching {
            kotlinxJsonConfig.decodeFromString(StreamChatChunk.serializer(), payload)
        }.getOrNull() ?: continue

        val piece: String = extractChatDeltaContent(chunk)
        if (piece.isNotEmpty()) {
            fullBuilder.append(piece)
            emit(ChatResult.Debug(message = fullBuilder.toString()))
        }

        if (isFinished(chunk)) break
    }

    val fullMessage: String = fullBuilder.toString().trim()
    val toolResult: ChatResult? = parseToolJson(fullMessage)

    if (toolResult != null) {
        emit(toolResult)
    } else {
        emit(ChatResult.Message(message = fullMessage))
    }
}

fun parseToolJson(input: String): ChatResult? {
    try {
        val json = when {

            // ако отговорът започва с tool_code, разбираме, че това не е отговор към потребителя
            // а функционално извикване и по-нататък следва json с неговия тип и параметри
            input.startsWith("tool_code") -> input.replace("tool_code", "")
                .trim()

            input.startsWith("```tool_code") -> input.replace("```tool_code", "")
                .removeSuffix("```")
                .trim()

            else -> return null
        }
        val rawTool = runCatching {
            gson.fromJson(json, RawTool::class.java)
        }.getOrNull() ?: return null

        val jsonObject = rawTool.arguments?.asJsonObject ?: return null

    // определяме коя точно функция иска да извика моделът
        // и преобразуваме json в модел на данни.
        // Този код може да се подобри, като се върне грешка на модела, ако преобразуването
        // не успее, с описание, че моделът трябва да редактира отговора си и да го поправи
        return when (rawTool.name?.lowercase()) {
            "search_docs" -> {
                gson.fromJson(jsonObject, ChatResult.SearchDocs::class.java).copy(
                    message = json
                )
            }
            "get_weather" -> gson.fromJson(jsonObject, ChatResult.GetWeather::class.java).copy(
                message = json
            )
            else -> null
        }
    } catch (e: Exception) {
        println(e.message ?: "error")
        return null
    }
}

ext {
    ktor_version = "3.2.3"
    kotlinx_serialization_json_version = "1.8.1"
    logback_version = "1.4.14"
}

dependencies {
    implementation "io.ktor:ktor-client-core:$ktor_version"
    implementation "io.ktor:ktor-client-cio:$ktor_version"
    implementation "io.ktor:ktor-client-content-negotiation:$ktor_version"
    implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
    implementation "io.ktor:ktor-client-logging:$ktor_version"
    implementation "org.jetbrains.kotlinx:kotlinx-serialization-json:$kotlinx_serialization_json_version"
    implementation "ch.qos.logback:logback-classic:$logback_version"
    implementation "com.google.code.gson:gson:2.13.1"
}

import io.ktor.client.call.*
import io.ktor.client.request.*

suspend fun getCurrentWeatherByCity(
    apiKey: String = WEATHER_API_KEY,
    city: String,
    units: String = "metric",
    lang: String = "en"
): WeatherResponse {
    return client.get(OWM_BASE) {
        url {
            parameters.append("q", city)
            parameters.append("appid", apiKey)
            parameters.append("units", units)
            parameters.append("lang", lang)
        }
    }.body()
}

import com.google.gson.JsonElement
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonObject


@Serializable
data class ChatMessage(
    @SerialName("role")
    val role: String,
    @SerialName("content")
    val content: String
)

@Serializable
data class ChatCompletionRequest(
    @SerialName("model")
    val model: String,
    @SerialName("messages")
    val messages: List<ChatMessage>,
    @SerialName("temperature")
    val temperature: Double? = null,
    @SerialName("max_tokens")
    val maxTokens: Int? = null,
    @SerialName("stream")
    val stream: Boolean? = null
)

@Serializable
data class StreamChatChunk(
    @SerialName("choices")
    val choices: List<StreamChatChoice>? = null
)


@Serializable
data class ChatChoiceMessage(
    @SerialName("content")
    val content: String? = null
)

@Serializable
data class ChatChoice(
    @SerialName("message")
    val message: ChatChoiceMessage? = null
)

@Serializable
data class ChatResponse(
    @SerialName("choices")
    val choices: List<ChatChoice> = emptyList()
)


@Serializable
data class StreamChatChoice(
    @SerialName("delta")
    val delta: JsonObject? = null,
    @SerialName("finish_reason")
    val finishReason: String? = null
)

data class RawTool(
    val name: String?,
    val arguments: JsonElement?
)

sealed interface ChatResult {

    data class Debug(
        val message: String?,
    ) : ChatResult

    data class Message(
        val message: String?,
    ) : ChatResult

    data class SearchDocs(
        val message: String? = null,
        val query: String?,
        val topK: Int?,
        val minSimilarityScore: Float?
    ) : ChatResult

    data class GetWeather(
        val message: String? = null,
        val location: String?,
    ) : ChatResult

}

import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.contentOrNull
import kotlinx.serialization.json.jsonPrimitive

fun extractDataLine(line: String): String? {
    val t: String = line.trim()
    if (t.isEmpty()) return null
    if (t == "data: [DONE]" || t == "[DONE]") return null
    if (t.startsWith("event:")) return null
    return if (t.startsWith("data:")) t.removePrefix("data:").trim() else t
}

fun extractChatDeltaContent(chunk: StreamChatChunk): String {
    val choice: StreamChatChoice = chunk.choices?.firstOrNull() ?: return ""
    val delta = choice.delta ?: return ""
    val el: JsonElement? = delta["content"]
    return el?.jsonPrimitive?.contentOrNull ?: ""
}

fun extractTextDelta(chunk: StreamTextChunk): String {
    val choice: StreamTextChoice = chunk.choices?.firstOrNull() ?: return ""
    val direct: String? = choice.text
    if (direct != null) return direct
    val delta = choice.delta
    val el: JsonElement? = delta?.get("content") ?: delta?.get("text")
    return el?.jsonPrimitive?.contentOrNull ?: ""
}

fun isFinished(chunk: StreamChatChunk): Boolean {
    val reason: String? = chunk.choices?.firstOrNull()?.finishReason
    return !reason.isNullOrEmpty()
}

fun isFinished(chunk: StreamTextChunk): Boolean {
    val reason: String? = chunk.choices?.firstOrNull()?.finishReason
    return !reason.isNullOrEmpty()
}

import com.google.gson.FieldNamingPolicy
import com.google.gson.Gson
import com.google.gson.GsonBuilder
import io.ktor.client.*
import io.ktor.client.engine.cio.*
import io.ktor.client.plugins.*
import io.ktor.client.plugins.contentnegotiation.*
import io.ktor.serialization.kotlinx.json.*
import kotlinx.serialization.json.Json

const val BASE_URL: String = "http://localhost:1234"

const val OWM_BASE = "https://api.openweathermap.org/data/2.5"

const val WEATHER_API_KEY = "..."

const val MODEL_ID = "google/gemma-3n-e4b"

val gson: Gson = GsonBuilder()
    .setFieldNamingPolicy(FieldNamingPolicy.LOWER_CASE_WITH_UNDERSCORES)
    .create()

val kotlinxJsonConfig: Json = Json {
    ignoreUnknownKeys = true
    prettyPrint = false
}

val client: HttpClient = HttpClient(CIO) {
    install(ContentNegotiation) {
        json(kotlinxJsonConfig)
    }
    install(HttpTimeout) {
        requestTimeoutMillis = 60_000
        connectTimeoutMillis = 10_000
        socketTimeoutMillis = 60_000
    }
}

import dspy
import requests
from typing import Dict

dspy.enable_logging()

OPENWEATHER_API_KEY = "..."

# Задаваме доставчик за LLM
lm = dspy.LM(
    "openai/google/gemma-3n-e4b",
    api_base="http://localhost:1234/v1",
    api_key="lm-studio",
    temperature=0
)
dspy.configure(lm=lm)

# Проверяваме дали доставчикът е намерен
print("Provider:\n", type(lm.provider))

# Изпращаме на модела тестови данни, като използваме модула Predict
probe = dspy.Predict("question -> answer")
try:
    result = probe(question="Who are you", max_tokens=1)
    print("Model supported by provider, test answer:\n", result.answer)
except Exception as e:
    print("Model not supported by provider:\n", e)

def get_weather(city: str, units: str = "metric") -> str:
    url = "https://api.openweathermap.org/data/2.5/weather"
    params = {
        "q": city,
        "appid": OPENWEATHER_API_KEY,
        "units": units,
        "lang": "en"
    }
    response = requests.get(url, params=params)
    if response.status_code == 200:
        data: Dict = response.json()
        temp = data["main"]["temp"]
        description = data["weather"][0]["description"]
        return f"Weather in {city}: {temp}°C, {description}"
    else:
        return f"Error fetching weather: {response.text}"


# Добавяме формат на общуване въпрос -> отговор,
# инструмента get_weather и модула ReAct
# Фреймуъркът сам ще генерира промпта в съответствие
# с това кой модул и с какви параметри използваме
agent = dspy.ReAct(
    "question -> answer",
    tools=[get_weather]
)

result = agent(question="Check weather in London")

print("Answer:\n", result.answer)

# Можем да разгледаме историята на общуването с LLM,
# кои tools са били извикани и какъв prompt е бил генериран
history = dspy.inspect_history(n=10)

import requests
from typing import Dict
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain import hub
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.callbacks.base import BaseCallbackHandler

OPENWEATHER_API_KEY = "..."

# Callback за логване
class CustomLogger(BaseCallbackHandler):
    def on_tool_start(self, serialized, input_str, **kwargs):
        print(f"\nTOOL START:\n{serialized.get('name')}\nINPUT:\n{input_str}\n")

    def on_tool_end(self, output, **kwargs):
        print(f"\nTOOL END:\n{output}\n")

    def on_llm_start(self, serialized, prompts, **kwargs):
        print(f"\nLLM START:\n{serialized.get('name')}\nPROMPTS\n{prompts}\n")

    def on_llm_end(self, response, **kwargs):
        print(f"\nLLM END:\n{response}\n")

# Определяме инструмент (tool), който агентът може да извиква
@tool(description="Get current weather for a city.")
def get_weather(city: str) -> str:
    # Формираме заявка към OpenWeather API
    url = "https://api.openweathermap.org/data/2.5/weather"
    params = {"q": city, "appid": OPENWEATHER_API_KEY, "units": "metric", "lang": "en"}
    r = requests.get(url, params=params, timeout=15)
    if not r.ok:
        # Ако има грешка – връщаме текста на грешката
        return f"error: {r.status_code} {r.text}"
    # Парсваме JSON отговора
    data: Dict = r.json()
    # Връщаме кратък низ с времето
    return f"Weather in {data.get('name', city)}: {data['main']['temp']}°C, {data['weather'][0]['description']}"

# Свързваме LLM през LM Studio (локален сървър, съвместим с OpenAI API)
llm = ChatOpenAI(
    model="google/gemma-3n-e4b",
    openai_api_base="http://localhost:1234/v1",
    openai_api_key="lm-studio",
    temperature=0,
)

# Зареждаме готов шаблон на промпт ReAct от LangChain Hub
prompt = hub.pull("hwchase17/react")

# Създаваме агент в стил ReAct (моделът мисли + извиква инструменти)
agent = create_react_agent(
  llm,
  tools=[get_weather],
  prompt=prompt
)

# Обгръщаме агента в Executor, за да стартираме и управляваме работата му
executor = AgentExecutor(
    agent=agent,
    tools=[get_weather],
    verbose=True
)

# Стартираме агента с въпрос за времето в Лондон
result = executor.invoke(
    {"input": "Check weather in London"},
    config={"callbacks": [CustomLogger()]}
)

ext {
    ktor_version   = '3.2.3'
    lc_4_j_version   = '1.3.0'
    logback_version = '1.5.13'
    kotlinx_serialization_json_version = "1.8.1"
}

dependencies {
    implementation "io.ktor:ktor-client-core:$ktor_version"
    implementation "io.ktor:ktor-client-cio:$ktor_version"
    implementation "io.ktor:ktor-client-content-negotiation:$ktor_version"
    implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
    implementation "io.ktor:ktor-client-logging:$ktor_version"
    implementation "org.jetbrains.kotlinx:kotlinx-serialization-json:$kotlinx_serialization_json_version"
    implementation "ch.qos.logback:logback-classic:$logback_version"
    implementation "dev.langchain4j:langchain4j:$lc_4_j_version"
    implementation "dev.langchain4j:langchain4j-open-ai:$lc_4_j_version"
    implementation("dev.langchain4j:langchain4j-http-client-jdk:$lc_4_j_version")
}

import dev.langchain4j.http.client.jdk.JdkHttpClientBuilder
import dev.langchain4j.model.openai.OpenAiChatModel
import dev.langchain4j.service.AiServices
import kotlinx.coroutines.runBlocking
import java.net.http.HttpClient
import java.time.Duration

fun main() = runBlocking {

    val httpClientBuilder = HttpClient.newBuilder()
        .connectTimeout(Duration.ofSeconds(10))
        .version(HttpClient.Version.HTTP_1_1)

    val jdkClientBuilder = JdkHttpClientBuilder()
        .httpClientBuilder(httpClientBuilder)
        .connectTimeout(Duration.ofSeconds(10))
        .readTimeout(Duration.ofSeconds(90))

    val model = OpenAiChatModel.builder()
        .httpClientBuilder(jdkClientBuilder)
        .baseUrl(BASE_URL)
        .apiKey("lm-studio")
        .modelName(MODEL_ID)
        .temperature(0.2)
        .logRequests(true)
        .logResponses(true)
        .timeout(Duration.ofSeconds(90))
        .listeners(listOf(ChatLogger()))
        .build()

    val assistant = AiServices.builder(Assistant::class.java)
        .chatModel(model)
        .tools(WeatherTools)
        .build()

    val result = assistant.chat("Check weather in London")
    println(result)

}

import dev.langchain4j.service.UserMessage
import dev.langchain4j.service.V

interface Assistant {

    @UserMessage("{{input}}")
    fun chat(@V("input") message: String): String

}

import dev.langchain4j.agent.tool.Tool
import io.ktor.client.call.*
import io.ktor.client.request.*
import io.ktor.http.*
import kotlinx.coroutines.runBlocking

object WeatherTools {

    @Tool(
        name = "current_weather_city",
        value = ["Get current weather for a city"]
    )
    @JvmStatic
    fun currentWeatherByCity(city: String): String = runBlocking {
        client.use { httpClient ->
            val url = URLBuilder(OWM_BASE).apply {
                appendPathSegments("weather")
                parameters.append("q", city)
                parameters.append("appid", WEATHER_API_KEY)
                parameters.append("units", "metric")
                parameters.append("lang", "en")
            }.buildString()

            val response: WeatherResponse = httpClient.get(url).body()
            val name = response.name ?: city
            val temp = response.main?.temp
            val desc = response.weather.firstOrNull()?.description ?: "n/a"
            if (temp == null) {
                "Can't read temperature for $name"
            } else {
                "Weather in $name: $temp °C, $desc"
            }
        }
    }
}

import io.ktor.client.*
import io.ktor.client.engine.cio.*
import io.ktor.client.plugins.*
import io.ktor.client.plugins.contentnegotiation.*
import io.ktor.serialization.kotlinx.json.*
import kotlinx.serialization.json.Json

const val BASE_URL: String = "http://localhost:1234/v1"

const val OWM_BASE = "https://api.openweathermap.org/data/2.5"

const val WEATHER_API_KEY = "..."

const val MODEL_ID = "google/gemma-3n-e4b"

val kotlinxJsonConfig: Json = Json {
    ignoreUnknownKeys = true
    prettyPrint = false
}

val client: HttpClient = HttpClient(CIO) {
    install(ContentNegotiation) {
        json(kotlinxJsonConfig)
    }
    install(HttpTimeout) {
        requestTimeoutMillis = 60_000
        connectTimeoutMillis = 10_000
        socketTimeoutMillis = 60_000
    }
}

import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable

@Serializable
data class WeatherResponse(
    @SerialName("coord") val coord: Coord? = null,
    @SerialName("weather") val weather: List<WeatherItem> = emptyList(),
    @SerialName("base") val base: String? = null,
    @SerialName("main") val main: MainBlock? = null,
    @SerialName("visibility") val visibility: Int? = null,
    @SerialName("wind") val wind: Wind? = null,
    @SerialName("clouds") val clouds: Clouds? = null,
    @SerialName("dt") val dt: Long? = null,
    @SerialName("sys") val sys: Sys? = null,
    @SerialName("timezone") val timezone: Int? = null,
    @SerialName("id") val id: Long? = null,
    @SerialName("name") val name: String? = null,
    @SerialName("cod") val cod: Int? = null
)

@Serializable
data class Coord(
    @SerialName("lon") val lon: Double? = null,
    @SerialName("lat") val lat: Double? = null
)

@Serializable
data class WeatherItem(
    @SerialName("id") val id: Int? = null,
    @SerialName("main") val main: String? = null,
    @SerialName("description") val description: String? = null,
    @SerialName("icon") val icon: String? = null
)

@Serializable
data class MainBlock(
    @SerialName("temp") val temp: Double? = null,
    @SerialName("feels_like") val feelsLike: Double? = null,
    @SerialName("temp_min") val tempMin: Double? = null,
    @SerialName("temp_max") val tempMax: Double? = null,
    @SerialName("pressure") val pressure: Int? = null,
    @SerialName("humidity") val humidity: Int? = null
)

@Serializable
data class Wind(
    @SerialName("speed") val speed: Double? = null,
    @SerialName("deg") val deg: Int? = null,
    @SerialName("gust") val gust: Double? = null
)

@Serializable
data class Clouds(
    @SerialName("all") val all: Int? = null
)

@Serializable
data class Sys(
    @SerialName("type") val type: Int? = null,
    @SerialName("id") val id: Int? = null,
    @SerialName("country") val country: String? = null,
    @SerialName("sunrise") val sunrise: Long? = null,
    @SerialName("sunset") val sunset: Long? = null
)

import dev.langchain4j.model.chat.listener.ChatModelErrorContext
import dev.langchain4j.model.chat.listener.ChatModelListener
import dev.langchain4j.model.chat.listener.ChatModelRequestContext
import dev.langchain4j.model.chat.listener.ChatModelResponseContext

class ChatLogger : ChatModelListener {

    override fun onRequest(request: ChatModelRequestContext) {
        println("ON_LLM_START:")
        println(request)
    }

    override fun onResponse(response: ChatModelResponseContext) {
        println("ON_LLM_END:")
        println(response)
    }

    override fun onError(context: ChatModelErrorContext) {
        println("ON_LLM_ERROR:")
        println(context.error().message)
    }

}

localagentsRag = "0.2.0"
localagentsFc = "0.1.0"
tasksGenai = "0.10.25"
tasksText = "0.10.26.1"
tasksVision = "0.10.26.1"
tensorflowLite = "2.17.0"
kotlinxCoroutinesGuava = "1.10.2"

tasks-genai = { module = "com.google.mediapipe:tasks-genai", version.ref = "tasksGenai" }
tasks-text = { module = "com.google.mediapipe:tasks-text", version.ref = "tasksText" }
tasks-vision = { module = "com.google.mediapipe:tasks-vision", version.ref = "tasksVision" }
tensorflow-lite = { module = "org.tensorflow:tensorflow-lite", version.ref = "tensorflowLite" }
localagents-rag = { module = "com.google.ai.edge.localagents:localagents-rag", version.ref = "localagentsRag" }
localagents-fc = { module = "com.google.ai.edge.localagents:localagents-fc", version.ref = "localagentsFc" }
kotlinx-coroutines-guava = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-guava", version.ref = "kotlinxCoroutinesGuava" }

import com.google.ai.edge.localagents.rag.chunking.TextChunker
import com.google.ai.edge.localagents.rag.memory.DefaultSemanticTextMemory
import com.google.ai.edge.localagents.rag.memory.SqliteVectorStore
import com.google.ai.edge.localagents.rag.models.Embedder
import com.google.ai.edge.localagents.rag.prompt.PromptBuilder

interface MediaPipeEngineCommon {

    var chunker: TextChunker
    var embedder: Embedder<String>
    var vectorStore: SqliteVectorStore
    var promptBuilder: PromptBuilder
    var semanticMemory: DefaultSemanticTextMemory

    fun init(
        geckoModelPath: String, // Gecko_256_quant.tflite
        tokenizerModelPath: String, // sentencepiece.model
        useGpuForEmbeddings: Boolean = true,
    )

    fun saveTextToVectorStore(
        text: String,
        chunkOverlap: Int = 20,
        chunkTokenSize: Int = 128,
        chunkMaxSymbolsSize: Int = 1000,
        chunkBySentences: Boolean = false,
    ): String?

    fun readEmbeddingVectors(): List<VectorStoreEntity>

    suspend fun readEmbeddingVectors(
        query: String,
        topK: Int,
        minSimilarityScore: Float,
    ): List<VectorStoreEntity>

    fun makeSQLRequest(query: String): Boolean

}

import android.app.Application
import android.database.Cursor
import android.database.sqlite.SQLiteDatabase
import com.google.ai.edge.localagents.rag.chunking.TextChunker
import com.google.ai.edge.localagents.rag.memory.DefaultSemanticTextMemory
import com.google.ai.edge.localagents.rag.memory.SqliteVectorStore
import com.google.ai.edge.localagents.rag.memory.VectorStoreRecord
import com.google.ai.edge.localagents.rag.models.EmbedData
import com.google.ai.edge.localagents.rag.models.Embedder
import com.google.ai.edge.localagents.rag.models.EmbeddingRequest
import com.google.ai.edge.localagents.rag.models.GeckoEmbeddingModel
import com.google.ai.edge.localagents.rag.prompt.PromptBuilder
import com.google.common.collect.ImmutableList
import com.romankryvolapov.offlineailauncher.common.extensions.toDurationString
import com.romankryvolapov.offlineailauncher.common.models.common.LogUtil.logDebug
import com.romankryvolapov.offlineailauncher.common.models.common.LogUtil.logError
import kotlinx.coroutines.guava.await
import java.io.File
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.util.Optional

class MediaPipeEngineCommonImpl(
    private val application: Application
) : MediaPipeEngineCommon {

    companion object {
        private const val TAG = "CommonComponentsTag"
        private const val GECKO_EMBEDDING_MODEL_DIMENSION = 768
        private const val PROMPT_TEMPLATE: String =
            "You are an assistant for question-answering tasks. Here are the things I want to remember: {0} Use the things I want to remember, answer the following question the user has: {1}"

    }

    override lateinit var chunker: TextChunker
    override lateinit var embedder: Embedder<String>
    override lateinit var vectorStore: SqliteVectorStore
    override lateinit var promptBuilder: PromptBuilder
    override lateinit var semanticMemory: DefaultSemanticTextMemory

    override fun init(
        geckoModelPath: String,
        tokenizerModelPath: String,
        useGpuForEmbeddings: Boolean,
    ) {
        logDebug("init", TAG)
        chunker = TextChunker()

        // към embedder добавяме пътя до модела Gecko-110m-en
        // използвам версията Gecko_256_quant.tflite, тук 256 е максималният размер на текстовия вход
        // тази версия е оптимална от гледна точка на размера на частите текст и на бързодействието
        // важно – по-нататък в кода подаваме на какви части да се разбива текстът, това зависи от параметрите на модела
        embedder = GeckoEmbeddingModel(
            geckoModelPath,
            Optional.of(tokenizerModelPath),
            useGpuForEmbeddings,
        )

        // тук ще използвам вече готова SQLite база в приложението
        // в нея просто ще се добави още една таблица rag_vector_store
        // с колони text за текста и embeddings за вектора
        val database = File(application.getDatabasePath("database").absolutePath)
        if (!database.exists()) {
            logError("startEngine database not exists", TAG)
        }

        // Може да се направи и собствена реализация на базата данни, която наследява
        // интерфейса VectorStore<String>, но методът getNearestRecords трябва да бъде
        // реализиран правилно и да работи бързо, той търси най-близките вектори
        vectorStore = SqliteVectorStore(
            GECKO_EMBEDDING_MODEL_DIMENSION,
            database.absolutePath
        )
        semanticMemory = DefaultSemanticTextMemory(
            vectorStore,
            embedder
        )
        promptBuilder = PromptBuilder(
            PROMPT_TEMPLATE
        )
        logDebug("init ready", TAG)
    }

    override fun saveTextToVectorStore(
        text: String,

        // колко назад в текста да се връщаме преди фрагмента, тук 20
        chunkOverlap: Int,

        // Обърнете внимание, че размерът уж е в токени,
        // но се използва за chunker и може да не съответства
        // на размера в токени за embedder, тук chunkTokenSize е 128
        chunkTokenSize: Int,

        // при разбиване с помощта на chunkBySentences
        // размерът на изреченията може да е голям
        // ако надхвърли възможностите на embedder модела, той ще върне грешка
        // за целта се използва отрязване до максималния размер
        // тук той е 1000 символа
        chunkMaxSymbolsSize: Int,

        // използване на метода за разбиване по изречения
        chunkBySentences: Boolean,
    ): String? {
        logDebug("saveTextToVectorStore text length: ${text.length}", TAG)
        // таймер, за да разберем колко бързо работи
        val start = System.currentTimeMillis()
        val chunks: List<String> = if (chunkBySentences)
            chunker.chunkBySentences(
                text,
                chunkTokenSize,
            ).filter {
                it.isNotBlank()
            }.map { chunk ->
                if (chunk.length > chunkMaxSymbolsSize) {
                    logError("saveTextToVectorStore crop chunk", TAG)
                    chunk.substring(0, chunkMaxSymbolsSize)
                } else {
                    chunk
                }
            }
        else
            chunker.chunk(
                text,
                chunkTokenSize,
                chunkOverlap
            ).filter {
                it.isNotBlank()
            }.map { chunk ->
                if (chunk.length > chunkMaxSymbolsSize) {
                    logError("saveTextToVectorStore crop chunk", TAG)
                    chunk.substring(0, chunkMaxSymbolsSize)
                } else {
                    chunk
                }
            }
        val end = System.currentTimeMillis()
        val delta = end - start
        logDebug("saveTextToVectorStore chunks delta: ${delta.toDurationString()} size: ${chunks.size}", TAG)
        chunks.forEach {
            logDebug("length: ${it.length}", TAG)
        }
        if (chunks.isEmpty()) {
            logError("saveTextToVectorStore chunks.isEmpty()", TAG)
            return "Chunks is empty"
        }
        return try {
            // векторът се генерира вътре в semanticMemory
            val result: Boolean? = semanticMemory.recordBatchedMemoryItems(
                ImmutableList.copyOf(chunks)
            )?.get()
            val end = System.currentTimeMillis()
            val delta = end - start
            logDebug("saveTextToVectorStore ready delta: ${delta.toDurationString()} result: $result", TAG)
            null
        } catch (t: Throwable) {
            logError("saveTextToVectorStore failed: ${t.message}", t, TAG)
            t.message
        }
    }

    // търсене по заявката query, ще намери всички части от текста, подобни на заявката
    override suspend fun readEmbeddingVectors(
        query: String,

        // брой резултати от заявката към базата
        topK: Int,

        // доколко векторът на заявката query трябва да е подобен на запис в базата
        // 0.0 = търси във всички записи, сортирай по най-подобните
        // 1.0 = само идеално съвпадение
        // аз използвам стойности 0.6 - 0.8
        minSimilarityScore: Float,
    ): List<VectorStoreEntity> {
        logDebug("readEmbeddingVectors query: $query", TAG)
        val queryEmbedData: EmbedData<String> = EmbedData.create(
            query,
            EmbedData.TaskType.RETRIEVAL_QUERY
        )
        val embeddingRequest: EmbeddingRequest<String> = EmbeddingRequest
            .create(
                listOf(queryEmbedData)
            )
        val vector: ImmutableList<Float> = try {
            embedder.getEmbeddings(embeddingRequest).await()
        } catch (t: Throwable) {
            logError("readEmbeddingVectors: embedding failed: ${t.message}", t, TAG)
            return emptyList()
        }
        logDebug("searchDocsInternal vector size: ${vector.size}", TAG)
        if (vector.isEmpty()) {
            logError("readEmbeddingVectors vector.isEmpty()", TAG)
            return emptyList()
        }
        val hits: ImmutableList<VectorStoreRecord<String>> = try {
            vectorStore.getNearestRecords(
                vector,
                topK,
                minSimilarityScore
            )
        } catch (t: Throwable) {
            logError("readEmbeddingVectors: vector search failed: ${t.message}", t, TAG)
            return emptyList()
        }
        if (hits.isEmpty()) {
            logError("readEmbeddingVectors hits.isEmpty()", TAG)
            return emptyList()
        }
        val result = hits.map {
            VectorStoreEntity(
                id = null,
                text = it.data,
                embedding = it.embeddings
            )
        }
        logDebug("readEmbeddingVectors\nsize: ${result.size}\nresult: $result", TAG)
        return result
    }

    // просто извежда всички записи в базата
    override fun readEmbeddingVectors(): List<VectorStoreEntity> {
        logDebug("readEmbeddingPreview", TAG)
        var cursor: Cursor? = null
        var database: SQLiteDatabase? = null
        return try {
            val databaseFile = File(application.getDatabasePath("database").absolutePath)
            database = SQLiteDatabase.openDatabase(
                databaseFile.absolutePath,
                null,
                SQLiteDatabase.OPEN_READONLY
            )
            cursor = database.rawQuery("SELECT ROWID, text, embeddings FROM rag_vector_store", null)
            val result = mutableListOf<VectorStoreEntity>()
            while (cursor.moveToNext()) {
                val rowId = cursor.getLong(0)
                val text = cursor.getString(1)
                val blob = cursor.getBlob(2)
                val buffer = ByteBuffer.wrap(blob).order(ByteOrder.LITTLE_ENDIAN)
                val floats = mutableListOf<Float>()
                while (buffer.hasRemaining()) {
                    floats.add(buffer.float)
                }
                result.add(
                    VectorStoreEntity(
                        id = rowId,
                        text = text,
                        embedding = floats
                    )
                )
            }
            logDebug("readEmbeddingPreview\nsize: ${result.size}\nresult: $result", TAG)
            result
        } catch (t: Throwable) {
            logError("readEmbeddingPreview failed: ${t.message}", t, TAG)
            emptyList()
        } finally {
            cursor?.close()
            database?.close()
        }
    }

    // може да се напише собствена заявка и тя ще се изпълни,
    // например "DELETE FROM rag_vector_store"
    override fun makeSQLRequest(query: String): Boolean {
        logDebug("makeSQLRequest query: $query", TAG)
        var cursor: Cursor? = null
        var database: SQLiteDatabase? = null
        return try {
            val databaseFile = File(application.getDatabasePath("database").absolutePath)
            database = SQLiteDatabase.openDatabase(
                databaseFile.absolutePath,
                null,
                SQLiteDatabase.OPEN_READWRITE
            )
            cursor = database.rawQuery(query, null)
            val result = cursor.moveToFirst()
            logDebug("makeSQLRequest result: $result", TAG)
            result
        } catch (t: Throwable) {
            logError("makeSQLRequest failed: ${t.message}", t, TAG)
            false
        } finally {
            cursor?.close()
            database?.close()
        }
    }

}