feat: Add metadata filter argument to doc search

agents-api/agents_api/autogen/Docs.py

@@ -6,7 +6,7 @@
 from typing import Annotated, Any, Literal
 from uuid import UUID
 
-from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, StrictBool
 
 
 class BaseDocSearchRequest(BaseModel):
@@ -18,6 +18,7 @@ class BaseDocSearchRequest(BaseModel):
     """
     The language to be used for text-only search. Support for other languages coming soon.
     """
+    metadata_filter: dict[str, float | str | StrictBool | None]
 
 
 class CreateDocRequest(BaseModel):

agents-api/agents_api/models/docs/search_docs_by_embedding.py

@@ -1,5 +1,6 @@
 """This module contains functions for searching documents in the CozoDB based on embedding queries."""
 
+import json
 from typing import Any, Literal, TypeVar
 from uuid import UUID
 
@@ -51,6 +52,7 @@ def search_docs_by_embedding(
     ef: int = 50,
     mmr_lambda: float = 0.5,
     embedding_size: int = 1024,
+    metadata_filter: dict[str, Any] = {},
 ) -> tuple[list[str], dict]:
     """
     Searches for document snippets in CozoDB by embedding query.
@@ -62,11 +64,20 @@ def search_docs_by_embedding(
         k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3.
         confidence (float, optional): The confidence threshold for filtering results. Defaults to 0.8.
         mmr_lambda (float, optional): The lambda parameter for MMR. Defaults to 0.25.
+        embedding_size (int): Embedding vector length
+        metadata_filter (dict[str, Any]): Dictionary to filter agents based on metadata.
     """
 
     assert len(query_embedding) == embedding_size
     assert sum(query_embedding)
 
+    metadata_filter_str = ", ".join(
+        [
+            f"metadata->{json.dumps(k)} == {json.dumps(v)}"
+            for k, v in metadata_filter.items()
+        ]
+    )
+
     owners: list[list[str]] = [
         [owner_type, str(owner_id)] for owner_type, owner_id in owners
     ]
@@ -92,7 +103,8 @@ def search_docs_by_embedding(
                 owner_type,
                 owner_id,
                 doc_id
-            }}
+            }},
+            {metadata_filter_str}
 
         intersnippet_distance[
             doc_id,

agents-api/agents_api/models/docs/search_docs_by_text.py

@@ -1,5 +1,6 @@
 """This module contains functions for searching documents in the CozoDB based on embedding queries."""
 
+import json
 import re
 from typing import Any, Literal, TypeVar
 from uuid import UUID
@@ -49,6 +50,7 @@ def search_docs_by_text(
     owners: list[tuple[Literal["user", "agent"], UUID]],
     query: str,
     k: int = 3,
+    metadata_filter: dict[str, Any] = {},
 ) -> tuple[list[str], dict]:
     """
     Searches for document snippets in CozoDB by embedding query.
@@ -57,7 +59,14 @@ def search_docs_by_text(
         owners (list[tuple[Literal["user", "agent"], UUID]]): The type of the owner of the documents.
         query (str): The query string.
         k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3.
+        metadata_filter (dict[str, Any]): Dictionary to filter agents based on metadata.
     """
+    metadata_filter_str = ", ".join(
+        [
+            f"metadata->{json.dumps(k)} == {json.dumps(v)}"
+            for k, v in metadata_filter.items()
+        ]
+    )
 
     owners: list[list[str]] = [
         [owner_type, str(owner_id)] for owner_type, owner_id in owners
@@ -84,7 +93,8 @@ def search_docs_by_text(
                 owner_type,
                 owner_id,
                 doc_id
-            }}
+            }},
+            {metadata_filter_str}
 
         search_result[
             doc_id,

agents-api/agents_api/routers/docs/search_docs.py

@@ -25,21 +25,28 @@ def get_search_fn_and_params(
     search_fn, params = None, None
 
     match search_params:
-        case TextOnlyDocSearchRequest(text=query, limit=k):
+        case TextOnlyDocSearchRequest(
+            text=query, limit=k, metadata_filter=metadata_filter
+        ):
             search_fn = search_docs_by_text
             params = dict(
                 query=query,
                 k=k,
+                metadata_filter=metadata_filter,
             )
 
         case VectorDocSearchRequest(
-            vector=query_embedding, limit=k, confidence=confidence
+            vector=query_embedding,
+            limit=k,
+            confidence=confidence,
+            metadata_filter=metadata_filter,
         ):
             search_fn = search_docs_by_embedding
             params = dict(
                 query_embedding=query_embedding,
                 k=k,
                 confidence=confidence,
+                metadata_filter=metadata_filter,
             )
 
         case HybridDocSearchRequest(
@@ -48,6 +55,7 @@ def get_search_fn_and_params(
             limit=k,
             confidence=confidence,
             alpha=alpha,
+            metadata_filter=metadata_filter,
         ):
             search_fn = search_docs_hybrid
             params = dict(
@@ -56,6 +64,7 @@ def get_search_fn_and_params(
                 k=k,
                 embed_search_options=dict(confidence=confidence),
                 alpha=alpha,
+                metadata_filter=metadata_filter,
             )
 
     return search_fn, params

typespec/docs/models.tsp

@@ -82,6 +82,7 @@ model BaseDocSearchRequest {
 
     /** The language to be used for text-only search. Support for other languages coming soon. */
     lang: "en-US" = "en-US";
+    metadata_filter: MetadataFilter,
 }
 
 model VectorDocSearchRequest extends BaseDocSearchRequest {

typespec/tsp-output/@typespec/openapi3/openapi-0.4.0.yaml

@@ -2425,6 +2425,7 @@ components:
       required:
         - limit
         - lang
+        - metadata_filter
       properties:
         limit:
           type: integer
@@ -2438,6 +2439,14 @@ components:
             - en-US
           description: The language to be used for text-only search. Support for other languages coming soon.
           default: en-US
+        metadata_filter:
+          type: object
+          additionalProperties:
+            anyOf:
+              - type: number
+              - type: string
+              - type: boolean
+            nullable: true
     Docs.CreateDocRequest:
       type: object
       required:

typespec/tsp-output/@typespec/openapi3/openapi-1.0.0.yaml

@@ -2425,6 +2425,7 @@ components:
       required:
         - limit
         - lang
+        - metadata_filter
       properties:
         limit:
           type: integer
@@ -2438,6 +2439,14 @@ components:
             - en-US
           description: The language to be used for text-only search. Support for other languages coming soon.
           default: en-US
+        metadata_filter:
+          type: object
+          additionalProperties:
+            anyOf:
+              - type: number
+              - type: string
+              - type: boolean
+            nullable: true
     Docs.CreateDocRequest:
       type: object
       required: