From 4905a8ee66a9b0d5cd9be5c7d606685a6cd07b4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Han?= Date: Thu, 27 Mar 2025 15:59:34 +0100 Subject: [PATCH] feat(api): Extend Files API to support registering ressources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit proposes extending the Files API to support registering and unregistering ressources. This is particularly useful if you want to point to pre-existing data and let Llama Stack know about them. Signed-off-by: Sébastien Han --- docs/_static/llama-stack-spec.html | 146 +++++++++++++++++++++++++++++ docs/_static/llama-stack-spec.yaml | 111 ++++++++++++++++++++++ llama_stack/apis/files/files.py | 60 ++++++++++++ 3 files changed, 317 insertions(+) diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index 4990d845e..3dce3d4aa 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -815,6 +815,56 @@ } ] }, + "post": { + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/FileRegistrationResponse" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "Files" + ], + "description": "Register an existing file with the provider.", + "parameters": [ + { + "name": "bucket", + "in": "path", + "description": "Storage location", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "key", + "in": "path", + "description": "File path relative to the storage location", + "required": true, + "schema": { + "type": "string" + } + } + ] + }, "delete": { "responses": { "200": { @@ -2540,6 +2590,47 @@ } } ] + }, + "post": { + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/BucketRegistrationResponse" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "Files" + ], + "description": "Register an existing storage location with the provider.", + "parameters": [ + { + "name": "bucket", + "in": "path", + "description": "Storage location", + "required": true, + "schema": { + "type": "string" + } + } + ] } }, "/v1/models": { @@ -9385,6 +9476,61 @@ ], "title": "RegisterBenchmarkRequest" }, + "BucketRegistrationResponse": { + "type": "object", + "properties": { + "bucket": { + "type": "string", + "description": "The registered storage location URI (e.g., \"s3://my-bucket\" or \"file:///data\")" + }, + "created_at": { + "type": "integer", + "description": "Timestamp of registration" + }, + "status": { + "type": "string", + "description": "Current status of the storage location" + } + }, + "additionalProperties": false, + "required": [ + "bucket", + "created_at", + "status" + ], + "title": "BucketRegistrationResponse", + "description": "Response after registering a storage location." + }, + "FileRegistrationResponse": { + "type": "object", + "properties": { + "bucket": { + "type": "string", + "description": "The storage location URI (e.g., \"s3://my-bucket\" or \"file:///data\")" + }, + "key": { + "type": "string", + "description": "The file path relative to the storage location" + }, + "created_at": { + "type": "integer", + "description": "Timestamp of registration" + }, + "status": { + "type": "string", + "description": "Current status of the file" + } + }, + "additionalProperties": false, + "required": [ + "bucket", + "key", + "created_at", + "status" + ], + "title": "FileRegistrationResponse", + "description": "Response after registering a file." + }, "RegisterDatasetRequest": { "type": "object", "properties": { diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index ba3868560..1fd44789b 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -553,6 +553,42 @@ paths: required: true schema: type: string + post: + responses: + '200': + description: OK + content: + application/json: + schema: + $ref: '#/components/schemas/FileRegistrationResponse' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - Files + description: >- + Register an existing file with the provider. + parameters: + - name: bucket + in: path + description: Storage location + required: true + schema: + type: string + - name: key + in: path + description: >- + File path relative to the storage location + required: true + schema: + type: string delete: responses: '200': @@ -1739,6 +1775,35 @@ paths: required: true schema: type: string + post: + responses: + '200': + description: OK + content: + application/json: + schema: + $ref: '#/components/schemas/BucketRegistrationResponse' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - Files + description: >- + Register an existing storage location with the provider. + parameters: + - name: bucket + in: path + description: Storage location + required: true + schema: + type: string /v1/models: get: responses: @@ -6391,6 +6456,52 @@ components: - dataset_id - scoring_functions title: RegisterBenchmarkRequest + BucketRegistrationResponse: + type: object + properties: + bucket: + type: string + description: >- + The registered storage location URI (e.g., "s3://my-bucket" or "file:///data") + created_at: + type: integer + description: Timestamp of registration + status: + type: string + description: Current status of the storage location + additionalProperties: false + required: + - bucket + - created_at + - status + title: BucketRegistrationResponse + description: >- + Response after registering a storage location. + FileRegistrationResponse: + type: object + properties: + bucket: + type: string + description: >- + The storage location URI (e.g., "s3://my-bucket" or "file:///data") + key: + type: string + description: >- + The file path relative to the storage location + created_at: + type: integer + description: Timestamp of registration + status: + type: string + description: Current status of the file + additionalProperties: false + required: + - bucket + - key + - created_at + - status + title: FileRegistrationResponse + description: Response after registering a file. RegisterDatasetRequest: type: object properties: diff --git a/llama_stack/apis/files/files.py b/llama_stack/apis/files/files.py index ef8b65829..33f8cf0f3 100644 --- a/llama_stack/apis/files/files.py +++ b/llama_stack/apis/files/files.py @@ -77,6 +77,38 @@ class ListFileResponse(BaseModel): data: List[FileResponse] +@json_schema_type +class BucketRegistrationResponse(BaseModel): + """ + Response after registering a storage location. + + :param bucket: The registered storage location URI (e.g., "s3://my-bucket" or "file:///data") + :param created_at: Timestamp of registration + :param status: Current status of the storage location + """ + + bucket: str + created_at: int + status: str + + +@json_schema_type +class FileRegistrationResponse(BaseModel): + """ + Response after registering a file. + + :param bucket: The storage location URI (e.g., "s3://my-bucket" or "file:///data") + :param key: The file path relative to the storage location + :param created_at: Timestamp of registration + :param status: Current status of the file + """ + + bucket: str + key: str + created_at: int + status: str + + @runtime_checkable @trace_protocol class Files(Protocol): @@ -172,3 +204,31 @@ class Files(Protocol): :param key: Key under which the file is stored (valid chars: a-zA-Z0-9_-/.) """ ... + + @webmethod(route="/files/{bucket}", method="PUT") + async def register_bucket( + self, + bucket: str, + ) -> BucketRegistrationResponse: + """ + Register an existing storage location with the provider. + + :param bucket: Storage location + :raises: ValidationError if URI is invalid or contains invalid characters + """ + ... + + @webmethod(route="/files/{bucket}/{key:path}", method="PUT") + async def register_bucket_file( + self, + bucket: str, + key: str, + ) -> FileRegistrationResponse: + """ + Register an existing file with the provider. + + :param bucket: Storage location + :param key: File path relative to the storage location + :raises: ValidationError if URI is invalid or contains invalid characters + """ + ...