feat: Adding optional embeddings to content

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-11-03 14:48:52 -05:00
parent 97ccfb5e62
commit aefbb6f9ea
20 changed files with 1314 additions and 132 deletions

View file

@ -2691,7 +2691,12 @@ paths:
responses: responses:
'200': '200':
description: >- description: >-
<<<<<<< HEAD
A VectorStoreFileContentResponse representing the file contents. A VectorStoreFileContentResponse representing the file contents.
=======
File contents, optionally with embeddings and metadata based on extra_query
parameters.
>>>>>>> 639f0daa (feat: Adding optional embeddings to content)
content: content:
application/json: application/json:
schema: schema:
@ -2726,6 +2731,23 @@ paths:
required: true required: true
schema: schema:
type: string type: string
- name: extra_query
in: query
description: >-
Optional extra parameters to control response format. Set include_embeddings=true
to include embedding vectors. Set include_metadata=true to include chunk
metadata.
required: false
schema:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
deprecated: false deprecated: false
/v1/vector_stores/{vector_store_id}/search: /v1/vector_stores/{vector_store_id}/search:
post: post:
@ -10102,6 +10124,28 @@ components:
text: text:
type: string type: string
description: The actual text content description: The actual text content
embedding:
type: array
items:
type: number
description: >-
Optional embedding vector for this content chunk (when requested via extra_body)
chunk_metadata:
$ref: '#/components/schemas/ChunkMetadata'
description: >-
Optional chunk metadata (when requested via extra_body)
metadata:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
description: >-
Optional user-defined metadata (when requested via extra_body)
additionalProperties: false additionalProperties: false
required: required:
- type - type

View file

@ -2688,7 +2688,12 @@ paths:
responses: responses:
'200': '200':
description: >- description: >-
<<<<<<< HEAD
A VectorStoreFileContentResponse representing the file contents. A VectorStoreFileContentResponse representing the file contents.
=======
File contents, optionally with embeddings and metadata based on extra_query
parameters.
>>>>>>> 639f0daa (feat: Adding optional embeddings to content)
content: content:
application/json: application/json:
schema: schema:
@ -2723,6 +2728,23 @@ paths:
required: true required: true
schema: schema:
type: string type: string
- name: extra_query
in: query
description: >-
Optional extra parameters to control response format. Set include_embeddings=true
to include embedding vectors. Set include_metadata=true to include chunk
metadata.
required: false
schema:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
deprecated: false deprecated: false
/v1/vector_stores/{vector_store_id}/search: /v1/vector_stores/{vector_store_id}/search:
post: post:
@ -9386,6 +9408,28 @@ components:
text: text:
type: string type: string
description: The actual text content description: The actual text content
embedding:
type: array
items:
type: number
description: >-
Optional embedding vector for this content chunk (when requested via extra_body)
chunk_metadata:
$ref: '#/components/schemas/ChunkMetadata'
description: >-
Optional chunk metadata (when requested via extra_body)
metadata:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
description: >-
Optional user-defined metadata (when requested via extra_body)
additionalProperties: false additionalProperties: false
required: required:
- type - type

View file

@ -2691,7 +2691,12 @@ paths:
responses: responses:
'200': '200':
description: >- description: >-
<<<<<<< HEAD
A VectorStoreFileContentResponse representing the file contents. A VectorStoreFileContentResponse representing the file contents.
=======
File contents, optionally with embeddings and metadata based on extra_query
parameters.
>>>>>>> 639f0daa (feat: Adding optional embeddings to content)
content: content:
application/json: application/json:
schema: schema:
@ -2726,6 +2731,23 @@ paths:
required: true required: true
schema: schema:
type: string type: string
- name: extra_query
in: query
description: >-
Optional extra parameters to control response format. Set include_embeddings=true
to include embedding vectors. Set include_metadata=true to include chunk
metadata.
required: false
schema:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
deprecated: false deprecated: false
/v1/vector_stores/{vector_store_id}/search: /v1/vector_stores/{vector_store_id}/search:
post: post:
@ -10102,6 +10124,28 @@ components:
text: text:
type: string type: string
description: The actual text content description: The actual text content
embedding:
type: array
items:
type: number
description: >-
Optional embedding vector for this content chunk (when requested via extra_body)
chunk_metadata:
$ref: '#/components/schemas/ChunkMetadata'
description: >-
Optional chunk metadata (when requested via extra_body)
metadata:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
description: >-
Optional user-defined metadata (when requested via extra_body)
additionalProperties: false additionalProperties: false
required: required:
- type - type

View file

@ -10,7 +10,7 @@
# the root directory of this source tree. # the root directory of this source tree.
from typing import Annotated, Any, Literal, Protocol, runtime_checkable from typing import Annotated, Any, Literal, Protocol, runtime_checkable
from fastapi import Body from fastapi import Body, Query
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from llama_stack.apis.common.tracing import telemetry_traceable from llama_stack.apis.common.tracing import telemetry_traceable
@ -224,10 +224,16 @@ class VectorStoreContent(BaseModel):
:param type: Content type, currently only "text" is supported :param type: Content type, currently only "text" is supported
:param text: The actual text content :param text: The actual text content
:param embedding: Optional embedding vector for this content chunk (when requested via extra_body)
:param chunk_metadata: Optional chunk metadata (when requested via extra_body)
:param metadata: Optional user-defined metadata (when requested via extra_body)
""" """
type: Literal["text"] type: Literal["text"]
text: str text: str
embedding: list[float] | None = None
chunk_metadata: ChunkMetadata | None = None
metadata: dict[str, Any] | None = None
@json_schema_type @json_schema_type
@ -395,22 +401,6 @@ class VectorStoreListFilesResponse(BaseModel):
has_more: bool = False has_more: bool = False
@json_schema_type
class VectorStoreFileContentResponse(BaseModel):
"""Represents the parsed content of a vector store file.
:param object: The object type, which is always `vector_store.file_content.page`
:param data: Parsed content of the file
:param has_more: Indicates if there are more content pages to fetch
:param next_page: The token for the next page, if any
"""
object: Literal["vector_store.file_content.page"] = "vector_store.file_content.page"
data: list[VectorStoreContent]
has_more: bool
next_page: str | None = None
@json_schema_type @json_schema_type
class VectorStoreFileDeleteResponse(BaseModel): class VectorStoreFileDeleteResponse(BaseModel):
"""Response from deleting a vector store file. """Response from deleting a vector store file.
@ -732,12 +722,16 @@ class VectorIO(Protocol):
self, self,
vector_store_id: str, vector_store_id: str,
file_id: str, file_id: str,
) -> VectorStoreFileContentResponse: include_embeddings: Annotated[bool | None, Query(default=False)] = False,
include_metadata: Annotated[bool | None, Query(default=False)] = False,
) -> VectorStoreFileContentsResponse:
"""Retrieves the contents of a vector store file. """Retrieves the contents of a vector store file.
:param vector_store_id: The ID of the vector store containing the file to retrieve. :param vector_store_id: The ID of the vector store containing the file to retrieve.
:param file_id: The ID of the file to retrieve. :param file_id: The ID of the file to retrieve.
:returns: A VectorStoreFileContentResponse representing the file contents. :param include_embeddings: Whether to include embedding vectors in the response.
:param include_metadata: Whether to include chunk metadata in the response.
:returns: File contents, optionally with embeddings and metadata based on query parameters.
""" """
... ...

