mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-04 02:03:44 +00:00
feat: allow returning embeddings and metadata from /vector_stores/ methods; disallow changing Provider ID (#4046)
# What does this PR do?
- Updates `/vector_stores/{vector_store_id}/files/{file_id}/content` to
allow returning `embeddings` and `metadata` using the `extra_query`
- Updates the UI accordingly to display them.
- Update UI to support CRUD operations in the Vector Stores section and
adds a new modal exposing the functionality.
- Updates Vector Store update to fail if a user tries to update Provider
ID (which doesn't make sense to allow)
```python
In [1]: client.vector_stores.files.content(
vector_store_id=vector_store.id,
file_id=file.id,
extra_query={"include_embeddings": True, "include_metadata": True}
)
Out [1]: FileContentResponse(attributes={}, content=[Content(text='This is a test document to check if embeddings are generated properly.\n', type='text', embedding=[0.33760684728622437, ...,], chunk_metadata={'chunk_id': '62a63ae0-c202-f060-1b86-0a688995b8d3', 'document_id': 'file-27291dbc679642ac94ffac6d2810c339', 'source': None, 'created_timestamp': 1762053437, 'updated_timestamp': 1762053437, 'chunk_window': '0-13', 'chunk_tokenizer': 'DEFAULT_TIKTOKEN_TOKENIZER', 'chunk_embedding_model': 'sentence-transformers/nomic
-ai/nomic-embed-text-v1.5', 'chunk_embedding_dimension': 768, 'content_token_count': 13, 'metadata_token_count': 9}, metadata={'filename': 'test-embedding.txt', 'chunk_id': '62a63ae0-c202-f060-1b86-0a688995b8d3', 'document_id': 'file-27291dbc679642ac94ffac6d2810c339', 'token_count': 13, 'metadata_token_count': 9})], file_id='file-27291dbc679642ac94ffac6d2810c339', filename='test-embedding.txt')
```
Screenshots of UI are displayed below:
### List Vector Store with Added "Create New Vector Store"
<img width="1912" height="491" alt="Screenshot 2025-11-06 at 10 47
25 PM"
src="https://github.com/user-attachments/assets/a3a3ddd9-758d-4005-ac9c-5047f03916f3"
/>
### Create New Vector Store
<img width="1918" height="1048" alt="Screenshot 2025-11-06 at 10 47
49 PM"
src="https://github.com/user-attachments/assets/b4dc0d31-696f-4e68-b109-27915090f158"
/>
### Edit Vector Store
<img width="1916" height="1355" alt="Screenshot 2025-11-06 at 10 48
32 PM"
src="https://github.com/user-attachments/assets/ec879c63-4cf7-489f-bb1e-57ccc7931414"
/>
### Vector Store Files Contents page (with Embeddings)
<img width="1914" height="849" alt="Screenshot 2025-11-06 at 11 54
32 PM"
src="https://github.com/user-attachments/assets/3095520d-0e90-41f7-83bd-652f6c3fbf27"
/>
### Vector Store Files Contents Details page (with Embeddings)
<img width="1916" height="1221" alt="Screenshot 2025-11-06 at 11 55
00 PM"
src="https://github.com/user-attachments/assets/e71dbdc5-5b49-472b-a43a-5785f58d196c"
/>
<!-- If resolving an issue, uncomment and update the line below -->
<!-- Closes #[issue-number] -->
## Test Plan
Tests added for Middleware extension and Provider failures.
---------
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
parent
37853ca558
commit
eb3f9ac278
17 changed files with 1161 additions and 125 deletions
|
|
@ -704,34 +704,35 @@ class OpenAIVectorStoreMixin(ABC):
|
|||
# Unknown filter type, default to no match
|
||||
raise ValueError(f"Unsupported filter type: {filter_type}")
|
||||
|
||||
def _chunk_to_vector_store_content(self, chunk: Chunk) -> list[VectorStoreContent]:
|
||||
# content is InterleavedContent
|
||||
def _chunk_to_vector_store_content(
|
||||
self, chunk: Chunk, include_embeddings: bool = False, include_metadata: bool = False
|
||||
) -> list[VectorStoreContent]:
|
||||
def extract_fields() -> dict:
|
||||
"""Extract embedding and metadata fields from chunk based on include flags."""
|
||||
return {
|
||||
"embedding": chunk.embedding if include_embeddings else None,
|
||||
"chunk_metadata": chunk.chunk_metadata if include_metadata else None,
|
||||
"metadata": chunk.metadata if include_metadata else None,
|
||||
}
|
||||
|
||||
fields = extract_fields()
|
||||
|
||||
if isinstance(chunk.content, str):
|
||||
content = [
|
||||
VectorStoreContent(
|
||||
type="text",
|
||||
text=chunk.content,
|
||||
)
|
||||
]
|
||||
content_item = VectorStoreContent(type="text", text=chunk.content, **fields)
|
||||
content = [content_item]
|
||||
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"
|
||||
]
|
||||
content = []
|
||||
for item in chunk.content:
|
||||
if item.type == "text":
|
||||
content_item = VectorStoreContent(type="text", text=item.text, **fields)
|
||||
content.append(content_item)
|
||||
else:
|
||||
if chunk.content.type != "text":
|
||||
raise ValueError(f"Unsupported content type: {chunk.content.type}")
|
||||
content = [
|
||||
VectorStoreContent(
|
||||
type="text",
|
||||
text=chunk.content.text,
|
||||
)
|
||||
]
|
||||
|
||||
content_item = VectorStoreContent(type="text", text=chunk.content.text, **fields)
|
||||
content = [content_item]
|
||||
return content
|
||||
|
||||
async def openai_attach_file_to_vector_store(
|
||||
|
|
@ -820,13 +821,12 @@ class OpenAIVectorStoreMixin(ABC):
|
|||
message=str(e),
|
||||
)
|
||||
|
||||
# Create OpenAI vector store file metadata
|
||||
# Save vector store file to persistent storage AFTER insert_chunks
|
||||
# so that chunks include the embeddings that were generated
|
||||
file_info = vector_store_file_object.model_dump(exclude={"last_error"})
|
||||
file_info["filename"] = file_response.filename if file_response else ""
|
||||
|
||||
# Save vector store file to persistent storage (provider-specific)
|
||||
dict_chunks = [c.model_dump() for c in chunks]
|
||||
# This should be updated to include chunk_id
|
||||
await self._save_openai_vector_store_file(vector_store_id, file_id, file_info, dict_chunks)
|
||||
|
||||
# Update file_ids and file_counts in vector store metadata
|
||||
|
|
@ -921,21 +921,27 @@ class OpenAIVectorStoreMixin(ABC):
|
|||
self,
|
||||
vector_store_id: str,
|
||||
file_id: str,
|
||||
include_embeddings: bool | None = False,
|
||||
include_metadata: bool | None = False,
|
||||
) -> VectorStoreFileContentResponse:
|
||||
"""Retrieves the contents of a vector store file."""
|
||||
if vector_store_id not in self.openai_vector_stores:
|
||||
raise VectorStoreNotFoundError(vector_store_id)
|
||||
|
||||
# Parameters are already provided directly
|
||||
# include_embeddings and include_metadata are now function parameters
|
||||
|
||||
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]
|
||||
content = []
|
||||
for chunk in chunks:
|
||||
content.extend(self._chunk_to_vector_store_content(chunk))
|
||||
content.extend(
|
||||
self._chunk_to_vector_store_content(
|
||||
chunk, include_embeddings=include_embeddings or False, include_metadata=include_metadata or False
|
||||
)
|
||||
)
|
||||
return VectorStoreFileContentResponse(
|
||||
object="vector_store.file_content.page",
|
||||
data=content,
|
||||
has_more=False,
|
||||
next_page=None,
|
||||
)
|
||||
|
||||
async def openai_update_vector_store_file(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue