feat(vector_io): Return structured metadata from vector store search

Signed-off-by: Peter Gustafsson <peter.gustafsson6@gmail.com>

Author: Peter Gustafsson

client-sdks/stainless/openapi.yml

@@ -13782,6 +13782,11 @@ components:
           - type: 'null'
           description: The search mode to use (e.g., 'vector', 'keyword').
           default: vector
+        include_metadata:
+          type: boolean
+          title: Include Metadata
+          description: Whether to include structured chunk metadata in results.
+          default: false
       type: object
       required:
       - query

docs/static/deprecated-ogx-spec.yaml

@@ -7447,6 +7447,11 @@ components:
           - type: 'null'
           description: The search mode to use (e.g., 'vector', 'keyword').
           default: vector
+        include_metadata:
+          type: boolean
+          title: Include Metadata
+          description: Whether to include structured chunk metadata in results.
+          default: false
       type: object
       required:
       - query

docs/static/experimental-ogx-spec.yaml

@@ -8817,6 +8817,11 @@ components:
           - type: 'null'
           description: The search mode to use (e.g., 'vector', 'keyword').
           default: vector
+        include_metadata:
+          type: boolean
+          title: Include Metadata
+          description: Whether to include structured chunk metadata in results.
+          default: false
       type: object
       required:
       - query

docs/static/ogx-spec.yaml

@@ -12298,6 +12298,11 @@ components:
           - type: 'null'
           description: The search mode to use (e.g., 'vector', 'keyword').
           default: vector
+        include_metadata:
+          type: boolean
+          title: Include Metadata
+          description: Whether to include structured chunk metadata in results.
+          default: false
       type: object
       required:
       - query

docs/static/stainless-ogx-spec.yaml

@@ -13782,6 +13782,11 @@ components:
           - type: 'null'
           description: The search mode to use (e.g., 'vector', 'keyword').
           default: vector
+        include_metadata:
+          type: boolean
+          title: Include Metadata
+          description: Whether to include structured chunk metadata in results.
+          default: false
       type: object
       required:
       - query

src/ogx/providers/inline/vector_io/faiss/faiss.py

@@ -92,6 +92,20 @@ def _list_op(mv: Any, fv: Any, *, negate: bool) -> bool:
     return (mv not in fv) if negate else (mv in fv)
 
 
