diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index 5407f9808..0a5caa3d1 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -3241,6 +3241,87 @@ } }, "/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" + } + }, + { + "name": "limit", + "in": "query", + "required": false, + "schema": { + "type": "integer" + } + }, + { + "name": "order", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "after", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "before", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "filter", + "in": "query", + "required": false, + "schema": { + "$ref": "#/components/schemas/VectorStoreFileStatus" + } + } + ] + }, "post": { "responses": { "200": { @@ -3666,6 +3747,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": { @@ -3909,6 +4152,58 @@ ] } }, + "/v1/openai/v1/vector_stores/{vector_store_id}/files/{file_id}/content": { + "get": { + "responses": { + "200": { + "description": "A list of InterleavedContent representing the file contents.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/VectorStoreFileContentsResponse" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "VectorIO" + ], + "description": "Retrieves the contents of 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" + } + } + ] + } + }, "/v1/openai/v1/vector_stores/{vector_store_id}/search": { "post": { "responses": { @@ -12102,24 +12397,7 @@ "$ref": "#/components/schemas/VectorStoreFileLastError" }, "status": { - "oneOf": [ - { - "type": "string", - "const": "completed" - }, - { - "type": "string", - "const": "in_progress" - }, - { - "type": "string", - "const": "cancelled" - }, - { - "type": "string", - "const": "failed" - } - ] + "$ref": "#/components/schemas/VectorStoreFileStatus" }, "usage_bytes": { "type": "integer", @@ -12143,6 +12421,26 @@ "title": "VectorStoreFileObject", "description": "OpenAI Vector Store File object." }, + "VectorStoreFileStatus": { + "oneOf": [ + { + "type": "string", + "const": "completed" + }, + { + "type": "string", + "const": "in_progress" + }, + { + "type": "string", + "const": "cancelled" + }, + { + "type": "string", + "const": "failed" + } + ] + }, "OpenAIJSONSchema": { "type": "object", "properties": { @@ -12969,6 +13267,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 +13317,7 @@ "default": 0 }, "file_counts": { - "type": "object", - "additionalProperties": { - "type": "integer" - } + "$ref": "#/components/schemas/VectorStoreFileCounts" }, "status": { "type": "string", @@ -13120,6 +13444,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 +13696,39 @@ "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" + } + }, + "first_id": { + "type": "string" + }, + "last_id": { + "type": "string" + }, + "has_more": { + "type": "boolean", + "default": false + } + }, + "additionalProperties": false, + "required": [ + "object", + "data", + "has_more" + ], + "title": "VectorStoreListFilesResponse", + "description": "Response from listing vector stores." + }, "OpenAIModel": { "type": "object", "properties": { @@ -13429,6 +13810,75 @@ "type": "object", "title": "Response" }, + "VectorStoreContent": { + "type": "object", + "properties": { + "type": { + "type": "string", + "const": "text" + }, + "text": { + "type": "string" + } + }, + "additionalProperties": false, + "required": [ + "type", + "text" + ], + "title": "VectorStoreContent" + }, + "VectorStoreFileContentsResponse": { + "type": "object", + "properties": { + "file_id": { + "type": "string" + }, + "filename": { + "type": "string" + }, + "attributes": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } + }, + "content": { + "type": "array", + "items": { + "$ref": "#/components/schemas/VectorStoreContent" + } + } + }, + "additionalProperties": false, + "required": [ + "file_id", + "filename", + "attributes", + "content" + ], + "title": "VectorStoreFileContentsResponse", + "description": "Response from retrieving the contents of a vector store file." + }, "OpenaiSearchVectorStoreRequest": { "type": "object", "properties": { @@ -13501,24 +13951,6 @@ ], "title": "OpenaiSearchVectorStoreRequest" }, - "VectorStoreContent": { - "type": "object", - "properties": { - "type": { - "type": "string", - "const": "text" - }, - "text": { - "type": "string" - } - }, - "additionalProperties": false, - "required": [ - "type", - "text" - ], - "title": "VectorStoreContent" - }, "VectorStoreSearchResponse": { "type": "object", "properties": { @@ -13661,6 +14093,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..c115e1df2 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -2264,6 +2264,61 @@ 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 + - name: limit + in: query + required: false + schema: + type: integer + - name: order + in: query + required: false + schema: + type: string + - name: after + in: query + required: false + schema: + type: string + - name: before + in: query + required: false + schema: + type: string + - name: filter + in: query + required: false + schema: + $ref: '#/components/schemas/VectorStoreFileStatus' post: responses: '200': @@ -2572,6 +2627,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: @@ -2762,6 +2932,44 @@ paths: required: true schema: type: string + /v1/openai/v1/vector_stores/{vector_store_id}/files/{file_id}/content: + get: + responses: + '200': + description: >- + A list of InterleavedContent representing the file contents. + content: + application/json: + schema: + $ref: '#/components/schemas/VectorStoreFileContentsResponse' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - VectorIO + description: >- + Retrieves the contents of 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 /v1/openai/v1/vector_stores/{vector_store_id}/search: post: responses: @@ -8458,15 +8666,7 @@ components: last_error: $ref: '#/components/schemas/VectorStoreFileLastError' status: - oneOf: - - type: string - const: completed - - type: string - const: in_progress - - type: string - const: cancelled - - type: string - const: failed + $ref: '#/components/schemas/VectorStoreFileStatus' usage_bytes: type: integer default: 0 @@ -8484,6 +8684,16 @@ components: - vector_store_id title: VectorStoreFileObject description: OpenAI Vector Store File object. + VectorStoreFileStatus: + oneOf: + - type: string + const: completed + - type: string + const: in_progress + - type: string + const: cancelled + - type: string + const: failed OpenAIJSONSchema: type: object properties: @@ -9031,6 +9241,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 +9278,7 @@ components: type: integer default: 0 file_counts: - type: object - additionalProperties: - type: integer + $ref: '#/components/schemas/VectorStoreFileCounts' status: type: string default: completed @@ -9129,6 +9358,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 +9568,30 @@ 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' + first_id: + type: string + last_id: + type: string + has_more: + type: boolean + default: false + additionalProperties: false + required: + - object + - data + - has_more + title: VectorStoreListFilesResponse + description: Response from listing vector stores. OpenAIModel: type: object properties: @@ -9379,6 +9651,49 @@ components: Response: type: object title: Response + VectorStoreContent: + type: object + properties: + type: + type: string + const: text + text: + type: string + additionalProperties: false + required: + - type + - text + title: VectorStoreContent + VectorStoreFileContentsResponse: + type: object + properties: + file_id: + type: string + filename: + type: string + attributes: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + content: + type: array + items: + $ref: '#/components/schemas/VectorStoreContent' + additionalProperties: false + required: + - file_id + - filename + - attributes + - content + title: VectorStoreFileContentsResponse + description: >- + Response from retrieving the contents of a vector store file. OpenaiSearchVectorStoreRequest: type: object properties: @@ -9426,19 +9741,6 @@ components: required: - query title: OpenaiSearchVectorStoreRequest - VectorStoreContent: - type: object - properties: - type: - type: string - const: text - text: - type: string - additionalProperties: false - required: - - type - - text - title: VectorStoreContent VectorStoreSearchResponse: type: object properties: @@ -9524,6 +9826,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..dbea12d5f 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 @@ -168,6 +177,10 @@ class VectorStoreFileLastError(BaseModel): message: str +VectorStoreFileStatus = Literal["completed"] | Literal["in_progress"] | Literal["cancelled"] | Literal["failed"] +register_schema(VectorStoreFileStatus, name="VectorStoreFileStatus") + + @json_schema_type class VectorStoreFileObject(BaseModel): """OpenAI Vector Store File object.""" @@ -178,11 +191,41 @@ class VectorStoreFileObject(BaseModel): chunking_strategy: VectorStoreChunkingStrategy created_at: int last_error: VectorStoreFileLastError | None = None - status: Literal["completed"] | Literal["in_progress"] | Literal["cancelled"] | Literal["failed"] + status: VectorStoreFileStatus usage_bytes: int = 0 vector_store_id: str +@json_schema_type +class VectorStoreListFilesResponse(BaseModel): + """Response from listing vector stores.""" + + object: str = "list" + data: list[VectorStoreFileObject] + first_id: str | None = None + last_id: str | None = None + has_more: bool = False + + +@json_schema_type +class VectorStoreFileContentsResponse(BaseModel): + """Response from retrieving the contents of a vector store file.""" + + file_id: str + filename: str + attributes: dict[str, Any] + content: list[VectorStoreContent] + + +@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 +401,78 @@ 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, + limit: int | None = 20, + order: str | None = "desc", + after: str | None = None, + before: str | None = None, + filter: VectorStoreFileStatus | None = None, + ) -> 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}/content", method="GET") + async def openai_retrieve_vector_store_file_contents( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileContentsResponse: + """Retrieves the contents of 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 list of InterleavedContent representing the file contents. + """ + ... + + @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..643029d60 100644 --- a/llama_stack/distribution/routers/vector_io.py +++ b/llama_stack/distribution/routers/vector_io.py @@ -21,7 +21,13 @@ 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, + VectorStoreFileContentsResponse, + VectorStoreFileDeleteResponse, + VectorStoreFileObject, + VectorStoreFileStatus, +) from llama_stack.log import get_logger from llama_stack.providers.datatypes import HealthResponse, HealthStatus, RoutingTable @@ -279,6 +285,81 @@ class VectorIORouter(VectorIO): chunking_strategy=chunking_strategy, ) + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + limit: int | None = 20, + order: str | None = "desc", + after: str | None = None, + before: str | None = None, + filter: VectorStoreFileStatus | None = None, + ) -> 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, + limit=limit, + order=order, + after=after, + before=before, + filter=filter, + ) + + 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_retrieve_vector_store_file_contents( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileContentsResponse: + logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_contents: {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_contents( + 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..12f4d6ad0 100644 --- a/llama_stack/providers/inline/vector_io/faiss/faiss.py +++ b/llama_stack/providers/inline/vector_io/faiss/faiss.py @@ -45,6 +45,8 @@ 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}::" +OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX = f"openai_vector_stores_files_contents:{VERSION}::" class FaissIndex(EmbeddingIndex): @@ -283,3 +285,39 @@ 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], file_contents: list[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)) + content_key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}" + await self.kvstore.set(key=content_key, value=json.dumps(file_contents)) + + 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 _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]: + """Load vector store file contents from kvstore.""" + assert self.kvstore is not None + key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_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..d832e56f5 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,23 @@ 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, + PRIMARY KEY (store_id, file_id) + ); + """) + cur.execute(""" + CREATE TABLE IF NOT EXISTS openai_vector_store_files_contents ( + store_id TEXT, + file_id TEXT, + contents TEXT, + PRIMARY KEY (store_id, file_id) + ); + """) connection.commit() # Load any existing vector DB registrations. cur.execute("SELECT metadata FROM vector_dbs") @@ -615,6 +632,118 @@ 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], file_contents: list[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)), + ) + cur.execute( + "INSERT OR REPLACE INTO openai_vector_store_files_contents (store_id, file_id, contents) VALUES (?, ?, ?)", + (store_id, file_id, json.dumps(file_contents)), + ) + 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() + 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 _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]: + """Load vector store file contents from SQLite database.""" + + def _load(): + connection = _create_sqlite_connection(self.config.db_path) + cur = connection.cursor() + try: + cur.execute( + "SELECT contents FROM openai_vector_store_files_contents WHERE store_id = ? AND file_id = ?", + (store_id, file_id), + ) + row = cur.fetchone() + if row is None: + return None + (contents,) = row + return contents + finally: + cur.close() + connection.close() + + stored_contents = await asyncio.to_thread(_load) + return json.loads(stored_contents) if stored_contents 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..9f206aaee 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -24,7 +24,12 @@ 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, + VectorStoreFileContentsResponse, + 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 +268,38 @@ 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_retrieve_vector_store_file_contents( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileContentsResponse: + 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..737a46bac 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -26,7 +26,12 @@ 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, + VectorStoreFileContentsResponse, + 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 +267,41 @@ 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_retrieve_vector_store_file_contents( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileContentsResponse: + 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..e00fdf84e 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py +++ b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py @@ -24,7 +24,12 @@ 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, + VectorStoreFileContentsResponse, + 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 +268,38 @@ 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_retrieve_vector_store_file_contents( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileContentsResponse: + 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..9c0e1dbe7 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -4,6 +4,7 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. +import asyncio import logging import mimetypes import time @@ -12,6 +13,7 @@ from abc import ABC, abstractmethod from typing import Any from llama_stack.apis.files import Files +from llama_stack.apis.files.files import OpenAIFileObject from llama_stack.apis.vector_dbs import VectorDB from llama_stack.apis.vector_io import ( QueryChunksResponse, @@ -28,8 +30,13 @@ from llama_stack.apis.vector_io.vector_io import ( VectorStoreChunkingStrategy, VectorStoreChunkingStrategyAuto, VectorStoreChunkingStrategyStatic, + VectorStoreFileContentsResponse, + VectorStoreFileCounts, + VectorStoreFileDeleteResponse, VectorStoreFileLastError, VectorStoreFileObject, + VectorStoreFileStatus, + VectorStoreListFilesResponse, ) from llama_stack.providers.utils.memory.vector_store import content_from_data_and_mime_type, make_overlapped_chunks @@ -70,6 +77,33 @@ 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], file_contents: list[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 _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]: + """Load vector store file contents 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 +170,28 @@ class OpenAIVectorStoreMixin(ABC): await self.register_vector_db(vector_db) # Create OpenAI vector store metadata + status = "completed" + + # Start with no files attached and update later + file_counts = VectorStoreFileCounts( + cancelled=0, + completed=0, + failed=0, + in_progress=0, + total=0, + ) 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": [], "chunking_strategy": chunking_strategy, } @@ -165,18 +209,14 @@ class OpenAIVectorStoreMixin(ABC): # Store in memory cache self.openai_vector_stores[store_id] = store_info - return VectorStoreObject( - id=store_id, - created_at=created_at, - name=store_id, - usage_bytes=0, - file_counts={}, - status="completed", - expires_after=expires_after, - expires_at=None, - last_active_at=created_at, - metadata=metadata, - ) + # Now that our vector store is created, attach any files that were provided + file_ids = file_ids or [] + tasks = [self.openai_attach_file_to_vector_store(store_id, file_id) for file_id in file_ids] + await asyncio.gather(*tasks) + + # Get the updated store info and return it + store_info = self.openai_vector_stores[store_id] + return VectorStoreObject.model_validate(store_info) async def openai_list_vector_stores( self, @@ -346,33 +386,7 @@ class OpenAIVectorStoreMixin(ABC): if not self._matches_filters(chunk.metadata, filters): continue - # content is InterleavedContent - if isinstance(chunk.content, str): - content = [ - VectorStoreContent( - type="text", - text=chunk.content, - ) - ] - elif isinstance(chunk.content, list): - # TODO: Add support for other types of content - content = [ - VectorStoreContent( - type="text", - text=item.text, - ) - for item in chunk.content - if item.type == "text" - ] - else: - if chunk.content.type != "text": - raise ValueError(f"Unsupported content type: {chunk.content.type}") - content = [ - VectorStoreContent( - type="text", - text=chunk.content.text, - ) - ] + content = self._chunk_to_vector_store_content(chunk) response_data_item = VectorStoreSearchResponse( file_id=chunk.metadata.get("file_id", ""), @@ -448,6 +462,36 @@ class OpenAIVectorStoreMixin(ABC): # Unknown filter type, default to no match raise ValueError(f"Unsupported filter type: {filter_type}") + def _chunk_to_vector_store_content(self, chunk: Chunk) -> list[VectorStoreContent]: + # content is InterleavedContent + if isinstance(chunk.content, str): + content = [ + VectorStoreContent( + type="text", + text=chunk.content, + ) + ] + elif isinstance(chunk.content, list): + # TODO: Add support for other types of content + content = [ + VectorStoreContent( + type="text", + text=item.text, + ) + for item in chunk.content + if item.type == "text" + ] + else: + if chunk.content.type != "text": + raise ValueError(f"Unsupported content type: {chunk.content.type}") + content = [ + VectorStoreContent( + type="text", + text=chunk.content.text, + ) + ] + return content + async def openai_attach_file_to_vector_store( self, vector_store_id: str, @@ -455,14 +499,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") + attributes = attributes or {} chunking_strategy = chunking_strategy or VectorStoreChunkingStrategyAuto() + created_at = int(time.time()) + chunks: list[Chunk] = [] + file_response: OpenAIFileObject | None = None 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, ) @@ -504,12 +554,12 @@ class OpenAIVectorStoreMixin(ABC): code="server_error", message="No chunks were generated from the file", ) - return vector_store_file_object - - await self.insert_chunks( - vector_db_id=vector_store_id, - chunks=chunks, - ) + else: + await self.insert_chunks( + vector_db_id=vector_store_id, + chunks=chunks, + ) + vector_store_file_object.status = "completed" except Exception as e: logger.error(f"Error attaching file to vector store: {e}") vector_store_file_object.status = "failed" @@ -517,8 +567,171 @@ class OpenAIVectorStoreMixin(ABC): code="server_error", message=str(e), ) - return vector_store_file_object - vector_store_file_object.status = "completed" + # Create OpenAI vector store file metadata + file_info = vector_store_file_object.model_dump(exclude={"last_error"}) + file_info["filename"] = file_response.filename if file_response else "" + + # Save vector store file to persistent storage (provider-specific) + dict_chunks = [c.model_dump() for c in chunks] + await self._save_openai_vector_store_file(vector_store_id, file_id, file_info, dict_chunks) + + # Update file_ids and file_counts in vector store metadata + store_info = self.openai_vector_stores[vector_store_id].copy() + store_info["file_ids"].append(file_id) + store_info["file_counts"]["total"] += 1 + store_info["file_counts"][vector_store_file_object.status] += 1 + + # Save updated vector store to persistent storage + await self._save_openai_vector_store(vector_store_id, store_info) + + # Update vector store in-memory cache + self.openai_vector_stores[vector_store_id] = store_info return vector_store_file_object + + async def openai_list_files_in_vector_store( + self, + vector_store_id: str, + limit: int | None = 20, + order: str | None = "desc", + after: str | None = None, + before: str | None = None, + filter: VectorStoreFileStatus | None = None, + ) -> VectorStoreListFilesResponse: + """List files in a vector store.""" + limit = limit or 20 + order = order or "desc" + + 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: list[VectorStoreFileObject] = [] + for file_id in store_info["file_ids"]: + file_info = await self._load_openai_vector_store_file(vector_store_id, file_id) + file_object = VectorStoreFileObject(**file_info) + if filter and file_object.status != filter: + continue + file_objects.append(file_object) + + # Sort by created_at + reverse_order = order == "desc" + file_objects.sort(key=lambda x: x.created_at, reverse=reverse_order) + + # Apply cursor-based pagination + if after: + after_index = next((i for i, file in enumerate(file_objects) if file.id == after), -1) + if after_index >= 0: + 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)) + file_objects = file_objects[:before_index] + + # Apply limit + limited_files = file_objects[:limit] + + # Determine pagination info + has_more = len(file_objects) > limit + first_id = file_objects[0].id if file_objects else None + last_id = file_objects[-1].id if file_objects else None + + return VectorStoreListFilesResponse( + data=limited_files, + has_more=has_more, + first_id=first_id, + last_id=last_id, + ) + + 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_retrieve_vector_store_file_contents( + self, + vector_store_id: str, + file_id: str, + ) -> VectorStoreFileContentsResponse: + """Retrieves the contents of a vector store file.""" + if vector_store_id not in self.openai_vector_stores: + raise ValueError(f"Vector store {vector_store_id} not found") + + file_info = await self._load_openai_vector_store_file(vector_store_id, file_id) + dict_chunks = await self._load_openai_vector_store_file_contents(vector_store_id, file_id) + chunks = [Chunk.model_validate(c) for c in dict_chunks] + content = [] + for chunk in chunks: + content.extend(self._chunk_to_vector_store_content(chunk)) + return VectorStoreFileContentsResponse( + file_id=file_id, + filename=file_info.get("filename", ""), + attributes=file_info.get("attributes", {}), + content=content, + ) + + async def openai_update_vector_store_file( + 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) + + # TODO: We need to actually delete the embeddings from the underlying vector store... + # Also uncomment the corresponding integration test marked as xfail + # + # test_openai_vector_store_delete_file_removes_from_vector_store in + # tests/integration/vector_io/test_openai_vector_stores.py + + # 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 + + # Save updated vector store to persistent storage + await self._save_openai_vector_store(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..4856455c4 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,369 @@ 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(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"The secret string is foobazbar." + 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 + + # Search using OpenAI API to confirm our file attached + search_response = compat_client.vector_stores.search( + vector_store_id=vector_store.id, query="What is the secret string?", max_num_results=1 + ) + assert search_response is not None + assert len(search_response.data) > 0 + top_result = search_response.data[0] + top_content = top_result.content[0].text + assert "foobazbar" in top_content.lower() + + +def test_openai_vector_store_attach_files_on_creation(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store attach files on creation.""" + 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 some files and attach them to the vector store + valid_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") + valid_file_ids.append(file.id) + + # include an invalid file ID so we can test failed status + failed_file_id = "invalid_file_id" + file_ids = valid_file_ids + [failed_file_id] + num_failed = len(file_ids) - len(valid_file_ids) + + # Create a vector store + vector_store = compat_client.vector_stores.create( + name="test_store", + file_ids=file_ids, + ) + + assert vector_store.file_counts.completed == len(valid_file_ids) + assert vector_store.file_counts.total == len(file_ids) + assert vector_store.file_counts.cancelled == 0 + assert vector_store.file_counts.failed == num_failed + assert vector_store.file_counts.in_progress == 0 + + files_list = compat_client.vector_stores.files.list(vector_store_id=vector_store.id) + assert len(files_list.data) == len(file_ids) + assert set(file_ids) == {file.id for file in files_list.data} + for file in files_list.data: + if file.id in valid_file_ids: + assert file.status == "completed" + else: + assert file.status == "failed" + + failed_list = compat_client.vector_stores.files.list(vector_store_id=vector_store.id, filter="failed") + assert len(failed_list.data) == num_failed + assert failed_file_id == failed_list.data[0].id + + # Delete the invalid file + delete_response = compat_client.vector_stores.files.delete(vector_store_id=vector_store.id, file_id=failed_file_id) + assert delete_response.id == failed_file_id + + updated_vector_store = compat_client.vector_stores.retrieve(vector_store_id=vector_store.id) + assert updated_vector_store.file_counts.completed == len(valid_file_ids) + assert updated_vector_store.file_counts.total == len(valid_file_ids) + assert updated_vector_store.file_counts.failed == 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 not files_list.has_more + assert len(files_list.data) == 3 + assert set(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 files_list.first_id == files_list.data[0].id + assert not files_list.data[0].last_error + + first_page = compat_client.vector_stores.files.list(vector_store_id=vector_store.id, limit=2) + assert first_page.has_more + assert len(first_page.data) == 2 + assert first_page.first_id == first_page.data[0].id + assert first_page.last_id != first_page.data[-1].id + + next_page = compat_client.vector_stores.files.list( + vector_store_id=vector_store.id, limit=2, after=first_page.data[-1].id + ) + assert not next_page.has_more + assert len(next_page.data) == 1 + + 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_retrieve_file_contents(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store retrieve file contents.""" + skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) + + if isinstance(compat_client_with_empty_stores, LlamaStackClient): + pytest.skip("Vector Store Files retrieve contents 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" + file_name = "openai_test.txt" + attributes = {"foo": "bar"} + with BytesIO(test_content) as file_buffer: + file_buffer.name = file_name + 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=attributes, + ) + + assert file_attach_response.status == "completed" + + file_contents = compat_client.vector_stores.files.content( + vector_store_id=vector_store.id, + file_id=file.id, + ) + + assert file_contents + assert file_contents.content[0]["type"] == "text" + assert file_contents.content[0]["text"] == test_content.decode("utf-8") + assert file_contents.filename == file_name + assert file_contents.attributes == attributes + + +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 + + +# TODO: Remove this xfail once we have a way to remove embeddings from vector store +@pytest.mark.xfail(reason="Vector Store Files delete doesn't remove embeddings from vecntor store", strict=True) +def test_openai_vector_store_delete_file_removes_from_vector_store(compat_client_with_empty_stores, client_with_models): + """Test OpenAI vector store delete file removes from vector store.""" + 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"The secret string is foobazbar." + 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.status == "completed" + + # Search using OpenAI API to confirm our file attached + search_response = compat_client.vector_stores.search( + vector_store_id=vector_store.id, query="What is the secret string?", max_num_results=1 + ) + assert "foobazbar" in search_response.data[0].content[0].text.lower() + + # Delete the file + compat_client.vector_stores.files.delete(vector_store_id=vector_store.id, file_id=file.id) + + # Search using OpenAI API to confirm our file deleted + search_response = compat_client.vector_stores.search( + vector_store_id=vector_store.id, query="What is the secret string?", max_num_results=1 + ) + assert not search_response.data + + +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"