From cc64093ae4e0ecf7aa351b4389be26cbee86db93 Mon Sep 17 00:00:00 2001 From: slekkala1 Date: Tue, 30 Sep 2025 12:07:33 -0700 Subject: [PATCH] feat(api): Add Vector Store File batches api stub (#3615) # What does this PR do? Adding api stubs for vector store file batches apis https://github.com/llamastack/llama-stack/issues/3533 API Ref: https://platform.openai.com/docs/api-reference/vector-stores-file-batches ## Test Plan CI --- docs/static/llama-stack-spec.html | 454 ++++++++++++++++-- docs/static/llama-stack-spec.yaml | 360 ++++++++++++-- llama_stack/apis/vector_io/vector_io.py | 179 ++++++- llama_stack/core/routers/vector_io.py | 69 ++- .../utils/memory/openai_vector_store_mixin.py | 61 ++- 5 files changed, 1038 insertions(+), 85 deletions(-) diff --git a/docs/static/llama-stack-spec.html b/docs/static/llama-stack-spec.html index c755af948..97671f084 100644 --- a/docs/static/llama-stack-spec.html +++ b/docs/static/llama-stack-spec.html @@ -4741,6 +4741,59 @@ } } }, + "/v1/vector_stores/{vector_store_id}/file_batches/{batch_id}/cancel": { + "post": { + "responses": { + "200": { + "description": "A VectorStoreFileBatchObject representing the cancelled file batch.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFileBatchObject" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "summary": "Cancels a vector store file batch.", + "description": "Cancels a vector store file batch.", + "parameters": [ + { + "name": "batch_id", + "in": "path", + "description": "The ID of the file batch to cancel.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "vector_store_id", + "in": "path", + "description": "The ID of the vector store containing the file batch.", + "required": true, + "schema": { + "type": "string" + } + } + ] + } + }, "/v1/completions": { "post": { "responses": { @@ -4898,6 +4951,60 @@ } } }, + "/v1/vector_stores/{vector_store_id}/file_batches": { + "post": { + "responses": { + "200": { + "description": "A VectorStoreFileBatchObject representing the created file batch.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFileBatchObject" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "summary": "Create a vector store file batch.", + "description": "Create a vector store file batch.", + "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": { + "schema": { + "$ref": "#/components/schemas/OpenaiCreateVectorStoreFileBatchRequest" + } + } + }, + "required": true + } + } + }, "/v1/files/{file_id}": { "get": { "responses": { @@ -5460,6 +5567,104 @@ } } }, + "/v1/vector_stores/{vector_store_id}/file_batches/{batch_id}/files": { + "get": { + "responses": { + "200": { + "description": "A VectorStoreFilesListInBatchResponse containing the list of files in the batch.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFilesListInBatchResponse" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "summary": "Returns a list of vector store files in a batch.", + "description": "Returns a list of vector store files in a batch.", + "parameters": [ + { + "name": "batch_id", + "in": "path", + "description": "The ID of the file batch to list files from.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "vector_store_id", + "in": "path", + "description": "The ID of the vector store containing the file batch.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "after", + "in": "query", + "description": "A cursor for use in pagination. `after` is an object ID that defines your place in the list.", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "before", + "in": "query", + "description": "A cursor for use in pagination. `before` is an object ID that defines your place in the list.", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "filter", + "in": "query", + "description": "Filter by file status. One of in_progress, completed, failed, cancelled.", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "limit", + "in": "query", + "description": "A limit on the number of objects to be returned. Limit can range between 1 and 100, and the default is 20.", + "required": false, + "schema": { + "type": "integer" + } + }, + { + "name": "order", + "in": "query", + "description": "Sort order by the `created_at` timestamp of the objects. `asc` for ascending order and `desc` for descending order.", + "required": false, + "schema": { + "type": "string" + } + } + ] + } + }, "/v1/files/{file_id}/content": { "get": { "responses": { @@ -5504,6 +5709,59 @@ ] } }, + "/v1/vector_stores/{vector_store_id}/file_batches/{batch_id}": { + "get": { + "responses": { + "200": { + "description": "A VectorStoreFileBatchObject representing the file batch.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFileBatchObject" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "summary": "Retrieve a vector store file batch.", + "description": "Retrieve a vector store file batch.", + "parameters": [ + { + "name": "batch_id", + "in": "path", + "description": "The ID of the file batch to retrieve.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "vector_store_id", + "in": "path", + "description": "The ID of the vector store containing the file batch.", + "required": true, + "schema": { + "type": "string" + } + } + ] + } + }, "/v1/vector_stores/{vector_store_id}/files/{file_id}/content": { "get": { "responses": { @@ -14710,6 +14968,82 @@ } ] }, + "VectorStoreFileBatchObject": { + "type": "object", + "properties": { + "id": { + "type": "string", + "description": "Unique identifier for the file batch" + }, + "object": { + "type": "string", + "default": "vector_store.file_batch", + "description": "Object type identifier, always \"vector_store.file_batch\"" + }, + "created_at": { + "type": "integer", + "description": "Timestamp when the file batch was created" + }, + "vector_store_id": { + "type": "string", + "description": "ID of the vector store containing the file batch" + }, + "status": { + "$ref": "#/components/schemas/VectorStoreFileStatus", + "description": "Current processing status of the file batch" + }, + "file_counts": { + "$ref": "#/components/schemas/VectorStoreFileCounts", + "description": "File processing status counts for the batch" + } + }, + "additionalProperties": false, + "required": [ + "id", + "object", + "created_at", + "vector_store_id", + "status", + "file_counts" + ], + "title": "VectorStoreFileBatchObject", + "description": "OpenAI Vector Store File Batch object." + }, + "VectorStoreFileCounts": { + "type": "object", + "properties": { + "completed": { + "type": "integer", + "description": "Number of files that have been successfully processed" + }, + "cancelled": { + "type": "integer", + "description": "Number of files that had their processing cancelled" + }, + "failed": { + "type": "integer", + "description": "Number of files that failed to process" + }, + "in_progress": { + "type": "integer", + "description": "Number of files currently being processed" + }, + "total": { + "type": "integer", + "description": "Total number of files in the vector store" + } + }, + "additionalProperties": false, + "required": [ + "completed", + "cancelled", + "failed", + "in_progress", + "total" + ], + "title": "VectorStoreFileCounts", + "description": "File processing status counts for a vector store." + }, "OpenAIJSONSchema": { "type": "object", "properties": { @@ -15541,41 +15875,6 @@ "additionalProperties": false, "title": "OpenaiCreateVectorStoreRequest" }, - "VectorStoreFileCounts": { - "type": "object", - "properties": { - "completed": { - "type": "integer", - "description": "Number of files that have been successfully processed" - }, - "cancelled": { - "type": "integer", - "description": "Number of files that had their processing cancelled" - }, - "failed": { - "type": "integer", - "description": "Number of files that failed to process" - }, - "in_progress": { - "type": "integer", - "description": "Number of files currently being processed" - }, - "total": { - "type": "integer", - "description": "Total number of files in the vector store" - } - }, - "additionalProperties": false, - "required": [ - "completed", - "cancelled", - "failed", - "in_progress", - "total" - ], - "title": "VectorStoreFileCounts", - "description": "File processing status counts for a vector store." - }, "VectorStoreObject": { "type": "object", "properties": { @@ -15684,6 +15983,53 @@ "title": "VectorStoreObject", "description": "OpenAI Vector Store object." }, + "OpenaiCreateVectorStoreFileBatchRequest": { + "type": "object", + "properties": { + "file_ids": { + "type": "array", + "items": { + "type": "string" + }, + "description": "A list of File IDs that the vector store should use." + }, + "attributes": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + }, + "description": "(Optional) Key-value attributes to store with the files." + }, + "chunking_strategy": { + "$ref": "#/components/schemas/VectorStoreChunkingStrategy", + "description": "(Optional) The chunking strategy used to chunk the file(s). Defaults to auto." + } + }, + "additionalProperties": false, + "required": [ + "file_ids" + ], + "title": "OpenaiCreateVectorStoreFileBatchRequest" + }, "OpenAIFileDeleteResponse": { "type": "object", "properties": { @@ -16036,6 +16382,44 @@ "title": "VectorStoreListFilesResponse", "description": "Response from listing files in a vector store." }, + "VectorStoreFilesListInBatchResponse": { + "type": "object", + "properties": { + "object": { + "type": "string", + "default": "list", + "description": "Object type identifier, always \"list\"" + }, + "data": { + "type": "array", + "items": { + "$ref": "#/components/schemas/VectorStoreFileObject" + }, + "description": "List of vector store file objects in the batch" + }, + "first_id": { + "type": "string", + "description": "(Optional) ID of the first file in the list for pagination" + }, + "last_id": { + "type": "string", + "description": "(Optional) ID of the last file in the list for pagination" + }, + "has_more": { + "type": "boolean", + "default": false, + "description": "Whether there are more files available beyond this page" + } + }, + "additionalProperties": false, + "required": [ + "object", + "data", + "has_more" + ], + "title": "VectorStoreFilesListInBatchResponse", + "description": "Response from listing files in a vector store file batch." + }, "VectorStoreListResponse": { "type": "object", "properties": { diff --git a/docs/static/llama-stack-spec.yaml b/docs/static/llama-stack-spec.yaml index 901bade95..33a7e66d8 100644 --- a/docs/static/llama-stack-spec.yaml +++ b/docs/static/llama-stack-spec.yaml @@ -3369,6 +3369,44 @@ paths: schema: $ref: '#/components/schemas/OpenaiAttachFileToVectorStoreRequest' required: true + /v1/vector_stores/{vector_store_id}/file_batches/{batch_id}/cancel: + post: + responses: + '200': + description: >- + A VectorStoreFileBatchObject representing the cancelled file batch. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFileBatchObject' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + summary: Cancels a vector store file batch. + description: Cancels a vector store file batch. + parameters: + - name: batch_id + in: path + description: The ID of the file batch to cancel. + required: true + schema: + type: string + - name: vector_store_id + in: path + description: >- + The ID of the vector store containing the file batch. + required: true + schema: + type: string /v1/completions: post: responses: @@ -3490,6 +3528,44 @@ paths: schema: $ref: '#/components/schemas/OpenaiCreateVectorStoreRequest' required: true + /v1/vector_stores/{vector_store_id}/file_batches: + post: + responses: + '200': + description: >- + A VectorStoreFileBatchObject representing the created file batch. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFileBatchObject' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + summary: Create a vector store file batch. + description: Create a vector store file batch. + 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: + schema: + $ref: '#/components/schemas/OpenaiCreateVectorStoreFileBatchRequest' + required: true /v1/files/{file_id}: get: responses: @@ -3916,6 +3992,87 @@ paths: - file - purpose required: true + /v1/vector_stores/{vector_store_id}/file_batches/{batch_id}/files: + get: + responses: + '200': + description: >- + A VectorStoreFilesListInBatchResponse containing the list of files in + the batch. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFilesListInBatchResponse' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + summary: >- + Returns a list of vector store files in a batch. + description: >- + Returns a list of vector store files in a batch. + parameters: + - name: batch_id + in: path + description: >- + The ID of the file batch to list files from. + required: true + schema: + type: string + - name: vector_store_id + in: path + description: >- + The ID of the vector store containing the file batch. + required: true + schema: + type: string + - name: after + in: query + description: >- + A cursor for use in pagination. `after` is an object ID that defines your + place in the list. + required: false + schema: + type: string + - name: before + in: query + description: >- + A cursor for use in pagination. `before` is an object ID that defines + your place in the list. + required: false + schema: + type: string + - name: filter + in: query + description: >- + Filter by file status. One of in_progress, completed, failed, cancelled. + required: false + schema: + type: string + - name: limit + in: query + description: >- + A limit on the number of objects to be returned. Limit can range between + 1 and 100, and the default is 20. + required: false + schema: + type: integer + - name: order + in: query + description: >- + Sort order by the `created_at` timestamp of the objects. `asc` for ascending + order and `desc` for descending order. + required: false + schema: + type: string /v1/files/{file_id}/content: get: responses: @@ -3950,6 +4107,44 @@ paths: required: true schema: type: string + /v1/vector_stores/{vector_store_id}/file_batches/{batch_id}: + get: + responses: + '200': + description: >- + A VectorStoreFileBatchObject representing the file batch. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFileBatchObject' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + summary: Retrieve a vector store file batch. + description: Retrieve a vector store file batch. + parameters: + - name: batch_id + in: path + description: The ID of the file batch to retrieve. + required: true + schema: + type: string + - name: vector_store_id + in: path + description: >- + The ID of the vector store containing the file batch. + required: true + schema: + type: string /v1/vector_stores/{vector_store_id}/files/{file_id}/content: get: responses: @@ -10870,6 +11065,75 @@ components: const: cancelled - type: string const: failed + VectorStoreFileBatchObject: + type: object + properties: + id: + type: string + description: Unique identifier for the file batch + object: + type: string + default: vector_store.file_batch + description: >- + Object type identifier, always "vector_store.file_batch" + created_at: + type: integer + description: >- + Timestamp when the file batch was created + vector_store_id: + type: string + description: >- + ID of the vector store containing the file batch + status: + $ref: '#/components/schemas/VectorStoreFileStatus' + description: >- + Current processing status of the file batch + file_counts: + $ref: '#/components/schemas/VectorStoreFileCounts' + description: >- + File processing status counts for the batch + additionalProperties: false + required: + - id + - object + - created_at + - vector_store_id + - status + - file_counts + title: VectorStoreFileBatchObject + description: OpenAI Vector Store File Batch object. + VectorStoreFileCounts: + type: object + properties: + completed: + type: integer + description: >- + Number of files that have been successfully processed + cancelled: + type: integer + description: >- + Number of files that had their processing cancelled + failed: + type: integer + description: Number of files that failed to process + in_progress: + type: integer + description: >- + Number of files currently being processed + total: + type: integer + description: >- + Total number of files in the vector store + additionalProperties: false + required: + - completed + - cancelled + - failed + - in_progress + - total + title: VectorStoreFileCounts + description: >- + File processing status counts for a vector store. OpenAIJSONSchema: type: object properties: @@ -11432,38 +11696,6 @@ components: The ID of the provider to use for this vector store. additionalProperties: false title: OpenaiCreateVectorStoreRequest - VectorStoreFileCounts: - type: object - properties: - completed: - type: integer - description: >- - Number of files that have been successfully processed - cancelled: - type: integer - description: >- - Number of files that had their processing cancelled - failed: - type: integer - description: Number of files that failed to process - in_progress: - type: integer - description: >- - Number of files currently being processed - total: - type: integer - description: >- - Total number of files in the vector store - additionalProperties: false - required: - - completed - - cancelled - - failed - - in_progress - - total - title: VectorStoreFileCounts - description: >- - File processing status counts for a vector store. VectorStoreObject: type: object properties: @@ -11538,6 +11770,36 @@ components: - metadata title: VectorStoreObject description: OpenAI Vector Store object. + OpenaiCreateVectorStoreFileBatchRequest: + type: object + properties: + file_ids: + type: array + items: + type: string + description: >- + A list of File IDs that the vector store should use. + attributes: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + description: >- + (Optional) Key-value attributes to store with the files. + chunking_strategy: + $ref: '#/components/schemas/VectorStoreChunkingStrategy' + description: >- + (Optional) The chunking strategy used to chunk the file(s). Defaults to + auto. + additionalProperties: false + required: + - file_ids + title: OpenaiCreateVectorStoreFileBatchRequest OpenAIFileDeleteResponse: type: object properties: @@ -11841,6 +12103,40 @@ components: title: VectorStoreListFilesResponse description: >- Response from listing files in a vector store. + VectorStoreFilesListInBatchResponse: + type: object + properties: + object: + type: string + default: list + description: Object type identifier, always "list" + data: + type: array + items: + $ref: '#/components/schemas/VectorStoreFileObject' + description: >- + List of vector store file objects in the batch + first_id: + type: string + description: >- + (Optional) ID of the first file in the list for pagination + last_id: + type: string + description: >- + (Optional) ID of the last file in the list for pagination + has_more: + type: boolean + default: false + description: >- + Whether there are more files available beyond this page + additionalProperties: false + required: + - object + - data + - has_more + title: VectorStoreFilesListInBatchResponse + description: >- + Response from listing files in a vector store file batch. VectorStoreListResponse: type: object properties: diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index cea2a6917..e07175c49 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/llama_stack/apis/vector_io/vector_io.py @@ -318,7 +318,8 @@ class VectorStoreChunkingStrategyStatic(BaseModel): VectorStoreChunkingStrategy = Annotated[ - VectorStoreChunkingStrategyAuto | VectorStoreChunkingStrategyStatic, Field(discriminator="type") + VectorStoreChunkingStrategyAuto | VectorStoreChunkingStrategyStatic, + Field(discriminator="type"), ] register_schema(VectorStoreChunkingStrategy, name="VectorStoreChunkingStrategy") @@ -427,6 +428,44 @@ class VectorStoreFileDeleteResponse(BaseModel): deleted: bool = True +@json_schema_type +class VectorStoreFileBatchObject(BaseModel): + """OpenAI Vector Store File Batch object. + + :param id: Unique identifier for the file batch + :param object: Object type identifier, always "vector_store.file_batch" + :param created_at: Timestamp when the file batch was created + :param vector_store_id: ID of the vector store containing the file batch + :param status: Current processing status of the file batch + :param file_counts: File processing status counts for the batch + """ + + id: str + object: str = "vector_store.file_batch" + created_at: int + vector_store_id: str + status: VectorStoreFileStatus + file_counts: VectorStoreFileCounts + + +@json_schema_type +class VectorStoreFilesListInBatchResponse(BaseModel): + """Response from listing files in a vector store file batch. + + :param object: Object type identifier, always "list" + :param data: List of vector store file objects in the batch + :param first_id: (Optional) ID of the first file in the list for pagination + :param last_id: (Optional) ID of the last file in the list for pagination + :param has_more: Whether there are more files available beyond this page + """ + + object: str = "list" + data: list[VectorStoreFileObject] + first_id: str | None = None + last_id: str | None = None + has_more: bool = False + + class VectorDBStore(Protocol): def get_vector_db(self, vector_db_id: str) -> VectorDB | None: ... @@ -529,7 +568,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}", method="POST", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}", + method="POST", + level=LLAMA_STACK_API_V1, + ) async def openai_update_vector_store( self, vector_store_id: str, @@ -547,7 +590,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}", method="DELETE", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}", + method="DELETE", + level=LLAMA_STACK_API_V1, + ) async def openai_delete_vector_store( self, vector_store_id: str, @@ -559,7 +606,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}/search", method="POST", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}/search", + method="POST", + level=LLAMA_STACK_API_V1, + ) async def openai_search_vector_store( self, vector_store_id: str, @@ -568,7 +619,9 @@ class VectorIO(Protocol): max_num_results: int | None = 10, ranking_options: SearchRankingOptions | None = None, rewrite_query: bool | None = False, - search_mode: str | None = "vector", # Using str instead of Literal due to OpenAPI schema generator limitations + search_mode: ( + str | None + ) = "vector", # Using str instead of Literal due to OpenAPI schema generator limitations ) -> VectorStoreSearchResponsePage: """Search for chunks in a vector store. @@ -585,7 +638,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}/files", method="POST", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}/files", + method="POST", + level=LLAMA_STACK_API_V1, + ) async def openai_attach_file_to_vector_store( self, vector_store_id: str, @@ -603,7 +660,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}/files", method="GET", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}/files", + method="GET", + level=LLAMA_STACK_API_V1, + ) async def openai_list_files_in_vector_store( self, vector_store_id: str, @@ -625,7 +686,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}/files/{file_id}", method="GET", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}/files/{file_id}", + method="GET", + level=LLAMA_STACK_API_V1, + ) async def openai_retrieve_vector_store_file( self, vector_store_id: str, @@ -657,7 +722,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}/files/{file_id}", method="POST", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}/files/{file_id}", + method="POST", + level=LLAMA_STACK_API_V1, + ) async def openai_update_vector_store_file( self, vector_store_id: str, @@ -673,7 +742,11 @@ class VectorIO(Protocol): """ ... - @webmethod(route="/vector_stores/{vector_store_id}/files/{file_id}", method="DELETE", level=LLAMA_STACK_API_V1) + @webmethod( + route="/vector_stores/{vector_store_id}/files/{file_id}", + method="DELETE", + level=LLAMA_STACK_API_V1, + ) async def openai_delete_vector_store_file( self, vector_store_id: str, @@ -686,3 +759,89 @@ class VectorIO(Protocol): :returns: A VectorStoreFileDeleteResponse indicating the deletion status. """ ... + + @webmethod( + route="/vector_stores/{vector_store_id}/file_batches", + method="POST", + level=LLAMA_STACK_API_V1, + ) + async def openai_create_vector_store_file_batch( + self, + vector_store_id: str, + file_ids: list[str], + attributes: dict[str, Any] | None = None, + chunking_strategy: VectorStoreChunkingStrategy | None = None, + ) -> VectorStoreFileBatchObject: + """Create a vector store file batch. + + :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. + :returns: A VectorStoreFileBatchObject representing the created file batch. + """ + ... + + @webmethod( + route="/vector_stores/{vector_store_id}/file_batches/{batch_id}", + method="GET", + level=LLAMA_STACK_API_V1, + ) + async def openai_retrieve_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + ) -> VectorStoreFileBatchObject: + """Retrieve a vector store file batch. + + :param batch_id: The ID of the file batch to retrieve. + :param vector_store_id: The ID of the vector store containing the file batch. + :returns: A VectorStoreFileBatchObject representing the file batch. + """ + ... + + @webmethod( + route="/vector_stores/{vector_store_id}/file_batches/{batch_id}/files", + method="GET", + level=LLAMA_STACK_API_V1, + ) + async def openai_list_files_in_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + after: str | None = None, + before: str | None = None, + filter: str | None = None, + limit: int | None = 20, + order: str | None = "desc", + ) -> VectorStoreFilesListInBatchResponse: + """Returns a list of vector store files in a batch. + + :param batch_id: The ID of the file batch to list files from. + :param vector_store_id: The ID of the vector store containing the file batch. + :param after: A cursor for use in pagination. `after` is an object ID that defines your place in the list. + :param before: A cursor for use in pagination. `before` is an object ID that defines your place in the list. + :param filter: Filter by file status. One of in_progress, completed, failed, cancelled. + :param limit: A limit on the number of objects to be returned. Limit can range between 1 and 100, and the default is 20. + :param order: Sort order by the `created_at` timestamp of the objects. `asc` for ascending order and `desc` for descending order. + :returns: A VectorStoreFilesListInBatchResponse containing the list of files in the batch. + """ + ... + + @webmethod( + route="/vector_stores/{vector_store_id}/file_batches/{batch_id}/cancel", + method="POST", + level=LLAMA_STACK_API_V1, + ) + async def openai_cancel_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + ) -> VectorStoreFileBatchObject: + """Cancels a vector store file batch. + + :param batch_id: The ID of the file batch to cancel. + :param vector_store_id: The ID of the vector store containing the file batch. + :returns: A VectorStoreFileBatchObject representing the cancelled file batch. + """ + ... diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index 786b0e391..0e3f9d8d9 100644 --- a/llama_stack/core/routers/vector_io.py +++ b/llama_stack/core/routers/vector_io.py @@ -8,9 +8,7 @@ import asyncio import uuid from typing import Any -from llama_stack.apis.common.content_types import ( - InterleavedContent, -) +from llama_stack.apis.common.content_types import InterleavedContent from llama_stack.apis.models import ModelType from llama_stack.apis.vector_io import ( Chunk, @@ -19,9 +17,11 @@ from llama_stack.apis.vector_io import ( VectorIO, VectorStoreChunkingStrategy, VectorStoreDeleteResponse, + VectorStoreFileBatchObject, VectorStoreFileContentsResponse, VectorStoreFileDeleteResponse, VectorStoreFileObject, + VectorStoreFilesListInBatchResponse, VectorStoreFileStatus, VectorStoreListResponse, VectorStoreObject, @@ -193,7 +193,10 @@ class VectorIORouter(VectorIO): all_stores = all_stores[after_index + 1 :] if before: - before_index = next((i for i, store in enumerate(all_stores) if store.id == before), len(all_stores)) + before_index = next( + (i for i, store in enumerate(all_stores) if store.id == before), + len(all_stores), + ) all_stores = all_stores[:before_index] # Apply limit @@ -363,3 +366,61 @@ class VectorIORouter(VectorIO): status=HealthStatus.ERROR, message=f"Health check failed: {str(e)}" ) return health_statuses + + async def openai_create_vector_store_file_batch( + self, + vector_store_id: str, + file_ids: list[str], + attributes: dict[str, Any] | None = None, + chunking_strategy: VectorStoreChunkingStrategy | None = None, + ) -> VectorStoreFileBatchObject: + logger.debug(f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(file_ids)} files") + return await self.routing_table.openai_create_vector_store_file_batch( + vector_store_id=vector_store_id, + file_ids=file_ids, + attributes=attributes, + chunking_strategy=chunking_strategy, + ) + + async def openai_retrieve_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + ) -> VectorStoreFileBatchObject: + logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_batch: {batch_id}, {vector_store_id}") + return await self.routing_table.openai_retrieve_vector_store_file_batch( + batch_id=batch_id, + vector_store_id=vector_store_id, + ) + + async def openai_list_files_in_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + after: str | None = None, + before: str | None = None, + filter: str | None = None, + limit: int | None = 20, + order: str | None = "desc", + ) -> VectorStoreFilesListInBatchResponse: + logger.debug(f"VectorIORouter.openai_list_files_in_vector_store_file_batch: {batch_id}, {vector_store_id}") + return await self.routing_table.openai_list_files_in_vector_store_file_batch( + batch_id=batch_id, + vector_store_id=vector_store_id, + after=after, + before=before, + filter=filter, + limit=limit, + order=order, + ) + + async def openai_cancel_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + ) -> VectorStoreFileBatchObject: + logger.debug(f"VectorIORouter.openai_cancel_vector_store_file_batch: {batch_id}, {vector_store_id}") + return await self.routing_table.openai_cancel_vector_store_file_batch( + batch_id=batch_id, + vector_store_id=vector_store_id, + ) diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index 3acdcf293..36432767f 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -24,11 +24,13 @@ from llama_stack.apis.vector_io import ( VectorStoreChunkingStrategyStatic, VectorStoreContent, VectorStoreDeleteResponse, + VectorStoreFileBatchObject, VectorStoreFileContentsResponse, VectorStoreFileCounts, VectorStoreFileDeleteResponse, VectorStoreFileLastError, VectorStoreFileObject, + VectorStoreFilesListInBatchResponse, VectorStoreFileStatus, VectorStoreListFilesResponse, VectorStoreListResponse, @@ -107,7 +109,11 @@ class OpenAIVectorStoreMixin(ABC): self.openai_vector_stores.pop(store_id, None) async def _save_openai_vector_store_file( - self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]] + self, + store_id: str, + file_id: str, + file_info: dict[str, Any], + file_contents: list[dict[str, Any]], ) -> None: """Save vector store file metadata to persistent storage.""" assert self.kvstore @@ -301,7 +307,10 @@ class OpenAIVectorStoreMixin(ABC): all_stores = all_stores[after_index + 1 :] if before: - before_index = next((i for i, store in enumerate(all_stores) if store["id"] == before), len(all_stores)) + before_index = next( + (i for i, store in enumerate(all_stores) if store["id"] == before), + len(all_stores), + ) all_stores = all_stores[:before_index] # Apply limit @@ -397,7 +406,9 @@ class OpenAIVectorStoreMixin(ABC): max_num_results: int | None = 10, ranking_options: SearchRankingOptions | None = None, rewrite_query: bool | None = False, - search_mode: str | None = "vector", # Using str instead of Literal due to OpenAPI schema generator limitations + search_mode: ( + str | None + ) = "vector", # Using str instead of Literal due to OpenAPI schema generator limitations ) -> VectorStoreSearchResponsePage: """Search for chunks in a vector store.""" max_num_results = max_num_results or 10 @@ -685,7 +696,10 @@ class OpenAIVectorStoreMixin(ABC): file_objects = file_objects[after_index + 1 :] if before: - before_index = next((i for i, file in enumerate(file_objects) if file.id == before), len(file_objects)) + before_index = next( + (i for i, file in enumerate(file_objects) if file.id == before), + len(file_objects), + ) file_objects = file_objects[:before_index] # Apply limit @@ -805,3 +819,42 @@ class OpenAIVectorStoreMixin(ABC): id=file_id, deleted=True, ) + + async def openai_create_vector_store_file_batch( + self, + vector_store_id: str, + file_ids: list[str], + attributes: dict[str, Any] | None = None, + chunking_strategy: VectorStoreChunkingStrategy | None = None, + ) -> VectorStoreFileBatchObject: + """Create a vector store file batch.""" + raise NotImplementedError("openai_create_vector_store_file_batch is not implemented yet") + + async def openai_list_files_in_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + after: str | None = None, + before: str | None = None, + filter: str | None = None, + limit: int | None = 20, + order: str | None = "desc", + ) -> VectorStoreFilesListInBatchResponse: + """Returns a list of vector store files in a batch.""" + raise NotImplementedError("openai_list_files_in_vector_store_file_batch is not implemented yet") + + async def openai_retrieve_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + ) -> VectorStoreFileBatchObject: + """Retrieve a vector store file batch.""" + raise NotImplementedError("openai_retrieve_vector_store_file_batch is not implemented yet") + + async def openai_cancel_vector_store_file_batch( + self, + batch_id: str, + vector_store_id: str, + ) -> VectorStoreFileBatchObject: + """Cancel a vector store file batch.""" + raise NotImplementedError("openai_cancel_vector_store_file_batch is not implemented yet")