Use VectorStoreContent vs InterleavedContent in vector store files

This extracts the existing logic to convert chunks to
VectorStoreContent objects into a reusable method and uses that when
returning our list of Vector Store File contents.

It also adds an xfail test for deleting vector store files, as that's
not implemented yet but parking the implementation of that for now.

Signed-off-by: Ben Browning <bbrownin@redhat.com>
This commit is contained in:
Ben Browning 2025-06-18 15:01:42 -04:00
parent 65869d22a4
commit f0d56316a0
5 changed files with 127 additions and 74 deletions

View file

@ -13756,6 +13756,24 @@
"type": "object", "type": "object",
"title": "Response" "title": "Response"
}, },
"VectorStoreContent": {
"type": "object",
"properties": {
"type": {
"type": "string",
"const": "text"
},
"text": {
"type": "string"
}
},
"additionalProperties": false,
"required": [
"type",
"text"
],
"title": "VectorStoreContent"
},
"VectorStoreFileContentsResponse": { "VectorStoreFileContentsResponse": {
"type": "object", "type": "object",
"properties": { "properties": {
@ -13793,7 +13811,7 @@
"content": { "content": {
"type": "array", "type": "array",
"items": { "items": {
"$ref": "#/components/schemas/InterleavedContentItem" "$ref": "#/components/schemas/VectorStoreContent"
} }
} }
}, },
@ -13879,24 +13897,6 @@
], ],
"title": "OpenaiSearchVectorStoreRequest" "title": "OpenaiSearchVectorStoreRequest"
}, },
"VectorStoreContent": {
"type": "object",
"properties": {
"type": {
"type": "string",
"const": "text"
},
"text": {
"type": "string"
}
},
"additionalProperties": false,
"required": [
"type",
"text"
],
"title": "VectorStoreContent"
},
"VectorStoreSearchResponse": { "VectorStoreSearchResponse": {
"type": "object", "type": "object",
"properties": { "properties": {

View file

@ -9616,6 +9616,19 @@ components:
Response: Response:
type: object type: object
title: Response title: Response
VectorStoreContent:
type: object
properties:
type:
type: string
const: text
text:
type: string
additionalProperties: false
required:
- type
- text
title: VectorStoreContent
VectorStoreFileContentsResponse: VectorStoreFileContentsResponse:
type: object type: object
properties: properties:
@ -9636,7 +9649,7 @@ components:
content: content:
type: array type: array
items: items:
$ref: '#/components/schemas/InterleavedContentItem' $ref: '#/components/schemas/VectorStoreContent'
additionalProperties: false additionalProperties: false
required: required:
- file_id - file_id
@ -9693,19 +9706,6 @@ components:
required: required:
- query - query
title: OpenaiSearchVectorStoreRequest title: OpenaiSearchVectorStoreRequest
VectorStoreContent:
type: object
properties:
type:
type: string
const: text
text:
type: string
additionalProperties: false
required:
- type
- text
title: VectorStoreContent
VectorStoreSearchResponse: VectorStoreSearchResponse:
type: object type: object
properties: properties:

View file

@ -12,7 +12,6 @@ from typing import Annotated, Any, Literal, Protocol, runtime_checkable
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from llama_stack.apis.common.content_types import InterleavedContentItem
from llama_stack.apis.inference import InterleavedContent from llama_stack.apis.inference import InterleavedContent
from llama_stack.apis.vector_dbs import VectorDB from llama_stack.apis.vector_dbs import VectorDB
from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
@ -208,7 +207,7 @@ class VectorStoreFileContentsResponse(BaseModel):
file_id: str file_id: str
filename: str filename: str
attributes: dict[str, Any] attributes: dict[str, Any]
content: list[InterleavedContentItem] content: list[VectorStoreContent]
@json_schema_type @json_schema_type

View file

@ -12,7 +12,6 @@ import uuid
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Any from typing import Any
from llama_stack.apis.common.content_types import InterleavedContentItem, TextContentItem
from llama_stack.apis.files import Files from llama_stack.apis.files import Files
from llama_stack.apis.files.files import OpenAIFileObject from llama_stack.apis.files.files import OpenAIFileObject
from llama_stack.apis.vector_dbs import VectorDB from llama_stack.apis.vector_dbs import VectorDB
@ -386,33 +385,7 @@ class OpenAIVectorStoreMixin(ABC):
if not self._matches_filters(chunk.metadata, filters): if not self._matches_filters(chunk.metadata, filters):
continue continue
# content is InterleavedContent content = self._chunk_to_vector_store_content(chunk)
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( response_data_item = VectorStoreSearchResponse(
file_id=chunk.metadata.get("file_id", ""), file_id=chunk.metadata.get("file_id", ""),
@ -488,6 +461,36 @@ class OpenAIVectorStoreMixin(ABC):
# Unknown filter type, default to no match # Unknown filter type, default to no match
raise ValueError(f"Unsupported filter type: {filter_type}") raise ValueError(f"Unsupported filter type: {filter_type}")
def _chunk_to_vector_store_content(self, chunk: Chunk) -> list[VectorStoreContent]:
# 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,
)
]
return content
async def openai_attach_file_to_vector_store( async def openai_attach_file_to_vector_store(
self, self,
vector_store_id: str, vector_store_id: str,
@ -634,20 +637,14 @@ class OpenAIVectorStoreMixin(ABC):
file_info = await self._load_openai_vector_store_file(vector_store_id, file_id) file_info = await self._load_openai_vector_store_file(vector_store_id, file_id)
dict_chunks = await self._load_openai_vector_store_file_contents(vector_store_id, file_id) dict_chunks = await self._load_openai_vector_store_file_contents(vector_store_id, file_id)
chunks = [Chunk.model_validate(c) for c in dict_chunks] chunks = [Chunk.model_validate(c) for c in dict_chunks]
contents: list[InterleavedContentItem] = [] content = []
for chunk in chunks: for chunk in chunks:
content = chunk.content content.extend(self._chunk_to_vector_store_content(chunk))
if isinstance(content, str):
contents.append(TextContentItem(text=content))
elif isinstance(content, InterleavedContentItem):
contents.append(content)
else:
contents.extend(contents)
return VectorStoreFileContentsResponse( return VectorStoreFileContentsResponse(
file_id=file_id, file_id=file_id,
filename=file_info.get("filename", ""), filename=file_info.get("filename", ""),
attributes=file_info.get("attributes", {}), attributes=file_info.get("attributes", {}),
content=contents, content=content,
) )
async def openai_update_vector_store_file( async def openai_update_vector_store_file(
@ -684,6 +681,10 @@ class OpenAIVectorStoreMixin(ABC):
await self._delete_openai_vector_store_file_from_storage(vector_store_id, file_id) await self._delete_openai_vector_store_file_from_storage(vector_store_id, file_id)
# TODO: We need to actually delete the embeddings from the underlying vector store... # TODO: We need to actually delete the embeddings from the underlying vector store...
# Also uncomment the corresponding integration test marked as xfail
#
# test_openai_vector_store_delete_file_removes_from_vector_store in
# tests/integration/vector_io/test_openai_vector_stores.py
# Update in-memory cache # Update in-memory cache
store_info["file_ids"].remove(file_id) store_info["file_ids"].remove(file_id)

View file

@ -440,7 +440,7 @@ def test_openai_vector_store_search_with_max_num_results(
assert len(search_response.data) == 2 assert len(search_response.data) == 2
def test_openai_vector_store_attach_file_response_attributes(compat_client_with_empty_stores, client_with_models): def test_openai_vector_store_attach_file(compat_client_with_empty_stores, client_with_models):
"""Test OpenAI vector store attach file.""" """Test OpenAI vector store attach file."""
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
@ -453,7 +453,7 @@ def test_openai_vector_store_attach_file_response_attributes(compat_client_with_
vector_store = compat_client.vector_stores.create(name="test_store") vector_store = compat_client.vector_stores.create(name="test_store")
# Create a file # Create a file
test_content = b"This is a test file" test_content = b"The secret string is foobazbar."
with BytesIO(test_content) as file_buffer: with BytesIO(test_content) as file_buffer:
file_buffer.name = "openai_test.txt" file_buffer.name = "openai_test.txt"
file = compat_client.files.create(file=file_buffer, purpose="assistants") file = compat_client.files.create(file=file_buffer, purpose="assistants")
@ -480,6 +480,16 @@ def test_openai_vector_store_attach_file_response_attributes(compat_client_with_
assert updated_vector_store.file_counts.failed == 0 assert updated_vector_store.file_counts.failed == 0
assert updated_vector_store.file_counts.in_progress == 0 assert updated_vector_store.file_counts.in_progress == 0
# Search using OpenAI API to confirm our file attached
search_response = compat_client.vector_stores.search(
vector_store_id=vector_store.id, query="What is the secret string?", max_num_results=1
)
assert search_response is not None
assert len(search_response.data) > 0
top_result = search_response.data[0]
top_content = top_result.content[0].text
assert "foobazbar" in top_content.lower()
def test_openai_vector_store_attach_files_on_creation(compat_client_with_empty_stores, client_with_models): def test_openai_vector_store_attach_files_on_creation(compat_client_with_empty_stores, client_with_models):
"""Test OpenAI vector store attach files on creation.""" """Test OpenAI vector store attach files on creation."""
@ -689,6 +699,49 @@ def test_openai_vector_store_delete_file(compat_client_with_empty_stores, client
assert updated_vector_store.file_counts.in_progress == 0 assert updated_vector_store.file_counts.in_progress == 0
# TODO: Remove this xfail once we have a way to remove embeddings from vector store
@pytest.mark.xfail(reason="Vector Store Files delete doesn't remove embeddings from vecntor store", strict=True)
def test_openai_vector_store_delete_file_removes_from_vector_store(compat_client_with_empty_stores, client_with_models):
"""Test OpenAI vector store delete file removes from vector store."""
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
if isinstance(compat_client_with_empty_stores, LlamaStackClient):
pytest.skip("Vector Store Files attach is not yet supported with LlamaStackClient")
compat_client = compat_client_with_empty_stores
# Create a vector store
vector_store = compat_client.vector_stores.create(name="test_store")
# Create a file
test_content = b"The secret string is foobazbar."
with BytesIO(test_content) as file_buffer:
file_buffer.name = "openai_test.txt"
file = compat_client.files.create(file=file_buffer, purpose="assistants")
# Attach the file to the vector store
file_attach_response = compat_client.vector_stores.files.create(
vector_store_id=vector_store.id,
file_id=file.id,
)
assert file_attach_response.status == "completed"
# Search using OpenAI API to confirm our file attached
search_response = compat_client.vector_stores.search(
vector_store_id=vector_store.id, query="What is the secret string?", max_num_results=1
)
assert "foobazbar" in search_response.data[0].content[0].text.lower()
# Delete the file
compat_client.vector_stores.files.delete(vector_store_id=vector_store.id, file_id=file.id)
# Search using OpenAI API to confirm our file deleted
search_response = compat_client.vector_stores.search(
vector_store_id=vector_store.id, query="What is the secret string?", max_num_results=1
)
assert not search_response.data
def test_openai_vector_store_update_file(compat_client_with_empty_stores, client_with_models): def test_openai_vector_store_update_file(compat_client_with_empty_stores, client_with_models):
"""Test OpenAI vector store update file.""" """Test OpenAI vector store update file."""
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)