update vector_stores.search

docs/_static/llama-stack-spec.html

@@ -3864,7 +3864,7 @@
                         "content": {
                             "application/json": {
                                 "schema": {
-                                    "$ref": "#/components/schemas/VectorStoreSearchResponse"
+                                    "$ref": "#/components/schemas/VectorStoreSearchResponsePage"
                                 }
                             }
                         }
@@ -13132,7 +13132,70 @@
                 ],
                 "title": "OpenaiSearchVectorStoreRequest"
             },
+            "VectorStoreContent": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "text"
+                    },
+                    "text": {
+                        "type": "string"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type",
+                    "text"
+                ],
+                "title": "VectorStoreContent"
+            },
             "VectorStoreSearchResponse": {
+                "type": "object",
+                "properties": {
+                    "file_id": {
+                        "type": "string"
+                    },
+                    "filename": {
+                        "type": "string"
+                    },
+                    "score": {
+                        "type": "number"
+                    },
+                    "attributes": {
+                        "type": "object",
+                        "additionalProperties": {
+                            "oneOf": [
+                                {
+                                    "type": "string"
+                                },
+                                {
+                                    "type": "number"
+                                },
+                                {
+                                    "type": "boolean"
+                                }
+                            ]
+                        }
+                    },
+                    "content": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/VectorStoreContent"
+                        }
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "file_id",
+                    "filename",
+                    "score",
+                    "content"
+                ],
+                "title": "VectorStoreSearchResponse",
+                "description": "Response from searching a vector store."
+            },
+            "VectorStoreSearchResponsePage": {
                 "type": "object",
                 "properties": {
                     "object": {
@@ -13145,29 +13208,7 @@
                     "data": {
                         "type": "array",
                         "items": {
-                            "type": "object",
-                            "additionalProperties": {
-                                "oneOf": [
-                                    {
-                                        "type": "null"
-                                    },
-                                    {
-                                        "type": "boolean"
-                                    },
-                                    {
-                                        "type": "number"
-                                    },
-                                    {
-                                        "type": "string"
-                                    },
-                                    {
-                                        "type": "array"
-                                    },
-                                    {
-                                        "type": "object"
-                                    }
-                                ]
-                            }
+                            "$ref": "#/components/schemas/VectorStoreSearchResponse"
                         }
                     },
                     "has_more": {
@@ -13185,7 +13226,7 @@
                     "data",
                     "has_more"
                 ],
-                "title": "VectorStoreSearchResponse",
+                "title": "VectorStoreSearchResponsePage",
                 "description": "Response from searching a vector store."
             },
             "OpenaiUpdateVectorStoreRequest": {

docs/_static/llama-stack-spec.yaml

@@ -2734,7 +2734,7 @@ paths:
           content:
             application/json:
               schema:
-                $ref: '#/components/schemas/VectorStoreSearchResponse'
+                $ref: '#/components/schemas/VectorStoreSearchResponsePage'
         '400':
           $ref: '#/components/responses/BadRequest400'
         '429':
@@ -9190,7 +9190,48 @@ components:
       required:
         - query
       title: OpenaiSearchVectorStoreRequest
+    VectorStoreContent:
+      type: object
+      properties:
+        type:
+          type: string
+          const: text
+        text:
+          type: string
+      additionalProperties: false
+      required:
+        - type
+        - text
+      title: VectorStoreContent
     VectorStoreSearchResponse:
+      type: object
+      properties:
+        file_id:
+          type: string
+        filename:
+          type: string
+        score:
+          type: number
+        attributes:
+          type: object
+          additionalProperties:
+            oneOf:
+              - type: string
+              - type: number
+              - type: boolean
+        content:
+          type: array
+          items:
+            $ref: '#/components/schemas/VectorStoreContent'
+      additionalProperties: false
+      required:
+        - file_id
+        - filename
+        - score
+        - content
+      title: VectorStoreSearchResponse
+      description: Response from searching a vector store.
+    VectorStoreSearchResponsePage:
       type: object
       properties:
         object:
@@ -9201,15 +9242,7 @@ components:
         data:
           type: array
           items:
-            type: object
-            additionalProperties:
-              oneOf:
-                - type: 'null'
-                - type: boolean
-                - type: number
-                - type: string
-                - type: array
-                - type: object
+            $ref: '#/components/schemas/VectorStoreSearchResponse'
         has_more:
           type: boolean
           default: false
@@ -9221,7 +9254,7 @@ components:
         - search_query
         - data
         - has_more
-      title: VectorStoreSearchResponse
+      title: VectorStoreSearchResponsePage
       description: Response from searching a vector store.
     OpenaiUpdateVectorStoreRequest:
       type: object

llama_stack/apis/vector_io/vector_io.py

@@ -8,7 +8,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import Any, Protocol, runtime_checkable
+from typing import Any, Literal, Protocol, runtime_checkable
 
 from pydantic import BaseModel, Field
 
@@ -96,13 +96,30 @@ class VectorStoreSearchRequest(BaseModel):
     rewrite_query: bool = False
 
 
+@json_schema_type
+class VectorStoreContent(BaseModel):
+    type: Literal["text"]
+    text: str
+
+
 @json_schema_type
 class VectorStoreSearchResponse(BaseModel):
     """Response from searching a vector store."""
 
+    file_id: str
+    filename: str
+    score: float
+    attributes: dict[str, str | float | bool] | None = None
+    content: list[VectorStoreContent]
+
+
+@json_schema_type
+class VectorStoreSearchResponsePage(BaseModel):
+    """Response from searching a vector store."""
+
     object: str = "vector_store.search_results.page"
     search_query: str
-    data: list[dict[str, Any]]
+    data: list[VectorStoreSearchResponse]
     has_more: bool = False
     next_page: str | None = None
 
@@ -259,7 +276,7 @@ class VectorIO(Protocol):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         """Search for chunks in a vector store.
 
         Searches a vector store for relevant chunks based on a query and optional file attribute filters.

llama_stack/distribution/routers/vector_io.py

@@ -17,7 +17,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import RoutingTable
@@ -242,7 +242,7 @@ class VectorIORouter(VectorIO):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         logger.debug(f"VectorIORouter.openai_search_vector_store: {vector_store_id}")
         # Route based on vector store ID
         provider = self.routing_table.get_provider_impl(vector_store_id)

llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -21,7 +21,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
@@ -239,5 +239,5 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")

llama_stack/providers/remote/vector_io/milvus/milvus.py

@@ -23,7 +23,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
@@ -237,7 +237,7 @@ class MilvusVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
 
 

llama_stack/providers/remote/vector_io/qdrant/qdrant.py

@@ -21,7 +21,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
@@ -239,5 +239,5 @@ class QdrantVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -13,10 +13,12 @@ from typing import Any
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import (
     QueryChunksResponse,
+    VectorStoreContent,
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
     VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 
 logger = logging.getLogger(__name__)
@@ -85,7 +87,6 @@ class OpenAIVectorStoreMixin(ABC):
         provider_vector_db_id: str | None = None,
     ) -> VectorStoreObject:
         """Creates a vector store."""
-        print("IN OPENAI VECTOR STORE MIXIN, openai_create_vector_store")
         # store and vector_db have the same id
         store_id = name or str(uuid.uuid4())
         created_at = int(time.time())
@@ -281,7 +282,7 @@ class OpenAIVectorStoreMixin(ABC):
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
         # search_mode: Literal["keyword", "vector", "hybrid"] = "vector",
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         """Search for chunks in a vector store."""
         # TODO: Add support in the API for this
         search_mode = "vector"
@@ -312,7 +313,7 @@ class OpenAIVectorStoreMixin(ABC):
 
             # Convert response to OpenAI format
             data = []
-            for i, (chunk, score) in enumerate(zip(response.chunks, response.scores, strict=False)):
+            for chunk, score in zip(response.chunks, response.scores, strict=False):
                 # Apply score based filtering
                 if score < score_threshold:
                     continue
@@ -323,18 +324,46 @@ class OpenAIVectorStoreMixin(ABC):
                     if not self._matches_filters(chunk.metadata, filters):
                         continue
 
-                chunk_data = {
-                    "id": f"chunk_{i}",
-                    "object": "vector_store.search_result",
-                    "score": score,
-                    "content": chunk.content.content if hasattr(chunk.content, "content") else str(chunk.content),
-                    "metadata": chunk.metadata,
-                }
-                data.append(chunk_data)
+                # content is InterleavedContent
+                if isinstance(chunk.content, str):
+                    content = [
+                        VectorStoreContent(
+                            type="text",
+                            text=chunk.content,
+                        )
+                    ]
+                elif isinstance(chunk.content, list):
+                    # TODO: Add support for other types of content
+                    content = [
+                        VectorStoreContent(
+                            type="text",
+                            text=item.text,
+                        )
+                        for item in chunk.content
+                        if item.type == "text"
+                    ]
+                else:
+                    if chunk.content.type != "text":
+                        raise ValueError(f"Unsupported content type: {chunk.content.type}")
+                    content = [
+                        VectorStoreContent(
+                            type="text",
+                            text=chunk.content.text,
+                        )
+                    ]
+
+                response_data_item = VectorStoreSearchResponse(
+                    file_id=chunk.metadata.get("file_id", ""),
+                    filename=chunk.metadata.get("filename", ""),
+                    score=score,
+                    attributes=chunk.metadata,
+                    content=content,
+                )
+                data.append(response_data_item)
                 if len(data) >= max_num_results:
                     break
 
-            return VectorStoreSearchResponse(
+            return VectorStoreSearchResponsePage(
                 search_query=search_query,
                 data=data,
                 has_more=False,  # For simplicity, we don't implement pagination here
@@ -344,7 +373,7 @@ class OpenAIVectorStoreMixin(ABC):
         except Exception as e:
             logger.error(f"Error searching vector store {vector_store_id}: {e}")
             # Return empty results on error
-            return VectorStoreSearchResponse(
+            return VectorStoreSearchResponsePage(
                 search_query=search_query,
                 data=[],
                 has_more=False,