undo vector_store_id in params

This commit is contained in:
Ashwin Bharambe 2025-10-12 18:48:55 -07:00
parent 33d7119708
commit e6428492e6
11 changed files with 94 additions and 59 deletions

View file

@ -2623,7 +2623,17 @@
],
"summary": "Create a vector store file batch.",
"description": "Create a vector store file batch.\nGenerate an OpenAI-compatible vector store file batch for the given vector store.",
"parameters": [],
"parameters": [
{
"name": "vector_store_id",
"in": "path",
"description": "The ID of the vector store to create the file batch for.",
"required": true,
"schema": {
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
@ -12321,10 +12331,6 @@
"OpenAICreateVectorStoreFileBatchRequestWithExtraBody": {
"type": "object",
"properties": {
"vector_store_id": {
"type": "string",
"description": "The ID of the vector store to create the file batch for"
},
"file_ids": {
"type": "array",
"items": {
@ -12365,7 +12371,6 @@
},
"additionalProperties": false,
"required": [
"vector_store_id",
"file_ids"
],
"title": "OpenAICreateVectorStoreFileBatchRequestWithExtraBody",

View file

@ -1932,7 +1932,14 @@ paths:
Generate an OpenAI-compatible vector store file batch for the given vector
store.
parameters: []
parameters:
- name: vector_store_id
in: path
description: >-
The ID of the vector store to create the file batch for.
required: true
schema:
type: string
requestBody:
content:
application/json:
@ -9326,10 +9333,6 @@ components:
"OpenAICreateVectorStoreFileBatchRequestWithExtraBody":
type: object
properties:
vector_store_id:
type: string
description: >-
The ID of the vector store to create the file batch for
file_ids:
type: array
items:
@ -9355,7 +9358,6 @@ components:
auto
additionalProperties: false
required:
- vector_store_id
- file_ids
title: >-
OpenAICreateVectorStoreFileBatchRequestWithExtraBody

View file

@ -3357,7 +3357,17 @@
],
"summary": "Create a vector store file batch.",
"description": "Create a vector store file batch.\nGenerate an OpenAI-compatible vector store file batch for the given vector store.",
"parameters": [],
"parameters": [
{
"name": "vector_store_id",
"in": "path",
"description": "The ID of the vector store to create the file batch for.",
"required": true,
"schema": {
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
@ -12847,10 +12857,6 @@
"OpenAICreateVectorStoreFileBatchRequestWithExtraBody": {
"type": "object",
"properties": {
"vector_store_id": {
"type": "string",
"description": "The ID of the vector store to create the file batch for"
},
"file_ids": {
"type": "array",
"items": {
@ -12891,7 +12897,6 @@
},
"additionalProperties": false,
"required": [
"vector_store_id",
"file_ids"
],
"title": "OpenAICreateVectorStoreFileBatchRequestWithExtraBody",

View file

@ -2553,7 +2553,14 @@ paths:
Generate an OpenAI-compatible vector store file batch for the given vector
store.
parameters: []
parameters:
- name: vector_store_id
in: path
description: >-
The ID of the vector store to create the file batch for.
required: true
schema:
type: string
requestBody:
content:
application/json:
@ -9791,10 +9798,6 @@ components:
"OpenAICreateVectorStoreFileBatchRequestWithExtraBody":
type: object
properties:
vector_store_id:
type: string
description: >-
The ID of the vector store to create the file batch for
file_ids:
type: array
items:
@ -9820,7 +9823,6 @@ components:
auto
additionalProperties: false
required:
- vector_store_id
- file_ids
title: >-
OpenAICreateVectorStoreFileBatchRequestWithExtraBody

View file

@ -3357,7 +3357,17 @@
],
"summary": "Create a vector store file batch.",
"description": "Create a vector store file batch.\nGenerate an OpenAI-compatible vector store file batch for the given vector store.",
"parameters": [],
"parameters": [
{
"name": "vector_store_id",
"in": "path",
"description": "The ID of the vector store to create the file batch for.",
"required": true,
"schema": {
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
@ -14856,10 +14866,6 @@
"OpenAICreateVectorStoreFileBatchRequestWithExtraBody": {
"type": "object",
"properties": {
"vector_store_id": {
"type": "string",
"description": "The ID of the vector store to create the file batch for"
},
"file_ids": {
"type": "array",
"items": {
@ -14900,7 +14906,6 @@
},
"additionalProperties": false,
"required": [
"vector_store_id",
"file_ids"
],
"title": "OpenAICreateVectorStoreFileBatchRequestWithExtraBody",

View file

@ -2556,7 +2556,14 @@ paths:
Generate an OpenAI-compatible vector store file batch for the given vector
store.
parameters: []
parameters:
- name: vector_store_id
in: path
description: >-
The ID of the vector store to create the file batch for.
required: true
schema:
type: string
requestBody:
content:
application/json:
@ -11236,10 +11243,6 @@ components:
"OpenAICreateVectorStoreFileBatchRequestWithExtraBody":
type: object
properties:
vector_store_id:
type: string
description: >-
The ID of the vector store to create the file batch for
file_ids:
type: array
items:
@ -11265,7 +11268,6 @@ components:
auto
additionalProperties: false
required:
- vector_store_id
- file_ids
title: >-
OpenAICreateVectorStoreFileBatchRequestWithExtraBody

View file

@ -491,13 +491,11 @@ class OpenAICreateVectorStoreRequestWithExtraBody(BaseModel, extra="allow"):
class OpenAICreateVectorStoreFileBatchRequestWithExtraBody(BaseModel, extra="allow"):
"""Request to create a vector store file batch with extra_body support.
:param vector_store_id: The ID of the vector store to create the file batch for
:param file_ids: A list of File IDs that the vector store should use
:param attributes: (Optional) Key-value attributes to store with the files
:param chunking_strategy: (Optional) The chunking strategy used to chunk the file(s). Defaults to auto
"""
vector_store_id: str
file_ids: list[str]
attributes: dict[str, Any] | None = None
chunking_strategy: VectorStoreChunkingStrategy | None = None
@ -849,11 +847,13 @@ class VectorIO(Protocol):
)
async def openai_create_vector_store_file_batch(
self,
vector_store_id: str,
params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
) -> VectorStoreFileBatchObject:
"""Create a vector store file batch.
Generate an OpenAI-compatible vector store file batch for the given vector store.
:param vector_store_id: The ID of the vector store to create the file batch for.
:returns: A VectorStoreFileBatchObject representing the created file batch.
"""
...

View file

@ -513,6 +513,14 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
# Strip NOT_GIVENs to use the defaults in signature
body = {k: v for k, v in body.items() if v is not NOT_GIVEN}
# Check if there's an unwrapped body parameter among multiple parameters
# (e.g., path param + body param like: vector_store_id: str, params: Annotated[Model, Body(...)])
unwrapped_body_param = None
for param in params_list:
if is_unwrapped_body_param(param.annotation):
unwrapped_body_param = param
break
# Convert parameters to Pydantic models where needed
converted_body = {}
for param_name, param in sig.parameters.items():
@ -522,5 +530,11 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
converted_body[param_name] = value
else:
converted_body[param_name] = convert_to_pydantic(param.annotation, value)
elif unwrapped_body_param and param.name == unwrapped_body_param.name:
# This is the unwrapped body param - construct it from remaining body keys
base_type = get_args(param.annotation)[0]
# Extract only the keys that aren't already used by other params
remaining_keys = {k: v for k, v in body.items() if k not in converted_body}
converted_body[param.name] = base_type(**remaining_keys)
return converted_body

View file

@ -383,13 +383,14 @@ class VectorIORouter(VectorIO):
async def openai_create_vector_store_file_batch(
self,
vector_store_id: str,
params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
) -> VectorStoreFileBatchObject:
logger.debug(
f"VectorIORouter.openai_create_vector_store_file_batch: {params.vector_store_id}, {len(params.file_ids)} files"
f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(params.file_ids)} files"
)
provider = await self.routing_table.get_provider_impl(params.vector_store_id)
return await provider.openai_create_vector_store_file_batch(params)
provider = await self.routing_table.get_provider_impl(vector_store_id)
return await provider.openai_create_vector_store_file_batch(vector_store_id, params)
async def openai_retrieve_vector_store_file_batch(
self,

View file

@ -978,10 +978,10 @@ class OpenAIVectorStoreMixin(ABC):
async def openai_create_vector_store_file_batch(
self,
vector_store_id: str,
params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
) -> VectorStoreFileBatchObject:
"""Create a vector store file batch."""
vector_store_id = params.vector_store_id
if vector_store_id not in self.openai_vector_stores:
raise VectorStoreNotFoundError(vector_store_id)

View file

@ -327,7 +327,7 @@ async def test_create_vector_store_file_batch(vector_io_adapter):
vector_io_adapter._process_file_batch_async = AsyncMock()
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
assert batch.vector_store_id == store_id
@ -354,7 +354,7 @@ async def test_retrieve_vector_store_file_batch(vector_io_adapter):
# Create batch first
created_batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
# Retrieve batch
@ -387,7 +387,7 @@ async def test_cancel_vector_store_file_batch(vector_io_adapter):
# Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
# Cancel batch
@ -432,7 +432,7 @@ async def test_list_files_in_vector_store_file_batch(vector_io_adapter):
# Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
# List files
@ -451,9 +451,8 @@ async def test_file_batch_validation_errors(vector_io_adapter):
# Test nonexistent vector store
with pytest.raises(VectorStoreNotFoundError):
await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(
vector_store_id="nonexistent", file_ids=["file_1"]
),
vector_store_id="nonexistent",
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"]),
)
# Setup store for remaining tests
@ -470,7 +469,7 @@ async def test_file_batch_validation_errors(vector_io_adapter):
# Test wrong vector store for batch
vector_io_adapter.openai_attach_file_to_vector_store = AsyncMock()
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=["file_1"])
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"])
)
# Create wrong_store so it exists but the batch doesn't belong to it
@ -517,7 +516,7 @@ async def test_file_batch_pagination(vector_io_adapter):
# Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
# Test pagination with limit
@ -589,7 +588,7 @@ async def test_file_batch_status_filtering(vector_io_adapter):
# Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
# Test filtering by completed status
@ -631,7 +630,7 @@ async def test_cancel_completed_batch_fails(vector_io_adapter):
# Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
# Manually update status to completed
@ -665,7 +664,7 @@ async def test_file_batch_persistence_across_restarts(vector_io_adapter):
# Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
batch_id = batch.id
@ -720,7 +719,7 @@ async def test_cancelled_batch_persists_in_storage(vector_io_adapter):
# Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
batch_id = batch.id
@ -767,10 +766,10 @@ async def test_only_in_progress_batches_resumed(vector_io_adapter):
# Create multiple batches
batch1 = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=["file_1"])
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"])
)
batch2 = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=["file_2"])
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_2"])
)
# Complete one batch (should persist with completed status)
@ -783,7 +782,7 @@ async def test_only_in_progress_batches_resumed(vector_io_adapter):
# Create a third batch that stays in progress
batch3 = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=["file_3"])
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_3"])
)
# Simulate restart - clear memory and reload from persistence
@ -944,7 +943,7 @@ async def test_max_concurrent_files_per_batch(vector_io_adapter):
file_ids = [f"file_{i}" for i in range(8)] # 8 files, but limit should be 5
batch = await vector_io_adapter.openai_create_vector_store_file_batch(
params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(vector_store_id=store_id, file_ids=file_ids)
vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
)
# Give time for the semaphore logic to start processing files