View file

@ -389,6 +389,12 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
matched_func, path_params, route_path, webmethod = find_matching_route(options.method, path, self.route_impls) matched_func, path_params, route_path, webmethod = find_matching_route(options.method, path, self.route_impls)
body |= path_params body |= path_params
# Pass through params that aren't already handled as path params
if options.params:
extra_query_params = {k: v for k, v in options.params.items() if k not in path_params}
if extra_query_params:
body["extra_query"] = extra_query_params
body, field_names = self._handle_file_uploads(options, body) body, field_names = self._handle_file_uploads(options, body)
body = self._convert_body(matched_func, body, exclude_params=set(field_names)) body = self._convert_body(matched_func, body, exclude_params=set(field_names))

View file

@ -24,7 +24,7 @@ from llama_stack.apis.vector_io import (
VectorStoreChunkingStrategyStaticConfig, VectorStoreChunkingStrategyStaticConfig,
VectorStoreDeleteResponse, VectorStoreDeleteResponse,
VectorStoreFileBatchObject, VectorStoreFileBatchObject,
VectorStoreFileContentResponse, VectorStoreFileContentsResponse,
VectorStoreFileDeleteResponse, VectorStoreFileDeleteResponse,
VectorStoreFileObject, VectorStoreFileObject,
VectorStoreFilesListInBatchResponse, VectorStoreFilesListInBatchResponse,
@ -247,6 +247,13 @@ class VectorIORouter(VectorIO):
metadata: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None,
) -> VectorStoreObject: ) -> VectorStoreObject:
logger.debug(f"VectorIORouter.openai_update_vector_store: {vector_store_id}") logger.debug(f"VectorIORouter.openai_update_vector_store: {vector_store_id}")
# Check if provider_id is being changed (not supported)
if metadata and "provider_id" in metadata:
current_store = await self.routing_table.get_object_by_identifier("vector_store", vector_store_id)
if current_store and current_store.provider_id != metadata["provider_id"]:
raise ValueError("provider_id cannot be changed after vector store creation")
provider = await self.routing_table.get_provider_impl(vector_store_id) provider = await self.routing_table.get_provider_impl(vector_store_id)
return await provider.openai_update_vector_store( return await provider.openai_update_vector_store(
vector_store_id=vector_store_id, vector_store_id=vector_store_id,
@ -338,12 +345,19 @@ class VectorIORouter(VectorIO):
self, self,
vector_store_id: str, vector_store_id: str,
file_id: str, file_id: str,
) -> VectorStoreFileContentResponse: include_embeddings: bool | None = False,
logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_contents: {vector_store_id}, {file_id}") include_metadata: bool | None = False,
provider = await self.routing_table.get_provider_impl(vector_store_id) ) -> VectorStoreFileContentsResponse:
return await provider.openai_retrieve_vector_store_file_contents( logger.debug(
f"VectorIORouter.openai_retrieve_vector_store_file_contents: {vector_store_id}, {file_id}, "
f"include_embeddings={include_embeddings}, include_metadata={include_metadata}"
)
return await self.routing_table.openai_retrieve_vector_store_file_contents(
vector_store_id=vector_store_id, vector_store_id=vector_store_id,
file_id=file_id, file_id=file_id,
include_embeddings=include_embeddings,
include_metadata=include_metadata,
) )
async def openai_update_vector_store_file( async def openai_update_vector_store_file(

View file

@ -15,7 +15,7 @@ from llama_stack.apis.vector_io.vector_io import (
SearchRankingOptions, SearchRankingOptions,
VectorStoreChunkingStrategy, VectorStoreChunkingStrategy,
VectorStoreDeleteResponse, VectorStoreDeleteResponse,
VectorStoreFileContentResponse, VectorStoreFileContentsResponse,
VectorStoreFileDeleteResponse, VectorStoreFileDeleteResponse,
VectorStoreFileObject, VectorStoreFileObject,
VectorStoreFileStatus, VectorStoreFileStatus,
@ -195,12 +195,17 @@ class VectorStoresRoutingTable(CommonRoutingTableImpl):
self, self,
vector_store_id: str, vector_store_id: str,
file_id: str, file_id: str,
) -> VectorStoreFileContentResponse: include_embeddings: bool | None = False,
include_metadata: bool | None = False,
) -> VectorStoreFileContentsResponse:
await self.assert_action_allowed("read", "vector_store", vector_store_id) await self.assert_action_allowed("read", "vector_store", vector_store_id)
provider = await self.get_provider_impl(vector_store_id) provider = await self.get_provider_impl(vector_store_id)
return await provider.openai_retrieve_vector_store_file_contents( return await provider.openai_retrieve_vector_store_file_contents(
vector_store_id=vector_store_id, vector_store_id=vector_store_id,
file_id=file_id, file_id=file_id,
include_embeddings=include_embeddings,
include_metadata=include_metadata,
) )
async def openai_update_vector_store_file( async def openai_update_vector_store_file(

View file

@ -0,0 +1,49 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import json
import re
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from llama_stack.log import get_logger
logger = get_logger(name=__name__, category="core::middleware")
# Patterns for endpoints that need query parameter injection
QUERY_PARAM_ENDPOINTS = [
# /vector_stores/{vector_store_id}/files/{file_id}/content
re.compile(r"/vector_stores/[^/]+/files/[^/]+/content$"),
]
class QueryParamsMiddleware(BaseHTTPMiddleware):
"""Middleware to inject query parameters into extra_query for specific endpoints"""
async def dispatch(self, request: Request, call_next):
# Check if this is an endpoint that needs query parameter injection
if request.method == "GET" and any(pattern.search(str(request.url.path)) for pattern in QUERY_PARAM_ENDPOINTS):
# Extract all query parameters and convert to appropriate types
extra_query = {}
query_params = dict(request.query_params)
# Convert query parameters using JSON parsing for robust type conversion
for key, value in query_params.items():
try:
# parse as JSON to handles booleans, numbers, strings properly
extra_query[key] = json.loads(value)
except (json.JSONDecodeError, ValueError):
# if parsing fails, keep as string
extra_query[key] = value
if extra_query:
# Store the extra_query in request state so we can access it later
request.state.extra_query = extra_query
logger.debug(f"QueryParamsMiddleware extracted extra_query: {extra_query}")
response = await call_next(request)
return response

View file

@ -46,6 +46,7 @@ from llama_stack.core.request_headers import (
request_provider_data_context, request_provider_data_context,
user_from_scope, user_from_scope,
) )
from llama_stack.core.server.query_params_middleware import QueryParamsMiddleware
from llama_stack.core.server.routes import get_all_api_routes from llama_stack.core.server.routes import get_all_api_routes
from llama_stack.core.stack import ( from llama_stack.core.stack import (
Stack, Stack,
@ -263,6 +264,10 @@ def create_dynamic_typed_route(func: Any, method: str, route: str) -> Callable:
is_streaming = is_streaming_request(func.__name__, request, **kwargs) is_streaming = is_streaming_request(func.__name__, request, **kwargs)
# Inject extra_query from middleware if available
if hasattr(request.state, "extra_query") and request.state.extra_query:
kwargs["extra_query"] = request.state.extra_query
try: try:
if is_streaming: if is_streaming:
context_vars = [CURRENT_TRACE_CONTEXT, PROVIDER_DATA_VAR] context_vars = [CURRENT_TRACE_CONTEXT, PROVIDER_DATA_VAR]
@ -402,6 +407,9 @@ def create_app() -> StackApp:
if not os.environ.get("LLAMA_STACK_DISABLE_VERSION_CHECK"): if not os.environ.get("LLAMA_STACK_DISABLE_VERSION_CHECK"):
app.add_middleware(ClientVersionMiddleware) app.add_middleware(ClientVersionMiddleware)
# handle extra_query for specific GET requests
app.add_middleware(QueryParamsMiddleware)
impls = app.stack.impls impls = app.stack.impls
if config.server.auth: if config.server.auth:

View file

@ -30,7 +30,7 @@ from llama_stack.apis.vector_io import (
VectorStoreContent, VectorStoreContent,
VectorStoreDeleteResponse, VectorStoreDeleteResponse,
VectorStoreFileBatchObject, VectorStoreFileBatchObject,
VectorStoreFileContentResponse, VectorStoreFileContentsResponse,
VectorStoreFileCounts, VectorStoreFileCounts,
VectorStoreFileDeleteResponse, VectorStoreFileDeleteResponse,
VectorStoreFileLastError, VectorStoreFileLastError,
@ -450,7 +450,7 @@ class OpenAIVectorStoreMixin(ABC):
# Now that our vector store is created, attach any files that were provided # Now that our vector store is created, attach any files that were provided
file_ids = params.file_ids or [] file_ids = params.file_ids or []
tasks = [self.openai_attach_file_to_vector_store(vector_store_id, file_id) for file_id in file_ids] tasks = [self.openai_attach_file_to_vector_store(vector_store_id, file_id) for file_id in file_ids]
await asyncio.gather(*tasks) await asyncio.gather(*tasks, return_exceptions=True)
# Get the updated store info and return it # Get the updated store info and return it
store_info = self.openai_vector_stores[vector_store_id] store_info = self.openai_vector_stores[vector_store_id]
@ -704,34 +704,35 @@ 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]: def _extract_chunk_fields(self, chunk: Chunk, include_embeddings: bool, include_metadata: bool) -> dict:
# content is InterleavedContent """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,
}
def _chunk_to_vector_store_content(
self, chunk: Chunk, include_embeddings: bool = False, include_metadata: bool = False
) -> list[VectorStoreContent]:
fields = self._extract_chunk_fields(chunk, include_embeddings, include_metadata)
if isinstance(chunk.content, str): if isinstance(chunk.content, str):
content = [ content_item = VectorStoreContent(type="text", text=chunk.content, **fields)
VectorStoreContent( content = [content_item]
type="text",
text=chunk.content,
)
]
elif isinstance(chunk.content, list): elif isinstance(chunk.content, list):
# TODO: Add support for other types of content # TODO: Add support for other types of content
content = [ content = []
VectorStoreContent( for item in chunk.content:
type="text", if item.type == "text":
text=item.text, content_item = VectorStoreContent(type="text", text=item.text, **fields)
) content.append(content_item)
for item in chunk.content
if item.type == "text"
]
else: else:
if chunk.content.type != "text": if chunk.content.type != "text":
raise ValueError(f"Unsupported content type: {chunk.content.type}") raise ValueError(f"Unsupported content type: {chunk.content.type}")
content = [
VectorStoreContent( content_item = VectorStoreContent(type="text", text=chunk.content.text, **fields)
type="text", content = [content_item]
text=chunk.content.text,
)
]
return content return content
async def openai_attach_file_to_vector_store( async def openai_attach_file_to_vector_store(
@ -820,13 +821,12 @@ class OpenAIVectorStoreMixin(ABC):
message=str(e), 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 = vector_store_file_object.model_dump(exclude={"last_error"})
file_info["filename"] = file_response.filename if file_response else "" 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] 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) 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 # Update file_ids and file_counts in vector store metadata
@ -921,21 +921,28 @@ class OpenAIVectorStoreMixin(ABC):
self, self,
vector_store_id: str, vector_store_id: str,
file_id: str, file_id: str,
) -> VectorStoreFileContentResponse: include_embeddings: bool | None = False,
include_metadata: bool | None = False,
) -> VectorStoreFileContentsResponse:
"""Retrieves the contents of a vector store file.""" """Retrieves the contents of a vector store file."""
if vector_store_id not in self.openai_vector_stores: if vector_store_id not in self.openai_vector_stores:
raise VectorStoreNotFoundError(vector_store_id) raise VectorStoreNotFoundError(vector_store_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]
content = [] content = []
for chunk in chunks: for chunk in chunks:
content.extend(self._chunk_to_vector_store_content(chunk)) content.extend(
return VectorStoreFileContentResponse( self._chunk_to_vector_store_content(
object="vector_store.file_content.page", chunk, include_embeddings=include_embeddings or False, include_metadata=include_metadata or False
data=content, )
has_more=False, )
next_page=None, return VectorStoreFileContentsResponse(
file_id=file_id,
filename=file_info.get("filename", ""),
attributes=file_info.get("attributes", {}),
content=content,
) )
async def openai_update_vector_store_file( async def openai_update_vector_store_file(

View file

@ -8,6 +8,9 @@ import type {
import { useRouter } from "next/navigation"; import { useRouter } from "next/navigation";
import { usePagination } from "@/hooks/use-pagination"; import { usePagination } from "@/hooks/use-pagination";
import { Button } from "@/components/ui/button"; import { Button } from "@/components/ui/button";
import { Plus, Trash2, Search, Edit, X } from "lucide-react";
import { useState } from "react";
import { Input } from "@/components/ui/input";
import { import {
Table, Table,
TableBody, TableBody,
@ -17,9 +20,21 @@ import {
TableRow, TableRow,
} from "@/components/ui/table"; } from "@/components/ui/table";
import { Skeleton } from "@/components/ui/skeleton"; import { Skeleton } from "@/components/ui/skeleton";
import { useAuthClient } from "@/hooks/use-auth-client";
import {
VectorStoreEditor,
VectorStoreFormData,
} from "@/components/vector-stores/vector-store-editor";
export default function VectorStoresPage() { export default function VectorStoresPage() {
const router = useRouter(); const router = useRouter();
const client = useAuthClient();
const [deletingStores, setDeletingStores] = useState<Set<string>>(new Set());
const [searchTerm, setSearchTerm] = useState("");
const [showVectorStoreModal, setShowVectorStoreModal] = useState(false);
const [editingStore, setEditingStore] = useState<VectorStore | null>(null);
const [modalError, setModalError] = useState<string | null>(null);
const [showSuccessState, setShowSuccessState] = useState(false);
const { const {
data: stores, data: stores,
status, status,
@ -47,6 +62,142 @@ export default function VectorStoresPage() {
} }
}, [status, hasMore, loadMore]); }, [status, hasMore, loadMore]);
// Handle ESC key to close modal
React.useEffect(() => {
const handleEscape = (event: KeyboardEvent) => {
if (event.key === "Escape" && showVectorStoreModal) {
handleCancel();
}
};
document.addEventListener("keydown", handleEscape);
return () => document.removeEventListener("keydown", handleEscape);
}, [showVectorStoreModal]);
const handleDeleteVectorStore = async (storeId: string) => {
if (
!confirm(
"Are you sure you want to delete this vector store? This action cannot be undone."
)
) {
return;
}
setDeletingStores(prev => new Set([...prev, storeId]));
try {
await client.vectorStores.delete(storeId);
// Reload the data to reflect the deletion
window.location.reload();
} catch (err: unknown) {
console.error("Failed to delete vector store:", err);
const errorMessage = err instanceof Error ? err.message : "Unknown error";
alert(`Failed to delete vector store: ${errorMessage}`);
} finally {
setDeletingStores(prev => {
const newSet = new Set(prev);
newSet.delete(storeId);
return newSet;
});
}
};
const handleSaveVectorStore = async (formData: VectorStoreFormData) => {
try {
setModalError(null);
if (editingStore) {
// Update existing vector store
const updateParams: {
name?: string;
extra_body?: Record<string, unknown>;
} = {};
// Only include fields that have changed or are provided
if (formData.name && formData.name !== editingStore.name) {
updateParams.name = formData.name;
}
// Add all parameters to extra_body (except provider_id which can't be changed)
const extraBody: Record<string, unknown> = {};
if (formData.embedding_model) {
extraBody.embedding_model = formData.embedding_model;
}
if (formData.embedding_dimension) {
extraBody.embedding_dimension = formData.embedding_dimension;
}
if (Object.keys(extraBody).length > 0) {
updateParams.extra_body = extraBody;
}
await client.vectorStores.update(editingStore.id, updateParams);
// Show success state with close button
setShowSuccessState(true);
setModalError(
"✅ Vector store updated successfully! You can close this modal and refresh the page to see changes."
);
return;
}
const createParams: {
name?: string;
provider_id?: string;
extra_body?: Record<string, unknown>;
} = {
name: formData.name || undefined,
};
// Extract provider_id to top-level (like Python client does)
if (formData.provider_id) {
createParams.provider_id = formData.provider_id;
}
// Add remaining parameters to extra_body
const extraBody: Record<string, unknown> = {};
if (formData.provider_id) {
extraBody.provider_id = formData.provider_id;
}
if (formData.embedding_model) {
extraBody.embedding_model = formData.embedding_model;
}
if (formData.embedding_dimension) {
extraBody.embedding_dimension = formData.embedding_dimension;
}
if (Object.keys(extraBody).length > 0) {
createParams.extra_body = extraBody;
}
await client.vectorStores.create(createParams);
// Show success state with close button
setShowSuccessState(true);
setModalError(
"✅ Vector store created successfully! You can close this modal and refresh the page to see changes."
);
} catch (err: unknown) {
console.error("Failed to create vector store:", err);
const errorMessage =
err instanceof Error ? err.message : "Failed to create vector store";
setModalError(errorMessage);
}
};
const handleEditVectorStore = (store: VectorStore) => {
setEditingStore(store);
setShowVectorStoreModal(true);
setModalError(null);
};
const handleCancel = () => {
setShowVectorStoreModal(false);
setEditingStore(null);
setModalError(null);
setShowSuccessState(false);
};
const renderContent = () => { const renderContent = () => {
if (status === "loading") { if (status === "loading") {
return ( return (
@ -66,73 +217,190 @@ export default function VectorStoresPage() {
return <p>No vector stores found.</p>; return <p>No vector stores found.</p>;
} }
return ( // Filter stores based on search term
<div className="overflow-auto flex-1 min-h-0"> const filteredStores = stores.filter(store => {
<Table> if (!searchTerm) return true;
<TableHeader>
<TableRow>
<TableHead>ID</TableHead>
<TableHead>Name</TableHead>
<TableHead>Created</TableHead>
<TableHead>Completed</TableHead>
<TableHead>Cancelled</TableHead>
<TableHead>Failed</TableHead>
<TableHead>In Progress</TableHead>
<TableHead>Total</TableHead>
<TableHead>Usage Bytes</TableHead>
<TableHead>Provider ID</TableHead>
<TableHead>Provider Vector DB ID</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{stores.map(store => {
const fileCounts = store.file_counts;
const metadata = store.metadata || {};
const providerId = metadata.provider_id ?? "";
const providerDbId = metadata.provider_vector_db_id ?? "";
return ( const searchLower = searchTerm.toLowerCase();
<TableRow return (
key={store.id} store.id.toLowerCase().includes(searchLower) ||
onClick={() => router.push(`/logs/vector-stores/${store.id}`)} (store.name && store.name.toLowerCase().includes(searchLower)) ||
className="cursor-pointer hover:bg-muted/50" (store.metadata?.provider_id &&
> String(store.metadata.provider_id)
<TableCell> .toLowerCase()
<Button .includes(searchLower)) ||
variant="link" (store.metadata?.provider_vector_db_id &&
className="p-0 h-auto font-mono text-blue-600 hover:text-blue-800 dark:text-blue-400 dark:hover:text-blue-300" String(store.metadata.provider_vector_db_id)
onClick={() => .toLowerCase()
router.push(`/logs/vector-stores/${store.id}`) .includes(searchLower))
} );
> });
{store.id}
</Button> return (
</TableCell> <div className="space-y-4">
<TableCell>{store.name}</TableCell> {/* Search Bar */}
<TableCell> <div className="relative flex-1 max-w-md">
{new Date(store.created_at * 1000).toLocaleString()} <Search className="absolute left-3 top-1/2 transform -translate-y-1/2 text-muted-foreground h-4 w-4" />
</TableCell> <Input
<TableCell>{fileCounts.completed}</TableCell> placeholder="Search vector stores..."
<TableCell>{fileCounts.cancelled}</TableCell> value={searchTerm}
<TableCell>{fileCounts.failed}</TableCell> onChange={e => setSearchTerm(e.target.value)}
<TableCell>{fileCounts.in_progress}</TableCell> className="pl-10"
<TableCell>{fileCounts.total}</TableCell> />
<TableCell>{store.usage_bytes}</TableCell> </div>
<TableCell>{providerId}</TableCell>
<TableCell>{providerDbId}</TableCell> <div className="overflow-auto flex-1 min-h-0">
</TableRow> <Table>
); <TableHeader>
})} <TableRow>
</TableBody> <TableHead>ID</TableHead>
</Table> <TableHead>Name</TableHead>
<TableHead>Created</TableHead>
<TableHead>Completed</TableHead>
<TableHead>Cancelled</TableHead>
<TableHead>Failed</TableHead>
<TableHead>In Progress</TableHead>
<TableHead>Total</TableHead>
<TableHead>Usage Bytes</TableHead>
<TableHead>Provider ID</TableHead>
<TableHead>Provider Vector DB ID</TableHead>
<TableHead>Actions</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{filteredStores.map(store => {
const fileCounts = store.file_counts;
const metadata = store.metadata || {};
const providerId = metadata.provider_id ?? "";
const providerDbId = metadata.provider_vector_db_id ?? "";
return (
<TableRow
key={store.id}
onClick={() =>
router.push(`/logs/vector-stores/${store.id}`)
}
className="cursor-pointer hover:bg-muted/50"
>
<TableCell>
<Button
variant="link"
className="p-0 h-auto font-mono text-blue-600 hover:text-blue-800 dark:text-blue-400 dark:hover:text-blue-300"
onClick={() =>
router.push(`/logs/vector-stores/${store.id}`)
}
>
{store.id}
</Button>
</TableCell>
<TableCell>{store.name}</TableCell>
<TableCell>
{new Date(store.created_at * 1000).toLocaleString()}
</TableCell>
<TableCell>{fileCounts.completed}</TableCell>
<TableCell>{fileCounts.cancelled}</TableCell>
<TableCell>{fileCounts.failed}</TableCell>
<TableCell>{fileCounts.in_progress}</TableCell>
<TableCell>{fileCounts.total}</TableCell>
<TableCell>{store.usage_bytes}</TableCell>
<TableCell>{providerId}</TableCell>
<TableCell>{providerDbId}</TableCell>
<TableCell>
<div className="flex gap-2">
<Button
variant="outline"
size="sm"
onClick={e => {
e.stopPropagation();
handleEditVectorStore(store);
}}
>
<Edit className="h-4 w-4" />
</Button>
<Button
variant="outline"
size="sm"
onClick={e => {
e.stopPropagation();
handleDeleteVectorStore(store.id);
}}
disabled={deletingStores.has(store.id)}
>
{deletingStores.has(store.id) ? (
"Deleting..."
) : (
<>
<Trash2 className="h-4 w-4" />
</>
)}
</Button>
</div>
</TableCell>
</TableRow>
);
})}
</TableBody>
</Table>
</div>
</div> </div>
); );
}; };
return ( return (
<div className="space-y-4"> <div className="space-y-4">
<h1 className="text-2xl font-semibold">Vector Stores</h1> <div className="flex items-center justify-between">
<h1 className="text-2xl font-semibold">Vector Stores</h1>
<Button
onClick={() => setShowVectorStoreModal(true)}
disabled={status === "loading"}
>
<Plus className="h-4 w-4 mr-2" />
New Vector Store
</Button>
</div>
{renderContent()} {renderContent()}
{/* Create Vector Store Modal */}
{showVectorStoreModal && (
<div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50">
<div className="bg-background border rounded-lg shadow-lg max-w-2xl w-full mx-4 max-h-[90vh] overflow-hidden">
<div className="p-6 border-b flex items-center justify-between">
<h2 className="text-2xl font-bold">
{editingStore ? "Edit Vector Store" : "Create New Vector Store"}
</h2>
<Button
variant="ghost"
size="sm"
onClick={handleCancel}
className="p-1 h-auto"
>
<X className="h-4 w-4" />
</Button>
</div>
<div className="p-6 overflow-y-auto max-h-[calc(90vh-120px)]">
<VectorStoreEditor
onSave={handleSaveVectorStore}
onCancel={handleCancel}
error={modalError}
showSuccessState={showSuccessState}
isEditing={!!editingStore}
initialData={
editingStore
? {
name: editingStore.name || "",
embedding_model:
editingStore.metadata?.embedding_model || "",
embedding_dimension:
editingStore.metadata?.embedding_dimension || 768,
provider_id: editingStore.metadata?.provider_id || "",
}
: undefined
}
/>
</div>
</div>
</div>
)}
</div> </div>
); );
} }

View file

@ -2,7 +2,7 @@ import React from "react";
import { render, screen, fireEvent } from "@testing-library/react"; import { render, screen, fireEvent } from "@testing-library/react";
import "@testing-library/jest-dom"; import "@testing-library/jest-dom";
import { PromptEditor } from "./prompt-editor"; import { PromptEditor } from "./prompt-editor";
import type { Prompt, PromptFormData } from "./types"; import type { Prompt } from "./types";
describe("PromptEditor", () => { describe("PromptEditor", () => {
const mockOnSave = jest.fn(); const mockOnSave = jest.fn();

View file

@ -12,6 +12,20 @@ jest.mock("next/navigation", () => ({
}), }),
})); }));
// Mock NextAuth
jest.mock("next-auth/react", () => ({
useSession: () => ({
data: {
accessToken: "mock-access-token",
user: {
id: "mock-user-id",
email: "test@example.com",
},
},
status: "authenticated",
}),
}));
describe("VectorStoreDetailView", () => { describe("VectorStoreDetailView", () => {
const defaultProps = { const defaultProps = {
store: null, store: null,

View file

@ -1,16 +1,18 @@
"use client"; "use client";
import { useRouter } from "next/navigation"; import { useRouter } from "next/navigation";
import { useState, useEffect } from "react";
import type { VectorStore } from "llama-stack-client/resources/vector-stores/vector-stores"; import type { VectorStore } from "llama-stack-client/resources/vector-stores/vector-stores";
import type { VectorStoreFile } from "llama-stack-client/resources/vector-stores/files"; import type { VectorStoreFile } from "llama-stack-client/resources/vector-stores/files";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Skeleton } from "@/components/ui/skeleton"; import { Skeleton } from "@/components/ui/skeleton";
import { Button } from "@/components/ui/button"; import { Button } from "@/components/ui/button";
import { useAuthClient } from "@/hooks/use-auth-client";
import { Edit2, Trash2, X } from "lucide-react";
import { import {
DetailLoadingView, DetailLoadingView,
DetailErrorView, DetailErrorView,
DetailNotFoundView, DetailNotFoundView,
DetailLayout,
PropertiesCard, PropertiesCard,
PropertyItem, PropertyItem,
} from "@/components/layout/detail-layout"; } from "@/components/layout/detail-layout";
@ -23,6 +25,7 @@ import {
TableHeader, TableHeader,
TableRow, TableRow,
} from "@/components/ui/table"; } from "@/components/ui/table";
import { VectorStoreEditor, VectorStoreFormData } from "./vector-store-editor";
interface VectorStoreDetailViewProps { interface VectorStoreDetailViewProps {
store: VectorStore | null; store: VectorStore | null;
@ -43,21 +46,122 @@ export function VectorStoreDetailView({
errorFiles, errorFiles,
id, id,
}: VectorStoreDetailViewProps) { }: VectorStoreDetailViewProps) {
const title = "Vector Store Details";
const router = useRouter(); const router = useRouter();
const client = useAuthClient();
const [isDeleting, setIsDeleting] = useState(false);
const [showEditModal, setShowEditModal] = useState(false);
const [modalError, setModalError] = useState<string | null>(null);
const [showSuccessState, setShowSuccessState] = useState(false);
// Handle ESC key to close modal
useEffect(() => {
const handleEscape = (event: KeyboardEvent) => {
if (event.key === "Escape" && showEditModal) {
handleCancel();
}
};
document.addEventListener("keydown", handleEscape);
return () => document.removeEventListener("keydown", handleEscape);
}, [showEditModal]);
const handleFileClick = (fileId: string) => { const handleFileClick = (fileId: string) => {
router.push(`/logs/vector-stores/${id}/files/${fileId}`); router.push(`/logs/vector-stores/${id}/files/${fileId}`);
}; };
const handleEditVectorStore = () => {
setShowEditModal(true);
setModalError(null);
setShowSuccessState(false);
};
const handleCancel = () => {
setShowEditModal(false);
setModalError(null);
setShowSuccessState(false);
};
const handleSaveVectorStore = async (formData: VectorStoreFormData) => {
try {
setModalError(null);
// Update existing vector store (same logic as list page)
const updateParams: {
name?: string;
extra_body?: Record<string, unknown>;
} = {};
// Only include fields that have changed or are provided
if (formData.name && formData.name !== store?.name) {
updateParams.name = formData.name;
}
// Add all parameters to extra_body (except provider_id which can't be changed)
const extraBody: Record<string, unknown> = {};
if (formData.embedding_model) {
extraBody.embedding_model = formData.embedding_model;
}
if (formData.embedding_dimension) {
extraBody.embedding_dimension = formData.embedding_dimension;
}
if (Object.keys(extraBody).length > 0) {
updateParams.extra_body = extraBody;
}
await client.vectorStores.update(id, updateParams);
// Show success state
setShowSuccessState(true);
setModalError(
"✅ Vector store updated successfully! You can close this modal and refresh the page to see changes."
);
} catch (err: unknown) {
console.error("Failed to update vector store:", err);
const errorMessage =
err instanceof Error ? err.message : "Failed to update vector store";
setModalError(errorMessage);
}
};
const handleDeleteVectorStore = async () => {
if (
!confirm(
"Are you sure you want to delete this vector store? This action cannot be undone."
)
) {
return;
}
setIsDeleting(true);
try {
await client.vectorStores.delete(id);
// Redirect to the vector stores list after successful deletion
router.push("/logs/vector-stores");
} catch (err: unknown) {
console.error("Failed to delete vector store:", err);
const errorMessage = err instanceof Error ? err.message : "Unknown error";
alert(`Failed to delete vector store: ${errorMessage}`);
} finally {
setIsDeleting(false);
}
};
if (errorStore) { if (errorStore) {
return <DetailErrorView title={title} id={id} error={errorStore} />; return (
<DetailErrorView
title="Vector Store Details"
id={id}
error={errorStore}
/>
);
} }
if (isLoadingStore) { if (isLoadingStore) {
return <DetailLoadingView title={title} />; return <DetailLoadingView />;
} }
if (!store) { if (!store) {
return <DetailNotFoundView title={title} id={id} />; return <DetailNotFoundView title="Vector Store Details" id={id} />;
} }
const mainContent = ( const mainContent = (
@ -138,6 +242,73 @@ export function VectorStoreDetailView({
); );
return ( return (
<DetailLayout title={title} mainContent={mainContent} sidebar={sidebar} /> <>
<div className="flex items-center justify-between mb-6">
<h1 className="text-2xl font-bold">Vector Store Details</h1>
<div className="flex gap-2">
<Button
variant="outline"
onClick={handleEditVectorStore}
disabled={isDeleting}
>
<Edit2 className="h-4 w-4 mr-2" />
Edit
</Button>
<Button
variant="destructive"
onClick={handleDeleteVectorStore}
disabled={isDeleting}
>
{isDeleting ? (
"Deleting..."
) : (
<>
<Trash2 className="h-4 w-4 mr-2" />
Delete
</>
)}
</Button>
</div>
</div>
<div className="flex flex-col md:flex-row gap-6">
<div className="flex-grow md:w-2/3 space-y-6">{mainContent}</div>
<div className="md:w-1/3">{sidebar}</div>
</div>
{/* Edit Vector Store Modal */}
{showEditModal && (
<div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50">
<div className="bg-background border rounded-lg shadow-lg max-w-2xl w-full mx-4 max-h-[90vh] overflow-hidden">
<div className="p-6 border-b flex items-center justify-between">
<h2 className="text-2xl font-bold">Edit Vector Store</h2>
<Button
variant="ghost"
size="sm"
onClick={handleCancel}
className="p-1 h-auto"
>
<X className="h-4 w-4" />
</Button>
</div>
<div className="p-6 overflow-y-auto max-h-[calc(90vh-120px)]">
<VectorStoreEditor
onSave={handleSaveVectorStore}
onCancel={handleCancel}
error={modalError}
showSuccessState={showSuccessState}
isEditing={true}
initialData={{
name: store?.name || "",
embedding_model: store?.metadata?.embedding_model || "",
embedding_dimension:
store?.metadata?.embedding_dimension || 768,
provider_id: store?.metadata?.provider_id || "",
}}
/>
</div>
</div>
</div>
)}
</>
); );
} }

View file

@ -0,0 +1,235 @@
"use client";
import { useState, useEffect } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { Card, CardContent } from "@/components/ui/card";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { useAuthClient } from "@/hooks/use-auth-client";
import type { Model } from "llama-stack-client/resources/models";
export interface VectorStoreFormData {
name: string;
embedding_model?: string;
embedding_dimension?: number;
provider_id?: string;
}
interface VectorStoreEditorProps {
onSave: (formData: VectorStoreFormData) => Promise<void>;
onCancel: () => void;
error?: string | null;
initialData?: VectorStoreFormData;
showSuccessState?: boolean;
isEditing?: boolean;
}
export function VectorStoreEditor({
onSave,
onCancel,
error,
initialData,
showSuccessState,
isEditing = false,
}: VectorStoreEditorProps) {
const client = useAuthClient();
const [formData, setFormData] = useState<VectorStoreFormData>(
initialData || {
name: "",
embedding_model: "",
embedding_dimension: 768,
provider_id: "",
}
);
const [loading, setLoading] = useState(false);
const [models, setModels] = useState<Model[]>([]);
const [modelsLoading, setModelsLoading] = useState(true);
const [modelsError, setModelsError] = useState<string | null>(null);
const embeddingModels = models.filter(
model => model.custom_metadata?.model_type === "embedding"
);
useEffect(() => {
const fetchModels = async () => {
try {
setModelsLoading(true);
setModelsError(null);
const modelList = await client.models.list();
setModels(modelList);
// Set default embedding model if available
const embeddingModelsList = modelList.filter(model => {
return model.custom_metadata?.model_type === "embedding";
});
if (embeddingModelsList.length > 0 && !formData.embedding_model) {
setFormData(prev => ({
...prev,
embedding_model: embeddingModelsList[0].id,
}));
}
} catch (err) {
console.error("Failed to load models:", err);
setModelsError(
err instanceof Error ? err.message : "Failed to load models"
);
} finally {
setModelsLoading(false);
}
};
fetchModels();
}, [client]);
const handleSubmit = async (e: React.FormEvent) => {
e.preventDefault();
setLoading(true);
try {
await onSave(formData);
} finally {
setLoading(false);
}
};
return (
<Card>
<CardContent className="pt-6">
<form onSubmit={handleSubmit} className="space-y-4">
<div className="space-y-2">
<Label htmlFor="name">Name</Label>
<Input
id="name"
value={formData.name}
onChange={e => setFormData({ ...formData, name: e.target.value })}
placeholder="Enter vector store name"
required
/>
</div>
<div className="space-y-2">
<Label htmlFor="embedding_model">Embedding Model (Optional)</Label>
{modelsLoading ? (
<div className="text-sm text-muted-foreground">
Loading models... ({models.length} loaded)
</div>
) : modelsError ? (
<div className="text-sm text-destructive">
Error: {modelsError}
</div>
) : embeddingModels.length === 0 ? (
<div className="text-sm text-muted-foreground">
No embedding models available ({models.length} total models)
</div>
) : (
<Select
value={formData.embedding_model}
onValueChange={value =>
setFormData({ ...formData, embedding_model: value })
}
>
<SelectTrigger>
<SelectValue placeholder="Select Embedding Model" />
</SelectTrigger>
<SelectContent>
{embeddingModels.map((model, index) => (
<SelectItem key={model.id} value={model.id}>
{model.id}
</SelectItem>
))}
</SelectContent>
</Select>
)}
{formData.embedding_model && (
<p className="text-xs text-muted-foreground mt-1">
Dimension:{" "}
{embeddingModels.find(m => m.id === formData.embedding_model)
?.custom_metadata?.embedding_dimension || "Unknown"}
</p>
)}
</div>
<div className="space-y-2">
<Label htmlFor="embedding_dimension">
Embedding Dimension (Optional)
</Label>
<Input
id="embedding_dimension"
type="number"
value={formData.embedding_dimension}
onChange={e =>
setFormData({
...formData,
embedding_dimension: parseInt(e.target.value) || 768,
})
}
placeholder="768"
/>
</div>
<div className="space-y-2">
<Label htmlFor="provider_id">
Provider ID {isEditing ? "(Read-only)" : "(Optional)"}
</Label>
<Input
id="provider_id"
value={formData.provider_id}
onChange={e =>
setFormData({ ...formData, provider_id: e.target.value })
}
placeholder="e.g., faiss, chroma, sqlite"
disabled={isEditing}
/>
{isEditing && (
<p className="text-xs text-muted-foreground">
Provider ID cannot be changed after vector store creation
</p>
)}
</div>
{error && (
<div
className={`text-sm p-3 rounded ${
error.startsWith("✅")
? "text-green-700 bg-green-50 border border-green-200"
: "text-destructive bg-destructive/10"
}`}
>
{error}
</div>
)}
<div className="flex gap-2 pt-4">
{showSuccessState ? (
<Button type="button" onClick={onCancel}>
Close
</Button>
) : (
<>
<Button type="submit" disabled={loading}>
{loading
? initialData
? "Updating..."
: "Creating..."
: initialData
? "Update Vector Store"
: "Create Vector Store"}
</Button>
<Button type="button" variant="outline" onClick={onCancel}>
Cancel
</Button>
</>
)}
</div>
</form>
</CardContent>
</Card>
);
}

View file

@ -34,9 +34,35 @@ export class ContentsAPI {
async getFileContents( async getFileContents(
vectorStoreId: string, vectorStoreId: string,
fileId: string fileId: string,
includeEmbeddings: boolean = true,
includeMetadata: boolean = true
): Promise<VectorStoreContentsResponse> { ): Promise<VectorStoreContentsResponse> {
return this.client.vectorStores.files.content(vectorStoreId, fileId); try {
// Use query parameters to pass embeddings and metadata flags (OpenAI-compatible pattern)
const extraQuery: Record<string, boolean> = {};
if (includeEmbeddings) {
extraQuery.include_embeddings = true;
}
if (includeMetadata) {
extraQuery.include_metadata = true;
}
const result = await this.client.vectorStores.files.content(
vectorStoreId,
fileId,
{
query: {
include_embeddings: includeEmbeddings,
include_metadata: includeMetadata,
},
}
);
return result;
} catch (error) {
console.error("ContentsAPI.getFileContents error:", error);
throw error;
}
} }
async getContent( async getContent(
@ -70,11 +96,15 @@ export class ContentsAPI {
order?: string; order?: string;
after?: string; after?: string;
before?: string; before?: string;
includeEmbeddings?: boolean;
includeMetadata?: boolean;
} }
): Promise<VectorStoreListContentsResponse> { ): Promise<VectorStoreListContentsResponse> {
const fileContents = await this.client.vectorStores.files.content( const fileContents = await this.getFileContents(
vectorStoreId, vectorStoreId,
fileId fileId,
options?.includeEmbeddings ?? true,
options?.includeMetadata ?? true
); );
const contentItems: VectorStoreContentItem[] = []; const contentItems: VectorStoreContentItem[] = [];
@ -82,7 +112,7 @@ export class ContentsAPI {
const rawContent = content as Record<string, unknown>; const rawContent = content as Record<string, unknown>;
// Extract actual fields from the API response // Extract actual fields from the API response
const embedding = rawContent.embedding || undefined; const embedding = rawContent.embedding as number[] | undefined;
const created_timestamp = const created_timestamp =
rawContent.created_timestamp || rawContent.created_timestamp ||
rawContent.created_at || rawContent.created_at ||

View file

@ -11,6 +11,7 @@ import pytest
from llama_stack_client import BadRequestError from llama_stack_client import BadRequestError
from openai import BadRequestError as OpenAIBadRequestError from openai import BadRequestError as OpenAIBadRequestError
from llama_stack.apis.files import ExpiresAfter
from llama_stack.apis.vector_io import Chunk from llama_stack.apis.vector_io import Chunk
from llama_stack.core.library_client import LlamaStackAsLibraryClient from llama_stack.core.library_client import LlamaStackAsLibraryClient
from llama_stack.log import get_logger from llama_stack.log import get_logger
@ -1604,3 +1605,97 @@ def test_openai_vector_store_embedding_config_from_metadata(
assert "metadata_config_store" in store_names assert "metadata_config_store" in store_names
assert "consistent_config_store" in store_names assert "consistent_config_store" in store_names
@vector_provider_wrapper
def test_openai_vector_store_file_contents_with_extra_query(
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
):
"""Test that vector store file contents endpoint supports extra_query parameter."""
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
compat_client = compat_client_with_empty_stores
# Create a vector store
vector_store = compat_client.vector_stores.create(
name="test_extra_query_store",
extra_body={
"embedding_model": embedding_model_id,
"provider_id": vector_io_provider_id,
},
)
# Create and attach a file
test_content = b"This is test content for extra_query validation."
with BytesIO(test_content) as file_buffer:
file_buffer.name = "test_extra_query.txt"
file = compat_client.files.create(
file=file_buffer,
purpose="assistants",
expires_after=ExpiresAfter(anchor="created_at", seconds=86400),
)
file_attach_response = compat_client.vector_stores.files.create(
vector_store_id=vector_store.id,
file_id=file.id,
extra_body={"embedding_model": embedding_model_id},
)
assert file_attach_response.status == "completed"
# Wait for processing
time.sleep(2)
# Test that extra_query parameter is accepted and processed
content_with_extra_query = compat_client.vector_stores.files.content(
vector_store_id=vector_store.id,
file_id=file.id,
extra_query={"include_embeddings": True, "include_metadata": True},
)
# Test without extra_query for comparison
content_without_extra_query = compat_client.vector_stores.files.content(
vector_store_id=vector_store.id,
file_id=file.id,
)
# Validate that both calls succeed
assert content_with_extra_query is not None
assert content_without_extra_query is not None
assert len(content_with_extra_query.content) > 0
assert len(content_without_extra_query.content) > 0
# Validate that extra_query parameter is processed correctly
# Both should have the embedding/metadata fields available (may be None based on flags)
first_chunk_with_flags = content_with_extra_query.content[0]
first_chunk_without_flags = content_without_extra_query.content[0]
# The key validation: extra_query fields are present in the response
# Handle both dict and object responses (different clients may return different formats)
def has_field(obj, field):
if isinstance(obj, dict):
return field in obj
else:
return hasattr(obj, field)
# Validate that all expected fields are present in both responses
expected_fields = ["embedding", "chunk_metadata", "metadata", "text"]
for field in expected_fields:
assert has_field(first_chunk_with_flags, field), f"Field '{field}' missing from response with extra_query"
assert has_field(first_chunk_without_flags, field), f"Field '{field}' missing from response without extra_query"
# Validate content is the same
def get_field(obj, field):
if isinstance(obj, dict):
return obj[field]
else:
return getattr(obj, field)
assert get_field(first_chunk_with_flags, "text") == test_content.decode("utf-8")
assert get_field(first_chunk_without_flags, "text") == test_content.decode("utf-8")
with_flags_embedding = get_field(first_chunk_with_flags, "embedding")
without_flags_embedding = get_field(first_chunk_without_flags, "embedding")
# Validate that embeddings are included when requested and excluded when not requested
assert with_flags_embedding is not None, "Embeddings should be included when include_embeddings=True"
assert len(with_flags_embedding) > 0, "Embedding should be a non-empty list"
assert without_flags_embedding is None, "Embeddings should not be included when include_embeddings=False"

View file

@ -55,3 +55,65 @@ async def test_create_vector_stores_multiple_providers_missing_provider_id_error
with pytest.raises(ValueError, match="Multiple vector_io providers available"): with pytest.raises(ValueError, match="Multiple vector_io providers available"):
await router.openai_create_vector_store(request) await router.openai_create_vector_store(request)
async def test_update_vector_store_provider_id_change_fails():
"""Test that updating a vector store with a different provider_id fails with clear error."""
mock_routing_table = Mock()
# Mock an existing vector store with provider_id "faiss"
mock_existing_store = Mock()
mock_existing_store.provider_id = "inline::faiss"
mock_existing_store.identifier = "vs_123"
mock_routing_table.get_object_by_identifier = AsyncMock(return_value=mock_existing_store)
mock_routing_table.get_provider_impl = AsyncMock(
return_value=Mock(openai_update_vector_store=AsyncMock(return_value=Mock(id="vs_123")))
)
router = VectorIORouter(mock_routing_table)
# Try to update with different provider_id in metadata - this should fail
with pytest.raises(ValueError, match="provider_id cannot be changed after vector store creation"):
await router.openai_update_vector_store(
vector_store_id="vs_123",
name="updated_name",
metadata={"provider_id": "inline::sqlite"}, # Different provider_id
)
# Verify the existing store was looked up to check provider_id
mock_routing_table.get_object_by_identifier.assert_called_once_with("vector_store", "vs_123")
# Provider should not be called since validation failed
mock_routing_table.get_provider_impl.assert_not_called()
async def test_update_vector_store_same_provider_id_succeeds():
"""Test that updating a vector store with the same provider_id succeeds."""
mock_routing_table = Mock()
# Mock an existing vector store with provider_id "faiss"
mock_existing_store = Mock()
mock_existing_store.provider_id = "inline::faiss"
mock_existing_store.identifier = "vs_123"
mock_routing_table.get_object_by_identifier = AsyncMock(return_value=mock_existing_store)
mock_routing_table.get_provider_impl = AsyncMock(
return_value=Mock(openai_update_vector_store=AsyncMock(return_value=Mock(id="vs_123")))
)
router = VectorIORouter(mock_routing_table)
# Update with same provider_id should succeed
await router.openai_update_vector_store(
vector_store_id="vs_123",
name="updated_name",
metadata={"provider_id": "inline::faiss"}, # Same provider_id
)
# Verify the provider update method was called
mock_routing_table.get_provider_impl.assert_called_once_with("vs_123")
provider = await mock_routing_table.get_provider_impl("vs_123")
provider.openai_update_vector_store.assert_called_once_with(
vector_store_id="vs_123", name="updated_name", expires_after=None, metadata={"provider_id": "inline::faiss"}
)

View file

@ -0,0 +1,86 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from unittest.mock import AsyncMock, Mock
from fastapi import Request
from llama_stack.core.server.query_params_middleware import QueryParamsMiddleware
class TestQueryParamsMiddleware:
"""Test cases for the QueryParamsMiddleware."""
async def test_extracts_query_params_for_vector_store_content(self):
"""Test that middleware extracts query params for vector store content endpoints."""
middleware = QueryParamsMiddleware(Mock())
request = Mock(spec=Request)
request.method = "GET"
# Mock the URL properly
mock_url = Mock()
mock_url.path = "/v1/vector_stores/vs_123/files/file_456/content"
request.url = mock_url
request.query_params = {"include_embeddings": "true", "include_metadata": "false"}
# Create a fresh state object without any attributes
class MockState:
pass
request.state = MockState()
await middleware.dispatch(request, AsyncMock())
assert hasattr(request.state, "extra_query")
assert request.state.extra_query == {"include_embeddings": True, "include_metadata": False}
async def test_ignores_non_vector_store_endpoints(self):
"""Test that middleware ignores non-vector store endpoints."""
middleware = QueryParamsMiddleware(Mock())
request = Mock(spec=Request)
request.method = "GET"
# Mock the URL properly
mock_url = Mock()
mock_url.path = "/v1/inference/chat_completion"
request.url = mock_url
request.query_params = {"include_embeddings": "true"}
# Create a fresh state object without any attributes
class MockState:
pass
request.state = MockState()
await middleware.dispatch(request, AsyncMock())
assert not hasattr(request.state, "extra_query")
async def test_handles_json_parsing(self):
"""Test that middleware correctly parses JSON values and handles invalid JSON."""
middleware = QueryParamsMiddleware(Mock())
request = Mock(spec=Request)
request.method = "GET"
# Mock the URL properly
mock_url = Mock()
mock_url.path = "/v1/vector_stores/vs_123/files/file_456/content"
request.url = mock_url
request.query_params = {"config": '{"key": "value"}', "invalid": "not-json{", "number": "42"}
# Create a fresh state object without any attributes
class MockState:
pass
request.state = MockState()
await middleware.dispatch(request, AsyncMock())
expected = {"config": {"key": "value"}, "invalid": "not-json{", "number": 42}
assert request.state.extra_query == expected

View file

@ -104,12 +104,18 @@ async def test_paginated_response_url_setting():
route_handler = create_dynamic_typed_route(mock_api_method, "get", "/test/route") route_handler = create_dynamic_typed_route(mock_api_method, "get", "/test/route")
# Mock minimal request # Mock minimal request with proper state object
request = MagicMock() request = MagicMock()
request.scope = {"user_attributes": {}, "principal": ""} request.scope = {"user_attributes": {}, "principal": ""}
request.headers = {} request.headers = {}
request.body = AsyncMock(return_value=b"") request.body = AsyncMock(return_value=b"")
# Create a simple state object without auto-generating attributes
class MockState:
pass
request.state = MockState()
result = await route_handler(request) result = await route_handler(request)
assert isinstance(result, PaginatedResponse) assert isinstance(result, PaginatedResponse)