RAG (Retrieval-Augmented Generation) framework for the Embabel Agent platform providing content ingestion, chunking, hierarchical navigation, and semantic search capabilities
—
This guide covers semantic similarity search using vector embeddings, including filtering, result expansion, and performance optimization.
Vector search enables semantic similarity matching by:
Perform a simple semantic search to find relevant chunks.
import com.embabel.agent.rag.service.VectorSearch
import com.embabel.agent.rag.model.*
val searchOps: VectorSearch = // implementation
// Simple vector search
val results = searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = "machine learning algorithms",
topK = 10,
similarityThreshold = 0.7
),
clazz = Chunk::class.java
)
// Process results
results.forEach { result ->
println("Score: ${"%.3f".format(result.score)}")
println("Content: ${result.content.text}")
println("Metadata: ${result.content.metadata}")
println()
}val request = TextSimilaritySearchRequest(
query = "authentication methods",
topK = 20, // Maximum results to return
similarityThreshold = 0.75 // Minimum similarity (0.0-1.0)
)
// Higher threshold = stricter matching
// Lower threshold = more lenient matching
// topK limits results after threshold filteringApply metadata and entity filters to narrow search scope.
import com.embabel.agent.rag.service.FilteringVectorSearch
import com.embabel.agent.rag.filter.*
val searchOps: FilteringVectorSearch = // implementation
// Search with metadata filter
val results = searchOps.vectorSearchWithFilter(
request = TextSimilaritySearchRequest(
query = "authentication setup",
topK = 10
),
clazz = Chunk::class.java,
metadataFilter = PropertyFilter.eq("category", "security")
.and(PropertyFilter.gte("version", 2.0)),
entityFilter = null
)import com.embabel.agent.rag.filter.PropertyFilter
// Equality
val eqFilter = PropertyFilter.eq("status", "active")
// Comparison
val gtFilter = PropertyFilter.gt("priority", 5)
val gteFilter = PropertyFilter.gte("version", 2.0)
val ltFilter = PropertyFilter.lt("age", 30)
val lteFilter = PropertyFilter.lte("score", 100)
// Not equal
val neFilter = PropertyFilter.ne("status", "archived")
// Contains (substring)
val containsFilter = PropertyFilter.contains("title", "auth")
// In list
val inFilter = PropertyFilter.`in`("team", "platform", "infrastructure", "security")
// Not in list
val ninFilter = PropertyFilter.nin("status", "deleted", "archived")// AND operator
val andFilter = PropertyFilter.eq("category", "docs")
.and(PropertyFilter.gte("version", 3.0))
.and(PropertyFilter.ne("status", "draft"))
// OR operator
val orFilter = PropertyFilter.eq("type", "tutorial")
.or(PropertyFilter.eq("type", "guide"))
// NOT operator
val notFilter = !PropertyFilter.eq("archived", true)
// Complex boolean logic
val complexFilter = PropertyFilter.and(
PropertyFilter.or(
PropertyFilter.eq("type", "documentation"),
PropertyFilter.eq("type", "tutorial")
),
PropertyFilter.gte("lastUpdated", System.currentTimeMillis() - 86400000),
!PropertyFilter.eq("status", "archived")
)import com.embabel.agent.rag.filter.EntityFilter
// Filter by entity labels
val entityFilter = EntityFilter.hasAnyLabel("Person", "Employee")
// Combined with metadata filter
val results = searchOps.vectorSearchWithFilter(
request = TextSimilaritySearchRequest(
query = "software engineers",
topK = 10
),
clazz = NamedEntityData::class.java,
metadataFilter = PropertyFilter.eq("department", "engineering"),
entityFilter = EntityFilter.hasAnyLabel("Employee", "Contractor")
)Search across different content types simultaneously.
import com.embabel.agent.rag.model.*
// Define search types
val searchTypes = listOf(
Chunk::class.java,
Fact::class.java,
NamedEntity::class.java
)
// Search each type
val allResults = searchTypes.flatMap { type ->
searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = "machine learning",
topK = 5
),
clazz = type
)
}.sortedByDescending { it.score }
// Get top results across all types
val topResults = allResults.take(10)results.forEach { result ->
when (val content = result.content) {
is Chunk -> {
println("Chunk: ${content.text.take(100)}...")
println(" Section: ${content.metadata["container_section_title"]}")
}
is Fact -> {
println("Fact: ${content.assertion}")
println(" Authority: ${content.authority}")
}
is NamedEntity -> {
println("Entity: ${content.name}")
println(" Description: ${content.description}")
}
}
}Get additional context by expanding search results.
import com.embabel.agent.rag.service.ResultExpander
val expander: ResultExpander = // implementation
val searchResults = searchOps.vectorSearch(...)
// Get first result
val topResult = searchResults.first()
// Expand to include surrounding chunks (sequence)
val sequenceContext = expander.expandResult(
id = topResult.content.id,
method = ResultExpander.Method.SEQUENCE,
elementsToAdd = 2 // 2 chunks before and 2 after
)
println("Expanded to ${sequenceContext.size} chunks")
sequenceContext.forEach { element ->
println("- ${element.id}")
}// Zoom out to containing section
val parentContent = expander.expandResult(
id = topResult.content.id,
method = ResultExpander.Method.ZOOM_OUT,
elementsToAdd = 1
)
println("Parent section contains ${parentContent.size} elements")fun searchWithContext(
query: String,
searchOps: VectorSearch,
expander: ResultExpander,
contextChunks: Int = 1
): List<List<ContentElement>> {
// Perform search
val results = searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = query,
topK = 5
),
clazz = Chunk::class.java
)
// Expand each result
return results.map { result ->
expander.expandResult(
id = result.content.id,
method = ResultExpander.Method.SEQUENCE,
elementsToAdd = contextChunks
)
}
}Look up specific items when you know their identifier.
import com.embabel.agent.rag.service.FinderOperations
val finder: FinderOperations = // implementation
// Find by ID and class
val chunk = finder.findById("chunk-123", Chunk::class.java)
if (chunk != null) {
println("Found chunk: ${chunk.text}")
}
// Find by ID and type string
val entity = finder.findById<NamedEntity>("person-456", "Person")
if (entity != null) {
println("Found entity: ${entity.name}")
}
// Check type support
if (finder.supportsType("Chunk")) {
println("Chunk type is supported")
}Group similar items together using cluster analysis.
import com.embabel.agent.rag.service.*
val clusterFinder: ClusterFinder = // implementation
// Find clusters with custom parameters
val clusters = clusterFinder.findClusters(
ClusterRetrievalRequest<NamedEntityData>()
.withSimilarityThreshold(0.8) // Higher = tighter clusters
.withTopK(15) // Max similar items per cluster
)
// Process clusters
clusters.forEach { cluster ->
println("Cluster anchor: ${cluster.anchor.name}")
println("Similar items (${cluster.similar.size}):")
cluster.similar.forEach { similar ->
println(" - ${similar.content.name} (score: ${similar.score})")
}
println()
}Customize how search results are displayed.
import com.embabel.agent.rag.service.*
// Use default formatter
val formatter = SimpleRetrievableResultsFormatter
val results = searchOps.vectorSearch(...)
val formatted = formatter.formatResults(
SimilarityResults.fromList(results)
)
println(formatted)// Create custom formatter
val markdownFormatter = RetrievableResultsFormatter { results ->
buildString {
appendLine("# Search Results")
appendLine()
appendLine("Found ${results.results.size} matches")
appendLine()
results.results.forEachIndexed { index, result ->
appendLine("## ${index + 1}. Score: ${"%.3f".format(result.score)}")
when (val content = result.content) {
is Chunk -> {
appendLine("**Type:** Chunk")
appendLine("**Section:** ${content.metadata["container_section_title"]}")
appendLine()
appendLine(content.text)
}
is NamedEntity -> {
appendLine("**Type:** Entity")
appendLine("**Name:** ${content.name}")
appendLine()
appendLine(content.description)
}
}
appendLine()
appendLine("---")
appendLine()
}
}
}
val formatted = markdownFormatter.formatResults(
SimilarityResults.fromList(results)
)When using filters, you may need to fetch more results before filtering.
import com.embabel.agent.rag.tools.TopKInflationStrategy
// Multiplier strategy: fetch 3x more results before filtering
val strategy = TopKInflationStrategy.multiplier(
multiplier = 3,
maxTopK = 1000
)
// Offset strategy: fetch fixed extra amount
val offsetStrategy = TopKInflationStrategy.offset(
offset = 50,
maxTopK = 1000
)
// Expected pass rate strategy
val passRateStrategy = TopKInflationStrategy.expectedPassRate(
percentage = 0.3, // Expect 30% to pass filter
maxTopK = 1000
)fun batchSearch(
queries: List<String>,
searchOps: VectorSearch
): Map<String, List<SimilarityResult<Chunk>>> {
return queries.associateWith { query ->
searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = query,
topK = 5
),
clazz = Chunk::class.java
)
}
}
// Use
val queries = listOf(
"authentication methods",
"error handling",
"database configuration"
)
val batchResults = batchSearch(queries, searchOps)
batchResults.forEach { (query, results) ->
println("Query: $query → ${results.size} results")
}import java.util.concurrent.ConcurrentHashMap
class CachedVectorSearch(
private val delegate: VectorSearch
) : VectorSearch by delegate {
private val cache = ConcurrentHashMap<String, List<SimilarityResult<*>>>()
override fun <T : Retrievable> vectorSearch(
request: TextSimilaritySearchRequest,
clazz: Class<T>
): List<SimilarityResult<T>> {
val cacheKey = "${request.query}:${clazz.simpleName}:${request.topK}"
@Suppress("UNCHECKED_CAST")
return cache.getOrPut(cacheKey) {
delegate.vectorSearch(request, clazz)
} as List<SimilarityResult<T>>
}
fun clearCache() {
cache.clear()
}
}fun adaptiveSearch(
query: String,
searchOps: VectorSearch,
minResults: Int = 5
): List<SimilarityResult<Chunk>> {
var threshold = 0.8
var results = emptyList<SimilarityResult<Chunk>>()
// Lower threshold until we get enough results
while (results.size < minResults && threshold > 0.5) {
results = searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = query,
topK = 20,
similarityThreshold = threshold
),
clazz = Chunk::class.java
)
if (results.size < minResults) {
threshold -= 0.1
}
}
return results
}fun deduplicateResults(
results: List<SimilarityResult<Chunk>>,
similarityThreshold: Double = 0.95
): List<SimilarityResult<Chunk>> {
val deduplicated = mutableListOf<SimilarityResult<Chunk>>()
val seen = mutableSetOf<String>()
for (result in results) {
// Use normalized text as key
val normalized = result.content.text
.lowercase()
.replace(Regex("\\s+"), " ")
.trim()
if (!seen.contains(normalized)) {
deduplicated.add(result)
seen.add(normalized)
}
}
return deduplicated
}fun multiQuerySearch(
queries: List<String>,
searchOps: VectorSearch,
topK: Int = 5
): List<SimilarityResult<Chunk>> {
// Search with each query
val allResults = queries.flatMap { query ->
searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = query,
topK = topK
),
clazz = Chunk::class.java
)
}
// Deduplicate and re-rank
return allResults
.distinctBy { it.content.id }
.sortedByDescending { it.score }
.take(topK)
}
// Use
val results = multiQuerySearch(
queries = listOf(
"how to authenticate",
"authentication methods",
"login and authorization"
),
searchOps = searchOps,
topK = 10
)val results = searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = "very specific query",
topK = 10
),
clazz = Chunk::class.java
)
if (results.isEmpty()) {
println("No results found")
// Try with lower threshold
val relaxedResults = searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = "very specific query",
topK = 10,
similarityThreshold = 0.5
),
clazz = Chunk::class.java
)
if (relaxedResults.isEmpty()) {
println("Still no results with relaxed threshold")
// Try text search as fallback
}
}fun safeVectorSearch(
query: String,
searchOps: VectorSearch
): List<SimilarityResult<Chunk>> {
if (query.isBlank()) {
println("Warning: Empty query provided")
return emptyList()
}
return searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = query.trim(),
topK = 10
),
clazz = Chunk::class.java
)
}// Filter out low-quality results
fun filterQualityResults(
results: List<SimilarityResult<Chunk>>,
minScore: Double = 0.7,
minTextLength: Int = 50
): List<SimilarityResult<Chunk>> {
return results.filter { result ->
result.score >= minScore &&
result.content.text.length >= minTextLength
}
}fun <T : Retrievable> safeSearch(
query: String,
searchOps: VectorSearch,
clazz: Class<T>
): List<SimilarityResult<T>> {
// Check if type is supported (if searchOps implements TypeRetrievalOperations)
if (searchOps is TypeRetrievalOperations) {
val typeName = clazz.simpleName
if (!searchOps.supportsType(typeName)) {
println("Warning: Type $typeName not supported")
return emptyList()
}
}
return searchOps.vectorSearch(
request = TextSimilaritySearchRequest(
query = query,
topK = 10
),
clazz = clazz
)
}