Share search property filter dispatch

Signed-off-by: Arnab Nandy <arnab_nandy7@yahoo.com>

Author: arnabnandy7

embabel-agent-rag/embabel-agent-rag-core/src/main/kotlin/com/embabel/agent/rag/tools/CoreSearchTools.kt

@@ -78,6 +78,25 @@ internal fun List<SimilarityResult<out Retrievable>>.withNeighbours(
     return this + extra
 }
 
+@Suppress("UNCHECKED_CAST")
+internal fun <T : Retrievable> VectorSearch.vectorSearchWithFilter(
+    request: TextSimilaritySearchRequest,
+    clazz: Class<T>,
+    metadataFilter: PropertyFilter?,
+    entityFilter: EntityFilter?,
+): 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,
+        TopKInflationStrategy.DEFAULT,
+    ) { inflatedRequest ->
+        vectorSearch(inflatedRequest, clazz)
+    } as List<SimilarityResult<T>>
+}
+
 internal class VectorSearchTools @JvmOverloads constructor(
     private val vectorSearch: VectorSearch,
     private val searchFor: List<Class<out Retrievable>> = listOf(Chunk::class.java),
@@ -119,36 +138,10 @@ internal class VectorSearchTools @JvmOverloads constructor(
 
     private fun searchForAllTypes(request: TextSimilaritySearchRequest): List<SimilarityResult<out Retrievable>> {
         val allResults = searchFor.flatMap { clazz ->
-            searchWithFilter(request, clazz)
+            vectorSearch.vectorSearchWithFilter(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 vectorSearch.vectorSearch(request, clazz)
-        }
-
-        // If backend supports native filtering, use it
-        if (vectorSearch is FilteringVectorSearch) {
-            return vectorSearch.vectorSearchWithFilter(request, clazz, metadataFilter, entityFilter)
-        }
-
-        // Fallback: inflate topK, search, post-filter, take topK
-        // Note: PostFilteringSearch requires Datum constraint, so we cast
-        return PostFilteringSearch.search(
-            request,
-            metadataFilter,
-            entityFilter,
-            TopKInflationStrategy.DEFAULT
-        ) { inflatedRequest ->
-            vectorSearch.vectorSearch(inflatedRequest, clazz)
-        } as List<SimilarityResult<T>>
-    }
 }
 
 /**
@@ -378,6 +371,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.
@@ -478,36 +490,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

embabel-agent-rag/embabel-agent-rag-core/src/test/kotlin/com/embabel/agent/rag/tools/ToolishRagTest.kt

@@ -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
@@ -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 {
 
@@ -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