feat(api): Extend Files API to support registering ressources

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 <seb@redhat.com>
This commit is contained in:
Sébastien Han 2025-03-27 15:59:34 +01:00
parent 935e706b15
commit 4905a8ee66
No known key found for this signature in database
3 changed files with 317 additions and 0 deletions

View file

@ -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": { "delete": {
"responses": { "responses": {
"200": { "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": { "/v1/models": {
@ -9385,6 +9476,61 @@
], ],
"title": "RegisterBenchmarkRequest" "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": { "RegisterDatasetRequest": {
"type": "object", "type": "object",
"properties": { "properties": {

View file

@ -553,6 +553,42 @@ paths:
required: true required: true
schema: schema:
type: string 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: delete:
responses: responses:
'200': '200':
@ -1739,6 +1775,35 @@ paths:
required: true required: true
schema: schema:
type: string 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: /v1/models:
get: get:
responses: responses:
@ -6391,6 +6456,52 @@ components:
- dataset_id - dataset_id
- scoring_functions - scoring_functions
title: RegisterBenchmarkRequest 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: RegisterDatasetRequest:
type: object type: object
properties: properties:

View file

@ -77,6 +77,38 @@ class ListFileResponse(BaseModel):
data: List[FileResponse] 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 @runtime_checkable
@trace_protocol @trace_protocol
class Files(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_-/.) :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
"""
...