+def _metadata_index_key(value: Any) -> Any:
+    if isinstance(value, dict):
+        items = ((str(key), _metadata_index_key(item)) for key, item in value.items())
+        return ("dict", tuple(sorted(items)))
+    if isinstance(value, list):
+        return ("list", tuple(_metadata_index_key(item) for item in value))
+
+    try:
+        hash(value)
+    except TypeError:
+        return (type(value).__qualname__, repr(value))
+    return value
+
+
 _COMPARISON_OPS: dict[str, Any] = {
     "eq": lambda mv, fv: mv == fv,
     "ne": lambda mv, fv: mv != fv,
@@ -157,7 +171,8 @@ async def initialize(self) -> None:
                 # Rebuild inverted metadata index from loaded chunks
                 for pos, chunk in self.chunk_by_index.items():
                     for key, val in chunk.metadata.items():
-                        self._meta_index.setdefault(key, {}).setdefault(val, set()).add(pos)
+                        index_key = _metadata_index_key(val)
+                        self._meta_index.setdefault(key, {}).setdefault(index_key, set()).add(pos)
             except Exception as e:
                 logger.debug("Failed to deserialize Faiss index", error=str(e), exc_info=True)
                 raise ValueError(
@@ -204,7 +219,8 @@ async def add_chunks(self, embedded_chunks: list[EmbeddedChunk]):
             faiss_pos = indexlen + i
             self.chunk_by_index[faiss_pos] = embedded_chunk
             for key, val in embedded_chunk.metadata.items():
-                self._meta_index.setdefault(key, {}).setdefault(val, set()).add(faiss_pos)
+                index_key = _metadata_index_key(val)
+                self._meta_index.setdefault(key, {}).setdefault(index_key, set()).add(faiss_pos)
 
         async with self.chunk_id_lock:
             self.index.add(embeddings)
@@ -269,16 +285,16 @@ def _resolve_filter_positions(self, filter_obj: Filter) -> set[int]:
         # ComparisonFilter
         key, value, op_type = filter_obj.key, filter_obj.value, filter_obj.type
         if op_type == "eq":
-            return self._meta_index.get(key, {}).get(value, set()).copy()
+            return self._meta_index.get(key, {}).get(_metadata_index_key(value), set()).copy()
         if op_type == "in":
             result: set[int] = set()
             for v in value:
-                result |= self._meta_index.get(key, {}).get(v, set())
+                result |= self._meta_index.get(key, {}).get(_metadata_index_key(v), set())
             return result
         if op_type == "nin":
             excluded: set[int] = set()
             for v in value:
-                excluded |= self._meta_index.get(key, {}).get(v, set())
+                excluded |= self._meta_index.get(key, {}).get(_metadata_index_key(v), set())
             all_positions = {pos for s in self._meta_index.get(key, {}).values() for pos in s}
             return all_positions - excluded
         # Range ops and ne: linear scan over chunk_by_index metadata

src/ogx/providers/utils/memory/openai_vector_store_mixin.py

@@ -1148,7 +1148,7 @@ async def openai_search_vector_store(
             data = []
             for embedded_chunk, score in zip(response.chunks, response.scores, strict=False):
                 chunk = embedded_chunk
-                content = self._chunk_to_vector_store_content(chunk)
+                content = self._chunk_to_vector_store_content(chunk, include_metadata=request.include_metadata)
 
                 response_data_item = VectorStoreSearchResponse(
                     file_id=chunk.metadata.get("document_id", ""),

src/ogx_api/vector_io/models.py

@@ -286,6 +286,13 @@ class VectorStoreSearchResponse(BaseModel):
     attributes: dict[str, str | float | bool] | None = None
     content: list[VectorStoreContent]
 
+    @field_validator("attributes", mode="before")
+    @classmethod
+    def _validate_attributes(cls, v: dict[str, Any] | None) -> dict[str, str | float | bool] | None:
+        if v is None:
+            return None
+        return _sanitize_vector_store_attributes(v)
+
 
 @json_schema_type
 class VectorStoreSearchResponsePage(BaseModel):
@@ -822,6 +829,9 @@ class OpenAISearchVectorStoreRequest(BaseModel):
     ranking_options: SearchRankingOptions | None = Field(default=None, description="Options for ranking results.")
     rewrite_query: bool = Field(default=False, description="Whether to rewrite the query for better results.")
     search_mode: str | None = Field(default="vector", description="The search mode to use (e.g., 'vector', 'keyword').")
+    include_metadata: bool = Field(
+        default=False, description="Whether to include structured chunk metadata in results."
+    )
 
 
 @json_schema_type

tests/unit/providers/vector_io/test_faiss.py

@@ -173,6 +173,37 @@ async def test_meta_index_populated_on_add_chunks(faiss_index, sample_chunks, sa
     assert faiss_index._meta_index["document_id"]["mock-doc-2"] == {1}
 
 
+async def test_meta_index_supports_structured_metadata(faiss_index, embedding_dimension):
+    metadata = {
+        "headings": ["Introduction", "Architecture"],
+        "page_numbers": [1, 2],
+        "details": {"language": "en"},
+    }
+    chunk = EmbeddedChunk(
+        content="Chunk text",
+        chunk_id="structured-metadata-chunk",
+        metadata=metadata,
+        chunk_metadata=ChunkMetadata(chunk_id="structured-metadata-chunk", document_id="document-1"),
+        embedding=np.random.rand(embedding_dimension).astype(np.float32).tolist(),
+        embedding_model="test-embedding-model",
+        embedding_dimension=embedding_dimension,
+    )
+
+    await faiss_index.add_chunks([chunk])
+
+    from ogx.providers.utils.vector_io.filters import ComparisonFilter
+
+    headings = faiss_index._resolve_filter_positions(
+        ComparisonFilter(key="headings", value=["Introduction", "Architecture"], type="eq")
+    )
+    details = faiss_index._resolve_filter_positions(
+        ComparisonFilter(key="details", value={"language": "en"}, type="eq")
+    )
+    assert headings == {0}
+    assert details == {0}
+    assert faiss_index.chunk_by_index[0].metadata == metadata
+
+
 async def test_meta_index_updated_on_delete_chunks(faiss_index, sample_chunks, sample_embeddings, embedding_dimension):
     """After deleting a chunk, its position should be removed and remaining positions shifted."""
     from ogx.providers.utils.memory.vector_store import ChunkForDeletion

tests/unit/providers/vector_io/test_vector_io_stores_config.py

@@ -12,8 +12,10 @@
 import pytest
 
 from ogx_api import (
+    ChunkMetadata,
     EmbeddedChunk,
     OpenAICreateVectorStoreRequestWithExtraBody,
+    OpenAISearchVectorStoreRequest,
     QueryChunksResponse,
     VectorStore,
 )
@@ -196,8 +198,6 @@ async def mock_query_chunks(*args, **kwargs):
 
     # Test that rewrite_query=True doesn't cause an error (it's ignored at mixin level)
     # The mixin should process the search request without attempting to rewrite the query
-    from ogx_api import OpenAISearchVectorStoreRequest
-
     request = OpenAISearchVectorStoreRequest(
         query="test query",
         max_num_results=5,
@@ -213,6 +213,52 @@ async def mock_query_chunks(*args, **kwargs):
     assert result.search_query == ["test query"]  # Original query preserved
 
 
+async def test_search_vector_store_includes_structured_metadata_on_request(vector_io_adapter):
+    vector_store_id = "test_store_metadata"
+    vector_io_adapter.openai_vector_stores[vector_store_id] = {
+        "id": vector_store_id,
+        "name": "Test Store",
+        "description": "",
+        "vector_store_id": "test_db",
+        "embedding_model": "test/embedding",
+    }
+    metadata = {
+        "document_id": "file-123",
+        "filename": "guide.pdf",
+        "headings": ["Introduction", "Architecture"],
+        "page_numbers": [1, 2],
+    }
+    chunk = EmbeddedChunk(
+        content="Chunk text",
+        chunk_id="chunk-123",
+        metadata=metadata,
+        chunk_metadata=ChunkMetadata(chunk_id="chunk-123", document_id="file-123"),
+        embedding=[0.1, 0.2],
+        embedding_model="test/embedding",
+        embedding_dimension=2,
+    )
+    vector_io_adapter.query_chunks = AsyncMock(return_value=QueryChunksResponse(chunks=[chunk], scores=[0.9]))
+
+    default_result = await vector_io_adapter.openai_search_vector_store(
+        vector_store_id=vector_store_id,
+        request=OpenAISearchVectorStoreRequest(query="architecture"),
+    )
+    included_result = await vector_io_adapter.openai_search_vector_store(
+        vector_store_id=vector_store_id,
+        request=OpenAISearchVectorStoreRequest(query="architecture", include_metadata=True),
+    )
+
+    assert default_result.data[0].content[0].metadata is None
+    assert default_result.data[0].attributes == {
+        "document_id": "file-123",
+        "filename": "guide.pdf",
+        "headings": "Introduction, Architecture",
+        "page_numbers": "1, 2",
+    }
+    assert included_result.data[0].content[0].metadata == metadata
+    assert included_result.data[0].content[0].chunk_metadata == chunk.chunk_metadata
+
+
 async def test_search_vector_store_propagates_backend_errors(vector_io_adapter):
     """Test that exceptions from the vector store backend propagate to the caller."""
     vector_store_id = "test_store_error"
@@ -229,8 +275,6 @@ async def mock_query_chunks(*args, **kwargs):
 
     vector_io_adapter.query_chunks = mock_query_chunks
 
-    from ogx_api import OpenAISearchVectorStoreRequest
-
     request = OpenAISearchVectorStoreRequest(query="test query", max_num_results=5)
     with pytest.raises(KeyError, match="chunk_content"):
         await vector_io_adapter.openai_search_vector_store(