Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -79,19 +79,15 @@ internal fun List<SimilarityResult<out Retrievable>>.withNeighbours(
}

@Suppress("UNCHECKED_CAST")
internal fun <T : Retrievable> VectorSearch.searchWithFilter(
internal fun <T : Retrievable> VectorSearch.vectorSearchWithFilter(
request: TextSimilaritySearchRequest,
clazz: Class<T>,
metadataFilter: PropertyFilter?,
entityFilter: EntityFilter?,
): List<SimilarityResult<T>> {
if (metadataFilter == null && entityFilter == null) {
return vectorSearch(request, clazz)
}
if (this is FilteringVectorSearch) {
return vectorSearchWithFilter(request, clazz, metadataFilter, entityFilter)
}
return PostFilteringSearch.search(
): List<SimilarityResult<T>> = when {
metadataFilter == null && entityFilter == null -> vectorSearch(request, clazz)
this is FilteringVectorSearch -> vectorSearchWithFilter(request, clazz, metadataFilter, entityFilter)
else -> PostFilteringSearch.search(
request,
metadataFilter,
entityFilter,
Expand Down Expand Up @@ -147,7 +143,7 @@ internal class VectorSearchTools @JvmOverloads constructor(

private fun searchForAllTypes(request: TextSimilaritySearchRequest): List<SimilarityResult<out Retrievable>> {
val allResults = searchFor.flatMap { clazz ->
vectorSearch.searchWithFilter(request, clazz, metadataFilter, entityFilter)
vectorSearch.vectorSearchWithFilter(request, clazz, metadataFilter, entityFilter)
}
return deduplicateByIdKeepingHighestScore(allResults)
}
Expand Down Expand Up @@ -380,6 +376,25 @@ internal class SectionReadingTools @JvmOverloads constructor(
}
}

@Suppress("UNCHECKED_CAST")
internal fun <T : Retrievable> TextSearch.textSearchWithFilter(
request: TextSimilaritySearchRequest,
clazz: Class<T>,
metadataFilter: PropertyFilter?,
entityFilter: EntityFilter?,
): List<SimilarityResult<T>> = when {
metadataFilter == null && entityFilter == null -> textSearch(request, clazz)
this is FilteringTextSearch -> textSearchWithFilter(request, clazz, metadataFilter, entityFilter)
else -> PostFilteringSearch.search(
request,
metadataFilter,
entityFilter,
TopKInflationStrategy.DEFAULT,
) { inflatedRequest ->
textSearch(inflatedRequest, clazz)
} as List<SimilarityResult<T>>
}

/**
* Tools to perform text search operations with the syntax supported by
* the backing [TextSearch] store.
Expand Down Expand Up @@ -480,36 +495,11 @@ internal class TextSearchTools @JvmOverloads constructor(

private fun searchForAllTypes(request: TextSimilaritySearchRequest): List<SimilarityResult<out Retrievable>> {
val allResults = searchFor.flatMap { clazz ->
searchWithFilter(request, clazz)
textSearch.textSearchWithFilter(request, clazz, metadataFilter, entityFilter)
}
return deduplicateByIdKeepingHighestScore(allResults)
}

@Suppress("UNCHECKED_CAST")
private fun <T : Retrievable> searchWithFilter(
request: TextSimilaritySearchRequest,
clazz: Class<T>,
): List<SimilarityResult<T>> {
if (metadataFilter == null && entityFilter == null) {
return textSearch.textSearch(request, clazz)
}

// If backend supports native filtering, use it
if (textSearch is FilteringTextSearch) {
return textSearch.textSearchWithFilter(request, clazz, metadataFilter, entityFilter)
}

// Fallback: inflate topK, search, post-filter, take topK
return PostFilteringSearch.search(
request,
metadataFilter,
entityFilter,
TopKInflationStrategy.DEFAULT
) { inflatedRequest ->
textSearch.textSearch(inflatedRequest, clazz)
} as List<SimilarityResult<T>>
}

companion object {
/**
* Compose the tool's top-level description from the store's
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package com.embabel.agent.rag.tools

import com.embabel.agent.api.tool.Tool
import com.embabel.agent.filter.PropertyFilter
import com.embabel.agent.rag.model.Chunk
import com.embabel.agent.rag.model.ContentElement
import com.embabel.agent.rag.model.NamedEntityData.Companion.ENTITY_LABEL
Expand Down Expand Up @@ -462,6 +463,28 @@ class ToolishRagTest {
}
}

@Nested
inner class SearchWithFilterExtensionsTests {

@Test
fun `filter helpers should be callable on CoreSearchOperations`() {
val searchOperations = mockk<CoreSearchOperations>()
val request = TextSimilaritySearchRequest("test query", 0.5, 5)
every {
searchOperations.vectorSearch(request, Chunk::class.java)
} returns emptyList()
every {
searchOperations.textSearch(request, Chunk::class.java)
} returns emptyList()

val vectorResults = searchOperations.vectorSearchWithFilter(request, Chunk::class.java, null, null)
val textResults = searchOperations.textSearchWithFilter(request, Chunk::class.java, null, null)

assertTrue(vectorResults.isEmpty())
assertTrue(textResults.isEmpty())
}
}

@Nested
inner class TextSearchToolsTests {

Expand Down Expand Up @@ -506,6 +529,78 @@ class ToolishRagTest {
assertEquals("0 results:", result)
}

@Test
fun `textSearch should use native filtering when supported`() {
val textSearch = mockk<FilteringTextSearch>()
val metadataFilter = PropertyFilter.eq("ownerId", "alice")
val chunk = createChunk("chunk1", "Alice's content")
every {
textSearch.textSearchWithFilter(
any<TextSimilaritySearchRequest>(),
Chunk::class.java,
metadataFilter,
null,
)
} returns listOf(SimpleSimilaritySearchResult(match = chunk, score = 0.9))
val tools = TextSearchTools(textSearch, metadataFilter = metadataFilter)

val result = tools.textSearch("test query", 5, 0.5)

verify(exactly = 1) {
textSearch.textSearchWithFilter(
match<TextSimilaritySearchRequest> { it.topK == 5 },
Chunk::class.java,
metadataFilter,
null,
)
}
assertTrue(result.contains("Alice's content"))
}

@Test
fun `textSearch should inflate topK and post-filter when native filtering is unavailable`() {
val textSearch = mockk<TextSearch>()
val metadataFilter = PropertyFilter.eq("ownerId", "alice")
val included = Chunk(
id = "included",
text = "Alice's content",
parentId = "parent",
metadata = mapOf("ownerId" to "alice"),
)
val excluded = Chunk(
id = "excluded",
text = "Bob's content",
parentId = "parent",
metadata = mapOf("ownerId" to "bob"),
)
val truncated = Chunk(
id = "truncated",
text = "Alice's lower-ranked content",
parentId = "parent",
metadata = mapOf("ownerId" to "alice"),
)
every {
textSearch.textSearch(any<TextSimilaritySearchRequest>(), Chunk::class.java)
} returns listOf(
SimpleSimilaritySearchResult(match = excluded, score = 0.9),
SimpleSimilaritySearchResult(match = included, score = 0.8),
SimpleSimilaritySearchResult(match = truncated, score = 0.7),
)
val tools = TextSearchTools(textSearch, metadataFilter = metadataFilter)

val result = tools.textSearch("test query", 1, 0.5)

verify(exactly = 1) {
textSearch.textSearch(
match<TextSimilaritySearchRequest> { it.topK == 3 },
Chunk::class.java,
)
}
assertTrue(result.contains("Alice's content"))
assertFalse(result.contains("Bob's content"))
assertFalse(result.contains("Alice's lower-ranked content"))
}

}

@Nested
Expand Down
Loading