diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index 5407f9808..bad747c5c 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -3241,6 +3241,47 @@ } }, "/v1/openai/v1/vector_stores/{vector_store_id}/files": { + "get": { + "responses": { + "200": { + "description": "A VectorStoreListFilesResponse containing the list of files.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreListFilesResponse" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "description": "List files in a vector store.", + "parameters": [ + { + "name": "vector_store_id", + "in": "path", + "description": "The ID of the vector store to list files from.", + "required": true, + "schema": { + "type": "string" + } + } + ] + }, "post": { "responses": { "200": { @@ -3666,6 +3707,168 @@ ] } }, + "/v1/openai/v1/vector_stores/{vector_store_id}/files/{file_id}": { + "get": { + "responses": { + "200": { + "description": "A VectorStoreFileObject representing the file.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFileObject" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "description": "Retrieves a vector store file.", + "parameters": [ + { + "name": "vector_store_id", + "in": "path", + "description": "The ID of the vector store containing the file to retrieve.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "file_id", + "in": "path", + "description": "The ID of the file to retrieve.", + "required": true, + "schema": { + "type": "string" + } + } + ] + }, + "post": { + "responses": { + "200": { + "description": "A VectorStoreFileObject representing the updated file.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFileObject" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "description": "Updates a vector store file.", + "parameters": [ + { + "name": "vector_store_id", + "in": "path", + "description": "The ID of the vector store containing the file to update.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "file_id", + "in": "path", + "description": "The ID of the file to update.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/OpenaiUpdateVectorStoreFileRequest" + } + } + }, + "required": true + } + }, + "delete": { + "responses": { + "200": { + "description": "A VectorStoreFileDeleteResponse indicating the deletion status.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFileDeleteResponse" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "description": "Delete a vector store file.", + "parameters": [ + { + "name": "vector_store_id", + "in": "path", + "description": "The ID of the vector store containing the file to delete.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "file_id", + "in": "path", + "description": "The ID of the file to delete.", + "required": true, + "schema": { + "type": "string" + } + } + ] + } + }, "/v1/openai/v1/embeddings": { "post": { "responses": { @@ -12969,6 +13172,35 @@ ], "title": "OpenaiCreateVectorStoreRequest" }, + "VectorStoreFileCounts": { + "type": "object", + "properties": { + "completed": { + "type": "integer" + }, + "cancelled": { + "type": "integer" + }, + "failed": { + "type": "integer" + }, + "in_progress": { + "type": "integer" + }, + "total": { + "type": "integer" + } + }, + "additionalProperties": false, + "required": [ + "completed", + "cancelled", + "failed", + "in_progress", + "total" + ], + "title": "VectorStoreFileCounts" + }, "VectorStoreObject": { "type": "object", "properties": { @@ -12990,10 +13222,7 @@ "default": 0 }, "file_counts": { - "type": "object", - "additionalProperties": { - "type": "integer" - } + "$ref": "#/components/schemas/VectorStoreFileCounts" }, "status": { "type": "string", @@ -13120,6 +13349,30 @@ "title": "VectorStoreDeleteResponse", "description": "Response from deleting a vector store." }, + "VectorStoreFileDeleteResponse": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "object": { + "type": "string", + "default": "vector_store.file.deleted" + }, + "deleted": { + "type": "boolean", + "default": true + } + }, + "additionalProperties": false, + "required": [ + "id", + "object", + "deleted" + ], + "title": "VectorStoreFileDeleteResponse", + "description": "Response from deleting a vector store file." + }, "OpenaiEmbeddingsRequest": { "type": "object", "properties": { @@ -13348,6 +13601,28 @@ "title": "OpenAIFileObject", "description": "OpenAI File object as defined in the OpenAI Files API." }, + "VectorStoreListFilesResponse": { + "type": "object", + "properties": { + "object": { + "type": "string", + "default": "list" + }, + "data": { + "type": "array", + "items": { + "$ref": "#/components/schemas/VectorStoreFileObject" + } + } + }, + "additionalProperties": false, + "required": [ + "object", + "data" + ], + "title": "VectorStoreListFilesResponse", + "description": "Response from listing vector stores." + }, "OpenAIModel": { "type": "object", "properties": { @@ -13661,6 +13936,42 @@ "additionalProperties": false, "title": "OpenaiUpdateVectorStoreRequest" }, + "OpenaiUpdateVectorStoreFileRequest": { + "type": "object", + "properties": { + "attributes": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + }, + "description": "The updated key-value attributes to store with the file." + } + }, + "additionalProperties": false, + "required": [ + "attributes" + ], + "title": "OpenaiUpdateVectorStoreFileRequest" + }, "DPOAlignmentConfig": { "type": "object", "properties": { diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index a354e4bd0..c02decfe2 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -2264,6 +2264,36 @@ paths: $ref: '#/components/schemas/LogEventRequest' required: true /v1/openai/v1/vector_stores/{vector_store_id}/files: + get: + responses: + '200': + description: >- + A VectorStoreListFilesResponse containing the list of files. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreListFilesResponse' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + description: List files in a vector store. + parameters: + - name: vector_store_id + in: path + description: >- + The ID of the vector store to list files from. + required: true + schema: + type: string post: responses: '200': @@ -2572,6 +2602,121 @@ paths: required: true schema: type: string + /v1/openai/v1/vector_stores/{vector_store_id}/files/{file_id}: + get: + responses: + '200': + description: >- + A VectorStoreFileObject representing the file. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFileObject' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + description: Retrieves a vector store file. + parameters: + - name: vector_store_id + in: path + description: >- + The ID of the vector store containing the file to retrieve. + required: true + schema: + type: string + - name: file_id + in: path + description: The ID of the file to retrieve. + required: true + schema: + type: string + post: + responses: + '200': + description: >- + A VectorStoreFileObject representing the updated file. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFileObject' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + description: Updates a vector store file. + parameters: + - name: vector_store_id + in: path + description: >- + The ID of the vector store containing the file to update. + required: true + schema: + type: string + - name: file_id + in: path + description: The ID of the file to update. + required: true + schema: + type: string + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/OpenaiUpdateVectorStoreFileRequest' + required: true + delete: + responses: + '200': + description: >- + A VectorStoreFileDeleteResponse indicating the deletion status. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFileDeleteResponse' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + description: Delete a vector store file. + parameters: + - name: vector_store_id + in: path + description: >- + The ID of the vector store containing the file to delete. + required: true + schema: + type: string + - name: file_id + in: path + description: The ID of the file to delete. + required: true + schema: + type: string /v1/openai/v1/embeddings: post: responses: @@ -9031,6 +9176,27 @@ components: required: - name title: OpenaiCreateVectorStoreRequest + VectorStoreFileCounts: + type: object + properties: + completed: + type: integer + cancelled: + type: integer + failed: + type: integer + in_progress: + type: integer + total: + type: integer + additionalProperties: false + required: + - completed + - cancelled + - failed + - in_progress + - total + title: VectorStoreFileCounts VectorStoreObject: type: object properties: @@ -9047,9 +9213,7 @@ components: type: integer default: 0 file_counts: - type: object - additionalProperties: - type: integer + $ref: '#/components/schemas/VectorStoreFileCounts' status: type: string default: completed @@ -9129,6 +9293,25 @@ components: - deleted title: VectorStoreDeleteResponse description: Response from deleting a vector store. + VectorStoreFileDeleteResponse: + type: object + properties: + id: + type: string + object: + type: string + default: vector_store.file.deleted + deleted: + type: boolean + default: true + additionalProperties: false + required: + - id + - object + - deleted + title: VectorStoreFileDeleteResponse + description: >- + Response from deleting a vector store file. OpenaiEmbeddingsRequest: type: object properties: @@ -9320,6 +9503,22 @@ components: title: OpenAIFileObject description: >- OpenAI File object as defined in the OpenAI Files API. + VectorStoreListFilesResponse: + type: object + properties: + object: + type: string + default: list + data: + type: array + items: + $ref: '#/components/schemas/VectorStoreFileObject' + additionalProperties: false + required: + - object + - data + title: VectorStoreListFilesResponse + description: Response from listing vector stores. OpenAIModel: type: object properties: @@ -9524,6 +9723,25 @@ components: Set of 16 key-value pairs that can be attached to an object. additionalProperties: false title: OpenaiUpdateVectorStoreRequest + OpenaiUpdateVectorStoreFileRequest: + type: object + properties: + attributes: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + description: >- + The updated key-value attributes to store with the file. + additionalProperties: false + required: + - attributes + title: OpenaiUpdateVectorStoreFileRequest DPOAlignmentConfig: type: object properties: diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index 20cc594cc..8e569bfeb 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/llama_stack/apis/vector_io/vector_io.py @@ -38,6 +38,15 @@ class QueryChunksResponse(BaseModel): scores: list[float] +@json_schema_type +class VectorStoreFileCounts(BaseModel): + completed: int + cancelled: int + failed: int + in_progress: int + total: int + + @json_schema_type class VectorStoreObject(BaseModel): """OpenAI Vector Store object.""" @@ -47,7 +56,7 @@ class VectorStoreObject(BaseModel): created_at: int name: str | None = None usage_bytes: int = 0 - file_counts: dict[str, int] = Field(default_factory=dict) + file_counts: VectorStoreFileCounts status: str = "completed" expires_after: dict[str, Any] | None = None expires_at: int | None = None @@ -183,6 +192,23 @@ class VectorStoreFileObject(BaseModel): vector_store_id: str +@json_schema_type +class VectorStoreListFilesResponse(BaseModel): + """Response from listing vector stores.""" + + object: str = "list" + data: list[VectorStoreFileObject] + + +@json_schema_type +class VectorStoreFileDeleteResponse(BaseModel): + """Response from deleting a vector store file.""" + + id: str + object: str = "vector_store.file.deleted" + deleted: bool = True + + class VectorDBStore(Protocol): def get_vector_db(self, vector_db_id: str) -> VectorDB | None: ... @@ -358,3 +384,59 @@ class VectorIO(Protocol): :returns: A VectorStoreFileObject representing the attached file. """ ... + + @webmethod(route="/openai/v1/vector_stores/{vector_store_id}/files", method="GET") + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + ) -> VectorStoreListFilesResponse: + """List files in a vector store. + + :param vector_store_id: The ID of the vector store to list files from. + :returns: A VectorStoreListFilesResponse containing the list of files. + """ + ... + + @webmethod(route="/openai/v1/vector_stores/{vector_store_id}/files/{file_id}", method="GET") + async def openai_retrieve_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + """Retrieves a vector store file. + + :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. + :returns: A VectorStoreFileObject representing the file. + """ + ... + + @webmethod(route="/openai/v1/vector_stores/{vector_store_id}/files/{file_id}", method="POST") + async def openai_update_vector_store_file( + self, + vector_store_id: str, + file_id: str, + attributes: dict[str, Any], + ) -> VectorStoreFileObject: + """Updates a vector store file. + + :param vector_store_id: The ID of the vector store containing the file to update. + :param file_id: The ID of the file to update. + :param attributes: The updated key-value attributes to store with the file. + :returns: A VectorStoreFileObject representing the updated file. + """ + ... + + @webmethod(route="/openai/v1/vector_stores/{vector_store_id}/files/{file_id}", method="DELETE") + async def openai_delete_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileDeleteResponse: + """Delete a vector store file. + + :param vector_store_id: The ID of the vector store containing the file to delete. + :param file_id: The ID of the file to delete. + :returns: A VectorStoreFileDeleteResponse indicating the deletion status. + """ + ... diff --git a/llama_stack/distribution/routers/vector_io.py b/llama_stack/distribution/routers/vector_io.py index 3001b8666..c09b1df2e 100644 --- a/llama_stack/distribution/routers/vector_io.py +++ b/llama_stack/distribution/routers/vector_io.py @@ -21,7 +21,11 @@ from llama_stack.apis.vector_io import ( VectorStoreObject, VectorStoreSearchResponsePage, ) -from llama_stack.apis.vector_io.vector_io import VectorStoreChunkingStrategy, VectorStoreFileObject +from llama_stack.apis.vector_io.vector_io import ( + VectorStoreChunkingStrategy, + VectorStoreFileDeleteResponse, + VectorStoreFileObject, +) from llama_stack.log import get_logger from llama_stack.providers.datatypes import HealthResponse, HealthStatus, RoutingTable @@ -279,6 +283,58 @@ class VectorIORouter(VectorIO): chunking_strategy=chunking_strategy, ) + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + ) -> list[VectorStoreFileObject]: + logger.debug(f"VectorIORouter.openai_list_files_in_vector_store: {vector_store_id}") + # Route based on vector store ID + provider = self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_list_files_in_vector_store( + vector_store_id=vector_store_id, + ) + + async def openai_retrieve_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file: {vector_store_id}, {file_id}") + # Route based on vector store ID + provider = self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_retrieve_vector_store_file( + vector_store_id=vector_store_id, + file_id=file_id, + ) + + async def openai_update_vector_store_file( + self, + vector_store_id: str, + file_id: str, + attributes: dict[str, Any], + ) -> VectorStoreFileObject: + logger.debug(f"VectorIORouter.openai_update_vector_store_file: {vector_store_id}, {file_id}") + # Route based on vector store ID + provider = self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_update_vector_store_file( + vector_store_id=vector_store_id, + file_id=file_id, + attributes=attributes, + ) + + async def openai_delete_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileDeleteResponse: + logger.debug(f"VectorIORouter.openai_delete_vector_store_file: {vector_store_id}, {file_id}") + # Route based on vector store ID + provider = self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_delete_vector_store_file( + vector_store_id=vector_store_id, + file_id=file_id, + ) + async def health(self) -> dict[str, HealthResponse]: health_statuses = {} timeout = 1 # increasing the timeout to 1 second for health checks diff --git a/llama_stack/providers/inline/vector_io/faiss/faiss.py b/llama_stack/providers/inline/vector_io/faiss/faiss.py index 0864ba3a7..83c74bce5 100644 --- a/llama_stack/providers/inline/vector_io/faiss/faiss.py +++ b/llama_stack/providers/inline/vector_io/faiss/faiss.py @@ -45,6 +45,7 @@ VERSION = "v3" VECTOR_DBS_PREFIX = f"vector_dbs:{VERSION}::" FAISS_INDEX_PREFIX = f"faiss_index:{VERSION}::" OPENAI_VECTOR_STORES_PREFIX = f"openai_vector_stores:{VERSION}::" +OPENAI_VECTOR_STORES_FILES_PREFIX = f"openai_vector_stores_files:{VERSION}::" class FaissIndex(EmbeddingIndex): @@ -283,3 +284,28 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPr assert self.kvstore is not None key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}" await self.kvstore.delete(key) + + async def _save_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: + """Save vector store file metadata to kvstore.""" + assert self.kvstore is not None + key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}" + await self.kvstore.set(key=key, value=json.dumps(file_info)) + + async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]: + """Load vector store file metadata from kvstore.""" + assert self.kvstore is not None + key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}" + stored_data = await self.kvstore.get(key) + return json.loads(stored_data) if stored_data else {} + + async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: + """Update vector store file metadata in kvstore.""" + assert self.kvstore is not None + key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}" + await self.kvstore.set(key=key, value=json.dumps(file_info)) + + async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None: + """Delete vector store file metadata from kvstore.""" + assert self.kvstore is not None + key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}" + await self.kvstore.delete(key) diff --git a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py b/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py index c6712882a..50cc262e4 100644 --- a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py +++ b/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py @@ -461,6 +461,14 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc metadata TEXT ); """) + # Create a table to persist OpenAI vector store files. + cur.execute(""" + CREATE TABLE IF NOT EXISTS openai_vector_store_files ( + store_id TEXT, + file_id TEXT, + metadata TEXT + ); + """) connection.commit() # Load any existing vector DB registrations. cur.execute("SELECT metadata FROM vector_dbs") @@ -615,6 +623,90 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc await asyncio.to_thread(_delete) + async def _save_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: + """Save vector store file metadata to SQLite database.""" + + def _store(): + connection = _create_sqlite_connection(self.config.db_path) + cur = connection.cursor() + try: + cur.execute( + "INSERT OR REPLACE INTO openai_vector_store_files (store_id, file_id, metadata) VALUES (?, ?, ?)", + (store_id, file_id, json.dumps(file_info)), + ) + connection.commit() + except Exception as e: + logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}") + raise + finally: + cur.close() + connection.close() + + try: + await asyncio.to_thread(_store) + except Exception as e: + logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}") + raise + + async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]: + """Load vector store file metadata from SQLite database.""" + + def _load(): + connection = _create_sqlite_connection(self.config.db_path) + cur = connection.cursor() + try: + cur.execute( + "SELECT metadata FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?", + (store_id, file_id), + ) + row = cur.fetchone() + print(f"!!! row is {row}") + if row is None: + return None + (metadata,) = row + return metadata + finally: + cur.close() + connection.close() + + stored_data = await asyncio.to_thread(_load) + return json.loads(stored_data) if stored_data else {} + + async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: + """Update vector store file metadata in SQLite database.""" + + def _update(): + connection = _create_sqlite_connection(self.config.db_path) + cur = connection.cursor() + try: + cur.execute( + "UPDATE openai_vector_store_files SET metadata = ? WHERE store_id = ? AND file_id = ?", + (json.dumps(file_info), store_id, file_id), + ) + connection.commit() + finally: + cur.close() + connection.close() + + await asyncio.to_thread(_update) + + async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None: + """Delete vector store file metadata from SQLite database.""" + + def _delete(): + connection = _create_sqlite_connection(self.config.db_path) + cur = connection.cursor() + try: + cur.execute( + "DELETE FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?", (store_id, file_id) + ) + connection.commit() + finally: + cur.close() + connection.close() + + await asyncio.to_thread(_delete) + async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: if vector_db_id not in self.cache: raise ValueError(f"Vector DB {vector_db_id} not found. Found: {list(self.cache.keys())}") diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index 12c1b5022..cb9e49409 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -24,7 +24,11 @@ from llama_stack.apis.vector_io import ( VectorStoreObject, VectorStoreSearchResponsePage, ) -from llama_stack.apis.vector_io.vector_io import VectorStoreChunkingStrategy, VectorStoreFileObject +from llama_stack.apis.vector_io.vector_io import ( + VectorStoreChunkingStrategy, + VectorStoreFileObject, + VectorStoreListFilesResponse, +) from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig from llama_stack.providers.utils.memory.vector_store import ( @@ -263,3 +267,31 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate): chunking_strategy: VectorStoreChunkingStrategy | None = None, ) -> VectorStoreFileObject: raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma") + + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + ) -> VectorStoreListFilesResponse: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma") + + async def openai_retrieve_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma") + + async def openai_update_vector_store_file( + self, + vector_store_id: str, + file_id: str, + attributes: dict[str, Any] | None = None, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma") + + async def openai_delete_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma") diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index 31a9535db..830b80bf1 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -26,7 +26,11 @@ from llama_stack.apis.vector_io import ( VectorStoreObject, VectorStoreSearchResponsePage, ) -from llama_stack.apis.vector_io.vector_io import VectorStoreChunkingStrategy, VectorStoreFileObject +from llama_stack.apis.vector_io.vector_io import ( + VectorStoreChunkingStrategy, + VectorStoreFileObject, + VectorStoreListFilesResponse, +) from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig from llama_stack.providers.utils.memory.vector_store import ( @@ -262,6 +266,34 @@ class MilvusVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate): ) -> VectorStoreFileObject: raise NotImplementedError("OpenAI Vector Stores API is not supported in Milvus") + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + ) -> VectorStoreListFilesResponse: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Milvus") + + async def openai_retrieve_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Milvus") + + async def openai_update_vector_store_file( + self, + vector_store_id: str, + file_id: str, + attributes: dict[str, Any] | None = None, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Milvus") + + async def openai_delete_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Milvus") + def generate_chunk_id(document_id: str, chunk_text: str) -> str: """Generate a unique chunk ID using a hash of document ID and chunk text.""" diff --git a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py index 1ebf861e2..2cf697f41 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py +++ b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py @@ -24,7 +24,11 @@ from llama_stack.apis.vector_io import ( VectorStoreObject, VectorStoreSearchResponsePage, ) -from llama_stack.apis.vector_io.vector_io import VectorStoreChunkingStrategy, VectorStoreFileObject +from llama_stack.apis.vector_io.vector_io import ( + VectorStoreChunkingStrategy, + VectorStoreFileObject, + VectorStoreListFilesResponse, +) from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig from llama_stack.providers.utils.memory.vector_store import ( @@ -263,3 +267,31 @@ class QdrantVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate): chunking_strategy: VectorStoreChunkingStrategy | None = None, ) -> VectorStoreFileObject: raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant") + + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + ) -> VectorStoreListFilesResponse: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant") + + async def openai_retrieve_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant") + + async def openai_update_vector_store_file( + self, + vector_store_id: str, + file_id: str, + attributes: dict[str, Any] | None = None, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant") + + async def openai_delete_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant") 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 13d2d7423..cd53c05c8 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -28,8 +28,11 @@ from llama_stack.apis.vector_io.vector_io import ( VectorStoreChunkingStrategy, VectorStoreChunkingStrategyAuto, VectorStoreChunkingStrategyStatic, + VectorStoreFileCounts, + VectorStoreFileDeleteResponse, VectorStoreFileLastError, VectorStoreFileObject, + VectorStoreListFilesResponse, ) from llama_stack.providers.utils.memory.vector_store import content_from_data_and_mime_type, make_overlapped_chunks @@ -70,6 +73,26 @@ class OpenAIVectorStoreMixin(ABC): """Delete vector store metadata from persistent storage.""" pass + @abstractmethod + async def _save_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: + """Save vector store file metadata to persistent storage.""" + pass + + @abstractmethod + async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]: + """Load vector store file metadata from persistent storage.""" + pass + + @abstractmethod + async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: + """Update vector store file metadata in persistent storage.""" + pass + + @abstractmethod + async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None: + """Delete vector store file metadata from persistent storage.""" + pass + @abstractmethod async def register_vector_db(self, vector_db: VectorDB) -> None: """Register a vector database (provider-specific implementation).""" @@ -136,18 +159,28 @@ class OpenAIVectorStoreMixin(ABC): await self.register_vector_db(vector_db) # Create OpenAI vector store metadata + status = "completed" + file_ids = file_ids or [] + file_counts = VectorStoreFileCounts( + cancelled=0, + completed=len(file_ids), + failed=0, + in_progress=0, + total=len(file_ids), + ) + # TODO: actually attach these files to the vector store... store_info = { "id": store_id, "object": "vector_store", "created_at": created_at, "name": store_id, "usage_bytes": 0, - "file_counts": {}, - "status": "completed", + "file_counts": file_counts.model_dump(), + "status": status, "expires_after": expires_after, "expires_at": None, "last_active_at": created_at, - "file_ids": file_ids or [], + "file_ids": file_ids, "chunking_strategy": chunking_strategy, } @@ -170,8 +203,8 @@ class OpenAIVectorStoreMixin(ABC): created_at=created_at, name=store_id, usage_bytes=0, - file_counts={}, - status="completed", + file_counts=file_counts, + status=status, expires_after=expires_after, expires_at=None, last_active_at=created_at, @@ -455,14 +488,20 @@ class OpenAIVectorStoreMixin(ABC): attributes: dict[str, Any] | None = None, chunking_strategy: VectorStoreChunkingStrategy | None = None, ) -> VectorStoreFileObject: + if vector_store_id not in self.openai_vector_stores: + raise ValueError(f"Vector store {vector_store_id} not found") + + store_info = self.openai_vector_stores[vector_store_id].copy() + attributes = attributes or {} chunking_strategy = chunking_strategy or VectorStoreChunkingStrategyAuto() + created_at = int(time.time()) vector_store_file_object = VectorStoreFileObject( id=file_id, attributes=attributes, chunking_strategy=chunking_strategy, - created_at=int(time.time()), + created_at=created_at, status="in_progress", vector_store_id=vector_store_id, ) @@ -510,6 +549,20 @@ class OpenAIVectorStoreMixin(ABC): vector_db_id=vector_store_id, chunks=chunks, ) + vector_store_file_object.status = "completed" + + # Create OpenAI vector store file metadata + file_info = vector_store_file_object.model_dump(exclude={"last_error"}) + + # Save to persistent storage (provider-specific) + await self._save_openai_vector_store_file(vector_store_id, file_id, file_info) + + # Update in-memory cache + store_info["file_ids"].append(file_id) + store_info["file_counts"]["completed"] += 1 + store_info["file_counts"]["total"] += 1 + self.openai_vector_stores[vector_store_id] = store_info + except Exception as e: logger.error(f"Error attaching file to vector store: {e}") vector_store_file_object.status = "failed" @@ -519,6 +572,84 @@ class OpenAIVectorStoreMixin(ABC): ) return vector_store_file_object - vector_store_file_object.status = "completed" - return vector_store_file_object + + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + ) -> VectorStoreListFilesResponse: + """List files in a vector store.""" + + if vector_store_id not in self.openai_vector_stores: + raise ValueError(f"Vector store {vector_store_id} not found") + + store_info = self.openai_vector_stores[vector_store_id] + + file_objects = [] + for file_id in store_info["file_ids"]: + file_info = await self._load_openai_vector_store_file(vector_store_id, file_id) + file_objects.append(VectorStoreFileObject(**file_info)) + + return VectorStoreListFilesResponse( + data=file_objects, + ) + + async def openai_retrieve_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileObject: + """Retrieves a vector store file.""" + if vector_store_id not in self.openai_vector_stores: + raise ValueError(f"Vector store {vector_store_id} not found") + + store_info = self.openai_vector_stores[vector_store_id] + if file_id not in store_info["file_ids"]: + raise ValueError(f"File {file_id} not found in vector store {vector_store_id}") + + file_info = await self._load_openai_vector_store_file(vector_store_id, file_id) + return VectorStoreFileObject(**file_info) + + async def openai_update_vector_store_file( + self, + vector_store_id: str, + file_id: str, + attributes: dict[str, Any], + ) -> VectorStoreFileObject: + """Updates a vector store file.""" + if vector_store_id not in self.openai_vector_stores: + raise ValueError(f"Vector store {vector_store_id} not found") + + store_info = self.openai_vector_stores[vector_store_id] + if file_id not in store_info["file_ids"]: + raise ValueError(f"File {file_id} not found in vector store {vector_store_id}") + + file_info = await self._load_openai_vector_store_file(vector_store_id, file_id) + file_info["attributes"] = attributes + await self._update_openai_vector_store_file(vector_store_id, file_id, file_info) + return VectorStoreFileObject(**file_info) + + async def openai_delete_vector_store_file( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileDeleteResponse: + """Deletes a vector store file.""" + if vector_store_id not in self.openai_vector_stores: + raise ValueError(f"Vector store {vector_store_id} not found") + + store_info = self.openai_vector_stores[vector_store_id].copy() + + file = await self.openai_retrieve_vector_store_file(vector_store_id, file_id) + await self._delete_openai_vector_store_file_from_storage(vector_store_id, file_id) + + # Update in-memory cache + store_info["file_ids"].remove(file_id) + store_info["file_counts"][file.status] -= 1 + store_info["file_counts"]["total"] -= 1 + self.openai_vector_stores[vector_store_id] = store_info + + return VectorStoreFileDeleteResponse( + id=file_id, + deleted=True, + ) diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index d9c4199ed..d092c6fcd 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/tests/integration/vector_io/test_openai_vector_stores.py @@ -6,8 +6,11 @@ import logging import time +from io import BytesIO import pytest +from llama_stack_client import BadRequestError, LlamaStackClient +from openai import BadRequestError as OpenAIBadRequestError from openai import OpenAI from llama_stack.apis.vector_io import Chunk @@ -73,11 +76,23 @@ def compat_client_with_empty_stores(compat_client): logger.warning("Failed to clear vector stores") pass + def clear_files(): + try: + response = compat_client.files.list() + for file in response.data: + compat_client.files.delete(file_id=file.id) + except Exception: + # If the API is not available or fails, just continue + logger.warning("Failed to clear files") + pass + clear_vector_stores() + clear_files() yield compat_client # Clean up after the test clear_vector_stores() + clear_files() def test_openai_create_vector_store(compat_client_with_empty_stores, client_with_models): @@ -423,3 +438,204 @@ def test_openai_vector_store_search_with_max_num_results( assert search_response is not None assert len(search_response.data) == 2 + + +def test_openai_vector_store_attach_file_response_attributes(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store attach file.""" + skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) + + if isinstance(compat_client_with_empty_stores, LlamaStackClient): + pytest.skip("Vector Store Files attach is not yet supported with LlamaStackClient") + + compat_client = compat_client_with_empty_stores + + # Create a vector store + vector_store = compat_client.vector_stores.create(name="test_store") + + # Create a file + test_content = b"This is a test file" + with BytesIO(test_content) as file_buffer: + file_buffer.name = "openai_test.txt" + file = compat_client.files.create(file=file_buffer, purpose="assistants") + + # Attach the file to the vector store + file_attach_response = compat_client.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + ) + + assert file_attach_response + assert file_attach_response.object == "vector_store.file" + assert file_attach_response.id == file.id + assert file_attach_response.vector_store_id == vector_store.id + assert file_attach_response.status == "completed" + assert file_attach_response.chunking_strategy.type == "auto" + assert file_attach_response.created_at > 0 + assert not file_attach_response.last_error + + updated_vector_store = compat_client.vector_stores.retrieve(vector_store_id=vector_store.id) + assert updated_vector_store.file_counts.completed == 1 + assert updated_vector_store.file_counts.total == 1 + assert updated_vector_store.file_counts.cancelled == 0 + assert updated_vector_store.file_counts.failed == 0 + assert updated_vector_store.file_counts.in_progress == 0 + + +def test_openai_vector_store_list_files(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store list files.""" + skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) + + if isinstance(compat_client_with_empty_stores, LlamaStackClient): + pytest.skip("Vector Store Files list is not yet supported with LlamaStackClient") + + compat_client = compat_client_with_empty_stores + + # Create a vector store + vector_store = compat_client.vector_stores.create(name="test_store") + + # Create some files and attach them to the vector store + file_ids = [] + for i in range(3): + with BytesIO(f"This is a test file {i}".encode()) as file_buffer: + file_buffer.name = f"openai_test_{i}.txt" + file = compat_client.files.create(file=file_buffer, purpose="assistants") + + compat_client.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + ) + file_ids.append(file.id) + + files_list = compat_client.vector_stores.files.list(vector_store_id=vector_store.id) + assert files_list + assert files_list.object == "list" + assert files_list.data + assert len(files_list.data) == 3 + assert file_ids == [file.id for file in files_list.data] + assert files_list.data[0].object == "vector_store.file" + assert files_list.data[0].vector_store_id == vector_store.id + assert files_list.data[0].status == "completed" + assert files_list.data[0].chunking_strategy.type == "auto" + assert files_list.data[0].created_at > 0 + assert not files_list.data[0].last_error + + updated_vector_store = compat_client.vector_stores.retrieve(vector_store_id=vector_store.id) + assert updated_vector_store.file_counts.completed == 3 + assert updated_vector_store.file_counts.total == 3 + assert updated_vector_store.file_counts.cancelled == 0 + assert updated_vector_store.file_counts.failed == 0 + assert updated_vector_store.file_counts.in_progress == 0 + + +def test_openai_vector_store_list_files_invalid_vector_store(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store list files with invalid vector store ID.""" + skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) + + if isinstance(compat_client_with_empty_stores, LlamaStackClient): + pytest.skip("Vector Store Files list is not yet supported with LlamaStackClient") + + compat_client = compat_client_with_empty_stores + + with pytest.raises((BadRequestError, OpenAIBadRequestError)): + compat_client.vector_stores.files.list(vector_store_id="abc123") + + +def test_openai_vector_store_delete_file(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store delete file.""" + skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) + + if isinstance(compat_client_with_empty_stores, LlamaStackClient): + pytest.skip("Vector Store Files list is not yet supported with LlamaStackClient") + + compat_client = compat_client_with_empty_stores + + # Create a vector store + vector_store = compat_client.vector_stores.create(name="test_store") + + # Create some files and attach them to the vector store + file_ids = [] + for i in range(3): + with BytesIO(f"This is a test file {i}".encode()) as file_buffer: + file_buffer.name = f"openai_test_{i}.txt" + file = compat_client.files.create(file=file_buffer, purpose="assistants") + + compat_client.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + ) + file_ids.append(file.id) + + files_list = compat_client.vector_stores.files.list(vector_store_id=vector_store.id) + assert len(files_list.data) == 3 + + # Delete the first file + delete_response = compat_client.vector_stores.files.delete(vector_store_id=vector_store.id, file_id=file_ids[0]) + assert delete_response + assert delete_response.id == file_ids[0] + assert delete_response.deleted is True + assert delete_response.object == "vector_store.file.deleted" + + updated_vector_store = compat_client.vector_stores.retrieve(vector_store_id=vector_store.id) + assert updated_vector_store.file_counts.completed == 2 + assert updated_vector_store.file_counts.total == 2 + assert updated_vector_store.file_counts.cancelled == 0 + assert updated_vector_store.file_counts.failed == 0 + assert updated_vector_store.file_counts.in_progress == 0 + + # Delete the second file + delete_response = compat_client.vector_stores.files.delete(vector_store_id=vector_store.id, file_id=file_ids[1]) + assert delete_response + assert delete_response.id == file_ids[1] + + updated_vector_store = compat_client.vector_stores.retrieve(vector_store_id=vector_store.id) + assert updated_vector_store.file_counts.completed == 1 + assert updated_vector_store.file_counts.total == 1 + assert updated_vector_store.file_counts.cancelled == 0 + assert updated_vector_store.file_counts.failed == 0 + assert updated_vector_store.file_counts.in_progress == 0 + + +def test_openai_vector_store_update_file(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store update file.""" + skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) + + if isinstance(compat_client_with_empty_stores, LlamaStackClient): + pytest.skip("Vector Store Files update is not yet supported with LlamaStackClient") + + compat_client = compat_client_with_empty_stores + + # Create a vector store + vector_store = compat_client.vector_stores.create(name="test_store") + + # Create a file + test_content = b"This is a test file" + with BytesIO(test_content) as file_buffer: + file_buffer.name = "openai_test.txt" + file = compat_client.files.create(file=file_buffer, purpose="assistants") + + # Attach the file to the vector store + file_attach_response = compat_client.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + attributes={"foo": "bar"}, + ) + + assert file_attach_response.status == "completed" + assert file_attach_response.attributes["foo"] == "bar" + + # Update the file's attributes + updated_response = compat_client.vector_stores.files.update( + vector_store_id=vector_store.id, + file_id=file.id, + attributes={"foo": "baz"}, + ) + + assert updated_response.status == "completed" + assert updated_response.attributes["foo"] == "baz" + + # Ensure we can retrieve the file and see the updated attributes + retrieved_file = compat_client.vector_stores.files.retrieve( + vector_store_id=vector_store.id, + file_id=file.id, + ) + assert retrieved_file.attributes["foo"] == "baz"