Embeddings

Embeddings

Generate and compare vector representations of text and code that capture semantic meaning, enabling efficient similarity comparisons. Koog's embedding system is split into two modules.

Modules

Dependencies

dependencies {
    implementation("ai.koog:koog-agents:0.8.0")
    implementation("ai.koog:embeddings-base:0.8.0")
    implementation("ai.koog:embeddings-llm:0.8.0")
}

LLMEmbedder

The LLMEmbedder class is the central abstraction for generating embeddings. It wraps an LLM client and a model reference, exposing two key methods:

Embedding Providers

Ollama (Local)

Requires Ollama installed locally. Models must be pulled first via ollama pull <model-id>:

import ai.koog.embeddings.llm.LLMEmbedder
import ai.koog.clients.ollama.OllamaClient
import ai.koog.agents.ext.llm.OllamaModels

val client = OllamaClient()
val embedder = LLMEmbedder(client, OllamaModels.Embeddings.NOMIC_EMBED_TEXT)
val embedding = embedder.embed("This is the text to embed")

Available Ollama Models

OpenAI

Uses OpenAILLMClient with an API key:

import ai.koog.clients.openai.OpenAILLMClient
import ai.koog.agents.ext.llm.OpenAIModels

val client = OpenAILLMClient(System.getenv("OPENAI_API_KEY"))
val embedder = LLMEmbedder(client, OpenAIModels.Embeddings.TextEmbeddingAda002)
val embedding = embedder.embed("This is the text to embed")

AWS Bedrock

Uses BedrockLLMClient with AWS credentials:

import ai.koog.clients.bedrock.BedrockLLMClient
import ai.koog.clients.bedrock.BedrockClientSettings
import ai.koog.agents.ext.llm.BedrockModels

val client = BedrockLLMClient(
    identityProvider = StaticCredentialsProvider {
        StaticCredentialsProvider.Credentials(
            accessKeyId = System.getenv("AWS_ACCESS_KEY_ID"),
            secretAccessKey = System.getenv("AWS_SECRET_ACCESS_KEY")
        )
    },
    settings = BedrockClientSettings()
)
val embedder = LLMEmbedder(client, BedrockModels.Embeddings.AmazonTitanEmbedText)
val embedding = embedder.embed("This is the text to embed")

Bedrock Embedding Models

Similarity Comparison

Code-to-Text Comparison

Embeddings bridge code and natural language. Compare a code snippet against descriptions to find semantic matches:

val client = OllamaClient()
val embedder = LLMEmbedder(client, OllamaModels.Embeddings.NOMIC_EMBED_TEXT)

val codeEmbedding = embedder.embed("""
    fun factorial(n: Int): Int = if (n <= 1) 1 else n * factorial(n - 1)
""".trimIndent())

val descriptionA = embedder.embed("A recursive function that calculates the product of all positive integers up to n")
val descriptionB = embedder.embed("A sorting algorithm that rearranges elements in ascending order")

val scoreA = embedder.diff(codeEmbedding, descriptionA) // Lower = more similar
val scoreB = embedder.diff(codeEmbedding, descriptionB) // Higher = less similar

// scoreA < scoreB → descriptionA is the better semantic match

Code-to-Code Comparison

Cross-language similarity detection:

val kotlinFib = embedder.embed("""
    fun fib(n: Int): Int = if (n <= 1) n else fib(n - 1) + fib(n - 2)
""".trimIndent())

val pythonFib = embedder.embed("""
    def fib(n): return n if n <= 1 else fib(n-1) + fib(n-2)
""".trimIndent())

val javaSort = embedder.embed("""
    void bubbleSort(int[] arr) { /* ... */ }
""".trimIndent())

val fibSimilarity = embedder.diff(kotlinFib, pythonFib)   // Low (similar algorithms)
val crossSimilarity = embedder.diff(kotlinFib, javaSort)   // High (different algorithms)

// fibSimilarity < crossSimilarity → same algorithm, different syntax are more similar

Model Selection Guide

Integration with RAG

Embeddings are the foundation for RAG (Retrieval-Augmented Generation). Use LLMEmbedder with a vector store to build a knowledge base:

// See rag-vector-store SKILL.md for full integration details
val embedder = LLMEmbedder(ollamaClient, OllamaModels.Embeddings.NOMIC_EMBED_TEXT)
val storage = EmbeddingStorage(
    documentEmbedder = JVMTextDocumentEmbedder(embedder),
    vectorStorageBackend = InMemoryVectorStorageBackend()
)

API References

Troubleshooting