mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-28 02:53:30 +00:00
feat: reference implementation for files API (#2330)
Some checks failed
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 2s
Integration Tests / test-matrix (http, post_training) (push) Failing after 9s
Integration Tests / test-matrix (http, agents) (push) Failing after 10s
Integration Tests / test-matrix (http, providers) (push) Failing after 8s
Integration Tests / test-matrix (http, inference) (push) Failing after 11s
Integration Tests / test-matrix (http, inspect) (push) Failing after 10s
Integration Tests / test-matrix (http, datasets) (push) Failing after 11s
Integration Tests / test-matrix (library, datasets) (push) Failing after 8s
Integration Tests / test-matrix (http, scoring) (push) Failing after 10s
Integration Tests / test-matrix (library, inference) (push) Failing after 8s
Integration Tests / test-matrix (library, agents) (push) Failing after 10s
Integration Tests / test-matrix (http, tool_runtime) (push) Failing after 11s
Integration Tests / test-matrix (library, inspect) (push) Failing after 8s
Test External Providers / test-external-providers (venv) (push) Failing after 7s
Integration Tests / test-matrix (library, post_training) (push) Failing after 9s
Integration Tests / test-matrix (library, scoring) (push) Failing after 8s
Integration Tests / test-matrix (library, tool_runtime) (push) Failing after 8s
Integration Tests / test-matrix (library, providers) (push) Failing after 9s
Unit Tests / unit-tests (3.11) (push) Failing after 7s
Unit Tests / unit-tests (3.10) (push) Failing after 7s
Unit Tests / unit-tests (3.12) (push) Failing after 8s
Unit Tests / unit-tests (3.13) (push) Failing after 8s
Update ReadTheDocs / update-readthedocs (push) Failing after 6s
Pre-commit / pre-commit (push) Successful in 53s
Some checks failed
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 2s
Integration Tests / test-matrix (http, post_training) (push) Failing after 9s
Integration Tests / test-matrix (http, agents) (push) Failing after 10s
Integration Tests / test-matrix (http, providers) (push) Failing after 8s
Integration Tests / test-matrix (http, inference) (push) Failing after 11s
Integration Tests / test-matrix (http, inspect) (push) Failing after 10s
Integration Tests / test-matrix (http, datasets) (push) Failing after 11s
Integration Tests / test-matrix (library, datasets) (push) Failing after 8s
Integration Tests / test-matrix (http, scoring) (push) Failing after 10s
Integration Tests / test-matrix (library, inference) (push) Failing after 8s
Integration Tests / test-matrix (library, agents) (push) Failing after 10s
Integration Tests / test-matrix (http, tool_runtime) (push) Failing after 11s
Integration Tests / test-matrix (library, inspect) (push) Failing after 8s
Test External Providers / test-external-providers (venv) (push) Failing after 7s
Integration Tests / test-matrix (library, post_training) (push) Failing after 9s
Integration Tests / test-matrix (library, scoring) (push) Failing after 8s
Integration Tests / test-matrix (library, tool_runtime) (push) Failing after 8s
Integration Tests / test-matrix (library, providers) (push) Failing after 9s
Unit Tests / unit-tests (3.11) (push) Failing after 7s
Unit Tests / unit-tests (3.10) (push) Failing after 7s
Unit Tests / unit-tests (3.12) (push) Failing after 8s
Unit Tests / unit-tests (3.13) (push) Failing after 8s
Update ReadTheDocs / update-readthedocs (push) Failing after 6s
Pre-commit / pre-commit (push) Successful in 53s
# What does this PR do? TSIA Added Files provider to the fireworks template. Might want to add to all templates as a follow-up. ## Test Plan llama-stack pytest tests/unit/files/test_files.py llama-stack llama stack build --template fireworks --image-type conda --run LLAMA_STACK_CONFIG=http://localhost:8321 pytest -s -v tests/integration/files/
This commit is contained in:
parent
ba25c5e7e1
commit
3c9a10d2fe
18 changed files with 3041 additions and 2315 deletions
|
@ -18,6 +18,7 @@ The `llamastack/distribution-fireworks` distribution consists of the following p
|
|||
| agents | `inline::meta-reference` |
|
||||
| datasetio | `remote::huggingface`, `inline::localfs` |
|
||||
| eval | `inline::meta-reference` |
|
||||
| files | `inline::localfs` |
|
||||
| inference | `remote::fireworks`, `inline::sentence-transformers` |
|
||||
| safety | `inline::llama-guard` |
|
||||
| scoring | `inline::basic`, `inline::llm-as-judge`, `inline::braintrust` |
|
||||
|
|
|
@ -18,7 +18,7 @@ from collections.abc import Callable
|
|||
from contextlib import asynccontextmanager
|
||||
from importlib.metadata import version as parse_version
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any
|
||||
from typing import Annotated, Any, get_origin
|
||||
|
||||
import rich.pretty
|
||||
import yaml
|
||||
|
@ -244,15 +244,23 @@ def create_dynamic_typed_route(func: Any, method: str, route: str) -> Callable:
|
|||
|
||||
path_params = extract_path_params(route)
|
||||
if method == "post":
|
||||
# Annotate parameters that are in the path with Path(...) and others with Body(...)
|
||||
new_params = [new_params[0]] + [
|
||||
# Annotate parameters that are in the path with Path(...) and others with Body(...),
|
||||
# but preserve existing File() and Form() annotations for multipart form data
|
||||
new_params = (
|
||||
[new_params[0]]
|
||||
+ [
|
||||
(
|
||||
param.replace(annotation=Annotated[param.annotation, FastapiPath(..., title=param.name)])
|
||||
if param.name in path_params
|
||||
else (
|
||||
param # Keep original annotation if it's already an Annotated type
|
||||
if get_origin(param.annotation) is Annotated
|
||||
else param.replace(annotation=Annotated[param.annotation, Body(..., embed=True)])
|
||||
)
|
||||
)
|
||||
for param in new_params[1:]
|
||||
]
|
||||
)
|
||||
|
||||
route_handler.__signature__ = sig.replace(parameters=new_params)
|
||||
|
||||
|
|
20
llama_stack/providers/inline/files/localfs/__init__.py
Normal file
20
llama_stack/providers/inline/files/localfs/__init__.py
Normal file
|
@ -0,0 +1,20 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
from typing import Any
|
||||
|
||||
from llama_stack.distribution.datatypes import Api
|
||||
|
||||
from .config import LocalfsFilesImplConfig
|
||||
from .files import LocalfsFilesImpl
|
||||
|
||||
__all__ = ["LocalfsFilesImpl", "LocalfsFilesImplConfig"]
|
||||
|
||||
|
||||
async def get_provider_impl(config: LocalfsFilesImplConfig, deps: dict[Api, Any]):
|
||||
impl = LocalfsFilesImpl(config)
|
||||
await impl.initialize()
|
||||
return impl
|
31
llama_stack/providers/inline/files/localfs/config.py
Normal file
31
llama_stack/providers/inline/files/localfs/config.py
Normal file
|
@ -0,0 +1,31 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig
|
||||
|
||||
|
||||
class LocalfsFilesImplConfig(BaseModel):
|
||||
storage_dir: str = Field(
|
||||
description="Directory to store uploaded files",
|
||||
)
|
||||
metadata_store: SqlStoreConfig = Field(
|
||||
description="SQL store configuration for file metadata",
|
||||
)
|
||||
ttl_secs: int = 365 * 24 * 60 * 60 # 1 year
|
||||
|
||||
@classmethod
|
||||
def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
|
||||
return {
|
||||
"storage_dir": "${env.FILES_STORAGE_DIR:" + __distro_dir__ + "/files}",
|
||||
"metadata_store": SqliteSqlStoreConfig.sample_run_config(
|
||||
__distro_dir__=__distro_dir__,
|
||||
db_name="files_metadata.db",
|
||||
),
|
||||
}
|
214
llama_stack/providers/inline/files/localfs/files.py
Normal file
214
llama_stack/providers/inline/files/localfs/files.py
Normal file
|
@ -0,0 +1,214 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import File, Form, Response, UploadFile
|
||||
|
||||
from llama_stack.apis.common.responses import Order
|
||||
from llama_stack.apis.files import (
|
||||
Files,
|
||||
ListOpenAIFileResponse,
|
||||
OpenAIFileDeleteResponse,
|
||||
OpenAIFileObject,
|
||||
OpenAIFilePurpose,
|
||||
)
|
||||
from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
|
||||
from llama_stack.providers.utils.sqlstore.sqlstore import SqlStore, sqlstore_impl
|
||||
|
||||
from .config import LocalfsFilesImplConfig
|
||||
|
||||
|
||||
class LocalfsFilesImpl(Files):
|
||||
def __init__(self, config: LocalfsFilesImplConfig) -> None:
|
||||
self.config = config
|
||||
self.sql_store: SqlStore | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the files provider by setting up storage directory and metadata database."""
|
||||
# Create storage directory if it doesn't exist
|
||||
storage_path = Path(self.config.storage_dir)
|
||||
storage_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Initialize SQL store for metadata
|
||||
self.sql_store = sqlstore_impl(self.config.metadata_store)
|
||||
await self.sql_store.create_table(
|
||||
"openai_files",
|
||||
{
|
||||
"id": ColumnDefinition(type=ColumnType.STRING, primary_key=True),
|
||||
"filename": ColumnType.STRING,
|
||||
"purpose": ColumnType.STRING,
|
||||
"bytes": ColumnType.INTEGER,
|
||||
"created_at": ColumnType.INTEGER,
|
||||
"expires_at": ColumnType.INTEGER,
|
||||
"file_path": ColumnType.STRING, # Path to actual file on disk
|
||||
},
|
||||
)
|
||||
|
||||
def _generate_file_id(self) -> str:
|
||||
"""Generate a unique file ID for OpenAI API."""
|
||||
return f"file-{uuid.uuid4().hex}"
|
||||
|
||||
def _get_file_path(self, file_id: str) -> Path:
|
||||
"""Get the filesystem path for a file ID."""
|
||||
return Path(self.config.storage_dir) / file_id
|
||||
|
||||
# OpenAI Files API Implementation
|
||||
async def openai_upload_file(
|
||||
self,
|
||||
file: Annotated[UploadFile, File()],
|
||||
purpose: Annotated[OpenAIFilePurpose, Form()],
|
||||
) -> OpenAIFileObject:
|
||||
"""Upload a file that can be used across various endpoints."""
|
||||
if not self.sql_store:
|
||||
raise RuntimeError("Files provider not initialized")
|
||||
|
||||
file_id = self._generate_file_id()
|
||||
file_path = self._get_file_path(file_id)
|
||||
|
||||
content = await file.read()
|
||||
file_size = len(content)
|
||||
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(content)
|
||||
|
||||
created_at = int(time.time())
|
||||
expires_at = created_at + self.config.ttl_secs
|
||||
|
||||
await self.sql_store.insert(
|
||||
"openai_files",
|
||||
{
|
||||
"id": file_id,
|
||||
"filename": file.filename or "uploaded_file",
|
||||
"purpose": purpose.value,
|
||||
"bytes": file_size,
|
||||
"created_at": created_at,
|
||||
"expires_at": expires_at,
|
||||
"file_path": file_path.as_posix(),
|
||||
},
|
||||
)
|
||||
|
||||
return OpenAIFileObject(
|
||||
id=file_id,
|
||||
filename=file.filename or "uploaded_file",
|
||||
purpose=purpose,
|
||||
bytes=file_size,
|
||||
created_at=created_at,
|
||||
expires_at=expires_at,
|
||||
)
|
||||
|
||||
async def openai_list_files(
|
||||
self,
|
||||
after: str | None = None,
|
||||
limit: int | None = 10000,
|
||||
order: Order | None = Order.desc,
|
||||
purpose: OpenAIFilePurpose | None = None,
|
||||
) -> ListOpenAIFileResponse:
|
||||
"""Returns a list of files that belong to the user's organization."""
|
||||
if not self.sql_store:
|
||||
raise RuntimeError("Files provider not initialized")
|
||||
|
||||
# TODO: Implement 'after' pagination properly
|
||||
if after:
|
||||
raise NotImplementedError("After pagination not yet implemented")
|
||||
|
||||
where = None
|
||||
if purpose:
|
||||
where = {"purpose": purpose.value}
|
||||
|
||||
rows = await self.sql_store.fetch_all(
|
||||
"openai_files",
|
||||
where=where,
|
||||
order_by=[("created_at", order.value if order else Order.desc.value)],
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
files = [
|
||||
OpenAIFileObject(
|
||||
id=row["id"],
|
||||
filename=row["filename"],
|
||||
purpose=OpenAIFilePurpose(row["purpose"]),
|
||||
bytes=row["bytes"],
|
||||
created_at=row["created_at"],
|
||||
expires_at=row["expires_at"],
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
return ListOpenAIFileResponse(
|
||||
data=files,
|
||||
has_more=False, # TODO: Implement proper pagination
|
||||
first_id=files[0].id if files else "",
|
||||
last_id=files[-1].id if files else "",
|
||||
)
|
||||
|
||||
async def openai_retrieve_file(self, file_id: str) -> OpenAIFileObject:
|
||||
"""Returns information about a specific file."""
|
||||
if not self.sql_store:
|
||||
raise RuntimeError("Files provider not initialized")
|
||||
|
||||
row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
|
||||
if not row:
|
||||
raise ValueError(f"File with id {file_id} not found")
|
||||
|
||||
return OpenAIFileObject(
|
||||
id=row["id"],
|
||||
filename=row["filename"],
|
||||
purpose=OpenAIFilePurpose(row["purpose"]),
|
||||
bytes=row["bytes"],
|
||||
created_at=row["created_at"],
|
||||
expires_at=row["expires_at"],
|
||||
)
|
||||
|
||||
async def openai_delete_file(self, file_id: str) -> OpenAIFileDeleteResponse:
|
||||
"""Delete a file."""
|
||||
if not self.sql_store:
|
||||
raise RuntimeError("Files provider not initialized")
|
||||
|
||||
row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
|
||||
if not row:
|
||||
raise ValueError(f"File with id {file_id} not found")
|
||||
|
||||
# Delete physical file
|
||||
file_path = Path(row["file_path"])
|
||||
if file_path.exists():
|
||||
file_path.unlink()
|
||||
|
||||
# Delete metadata from database
|
||||
await self.sql_store.delete("openai_files", where={"id": file_id})
|
||||
|
||||
return OpenAIFileDeleteResponse(
|
||||
id=file_id,
|
||||
deleted=True,
|
||||
)
|
||||
|
||||
async def openai_retrieve_file_content(self, file_id: str) -> Response:
|
||||
"""Returns the contents of the specified file."""
|
||||
if not self.sql_store:
|
||||
raise RuntimeError("Files provider not initialized")
|
||||
|
||||
# Get file metadata
|
||||
row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
|
||||
if not row:
|
||||
raise ValueError(f"File with id {file_id} not found")
|
||||
|
||||
# Read file content
|
||||
file_path = Path(row["file_path"])
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File content not found on disk: {file_path}")
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
|
||||
# Return as binary response with appropriate content type
|
||||
return Response(
|
||||
content=content,
|
||||
media_type="application/octet-stream",
|
||||
headers={"Content-Disposition": f'attachment; filename="{row["filename"]}"'},
|
||||
)
|
|
@ -4,8 +4,22 @@
|
|||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
from llama_stack.providers.datatypes import ProviderSpec
|
||||
from llama_stack.providers.datatypes import (
|
||||
Api,
|
||||
InlineProviderSpec,
|
||||
ProviderSpec,
|
||||
)
|
||||
from llama_stack.providers.utils.sqlstore.sqlstore import sql_store_pip_packages
|
||||
|
||||
|
||||
def available_providers() -> list[ProviderSpec]:
|
||||
return []
|
||||
return [
|
||||
InlineProviderSpec(
|
||||
api=Api.files,
|
||||
provider_type="inline::localfs",
|
||||
# TODO: make this dynamic according to the sql store type
|
||||
pip_packages=sql_store_pip_packages,
|
||||
module="llama_stack.providers.inline.files.localfs",
|
||||
config_class="llama_stack.providers.inline.files.localfs.config.LocalfsFilesImplConfig",
|
||||
),
|
||||
]
|
||||
|
|
|
@ -16,6 +16,8 @@ from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
|
|||
|
||||
from .api import SqlStore
|
||||
|
||||
sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"]
|
||||
|
||||
|
||||
class SqlStoreType(Enum):
|
||||
sqlite = "sqlite"
|
||||
|
|
|
@ -24,6 +24,8 @@ distribution_spec:
|
|||
- inline::basic
|
||||
- inline::llm-as-judge
|
||||
- inline::braintrust
|
||||
files:
|
||||
- inline::localfs
|
||||
tool_runtime:
|
||||
- remote::brave-search
|
||||
- remote::tavily-search
|
||||
|
|
|
@ -13,6 +13,7 @@ from llama_stack.distribution.datatypes import (
|
|||
ShieldInput,
|
||||
ToolGroupInput,
|
||||
)
|
||||
from llama_stack.providers.inline.files.localfs.config import LocalfsFilesImplConfig
|
||||
from llama_stack.providers.inline.inference.sentence_transformers import (
|
||||
SentenceTransformersInferenceConfig,
|
||||
)
|
||||
|
@ -36,6 +37,7 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
"eval": ["inline::meta-reference"],
|
||||
"datasetio": ["remote::huggingface", "inline::localfs"],
|
||||
"scoring": ["inline::basic", "inline::llm-as-judge", "inline::braintrust"],
|
||||
"files": ["inline::localfs"],
|
||||
"tool_runtime": [
|
||||
"remote::brave-search",
|
||||
"remote::tavily-search",
|
||||
|
@ -62,6 +64,11 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
provider_type="inline::faiss",
|
||||
config=FaissVectorIOConfig.sample_run_config(f"~/.llama/distributions/{name}"),
|
||||
)
|
||||
files_provider = Provider(
|
||||
provider_id="meta-reference-files",
|
||||
provider_type="inline::localfs",
|
||||
config=LocalfsFilesImplConfig.sample_run_config(f"~/.llama/distributions/{name}"),
|
||||
)
|
||||
|
||||
available_models = {
|
||||
"fireworks": MODEL_ENTRIES,
|
||||
|
@ -104,6 +111,7 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
provider_overrides={
|
||||
"inference": [inference_provider, embedding_provider],
|
||||
"vector_io": [vector_io_provider],
|
||||
"files": [files_provider],
|
||||
},
|
||||
default_models=default_models + [embedding_model],
|
||||
default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")],
|
||||
|
@ -116,6 +124,7 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
embedding_provider,
|
||||
],
|
||||
"vector_io": [vector_io_provider],
|
||||
"files": [files_provider],
|
||||
"safety": [
|
||||
Provider(
|
||||
provider_id="llama-guard",
|
||||
|
|
|
@ -4,6 +4,7 @@ apis:
|
|||
- agents
|
||||
- datasetio
|
||||
- eval
|
||||
- files
|
||||
- inference
|
||||
- safety
|
||||
- scoring
|
||||
|
@ -90,6 +91,14 @@ providers:
|
|||
provider_type: inline::braintrust
|
||||
config:
|
||||
openai_api_key: ${env.OPENAI_API_KEY:}
|
||||
files:
|
||||
- provider_id: meta-reference-files
|
||||
provider_type: inline::localfs
|
||||
config:
|
||||
storage_dir: ${env.FILES_STORAGE_DIR:~/.llama/distributions/fireworks/files}
|
||||
metadata_store:
|
||||
type: sqlite
|
||||
db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/fireworks}/files_metadata.db
|
||||
tool_runtime:
|
||||
- provider_id: brave-search
|
||||
provider_type: remote::brave-search
|
||||
|
|
|
@ -4,6 +4,7 @@ apis:
|
|||
- agents
|
||||
- datasetio
|
||||
- eval
|
||||
- files
|
||||
- inference
|
||||
- safety
|
||||
- scoring
|
||||
|
@ -85,6 +86,14 @@ providers:
|
|||
provider_type: inline::braintrust
|
||||
config:
|
||||
openai_api_key: ${env.OPENAI_API_KEY:}
|
||||
files:
|
||||
- provider_id: meta-reference-files
|
||||
provider_type: inline::localfs
|
||||
config:
|
||||
storage_dir: ${env.FILES_STORAGE_DIR:~/.llama/distributions/fireworks/files}
|
||||
metadata_store:
|
||||
type: sqlite
|
||||
db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/fireworks}/files_metadata.db
|
||||
tool_runtime:
|
||||
- provider_id: brave-search
|
||||
provider_type: remote::brave-search
|
||||
|
|
|
@ -41,6 +41,7 @@ dependencies = [
|
|||
"tiktoken",
|
||||
"pillow",
|
||||
"h11>=0.16.0",
|
||||
"python-multipart>=0.0.20",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
|
|
@ -130,6 +130,8 @@ python-dotenv==1.0.1
|
|||
# via llama-stack
|
||||
python-jose==3.4.0
|
||||
# via llama-stack
|
||||
python-multipart==0.0.20
|
||||
# via llama-stack
|
||||
pytz==2025.1
|
||||
# via pandas
|
||||
pyyaml==6.0.2
|
||||
|
|
5
tests/integration/files/__init__.py
Normal file
5
tests/integration/files/__init__.py
Normal file
|
@ -0,0 +1,5 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
51
tests/integration/files/test_files.py
Normal file
51
tests/integration/files/test_files.py
Normal file
|
@ -0,0 +1,51 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
|
||||
def test_openai_client_basic_operations(openai_client):
|
||||
"""Test basic file operations through OpenAI client."""
|
||||
client = openai_client
|
||||
|
||||
test_content = b"files test content"
|
||||
|
||||
try:
|
||||
# Upload file using OpenAI client
|
||||
with BytesIO(test_content) as file_buffer:
|
||||
file_buffer.name = "openai_test.txt"
|
||||
uploaded_file = client.files.create(file=file_buffer, purpose="assistants")
|
||||
|
||||
# Verify basic response structure
|
||||
assert uploaded_file.id.startswith("file-")
|
||||
assert hasattr(uploaded_file, "filename")
|
||||
|
||||
# List files
|
||||
files_list = client.files.list()
|
||||
file_ids = [f.id for f in files_list.data]
|
||||
assert uploaded_file.id in file_ids
|
||||
|
||||
# Retrieve file info
|
||||
retrieved_file = client.files.retrieve(uploaded_file.id)
|
||||
assert retrieved_file.id == uploaded_file.id
|
||||
|
||||
# Retrieve file content - OpenAI client returns httpx Response object
|
||||
content_response = client.files.content(uploaded_file.id)
|
||||
# The response is an httpx Response object with .content attribute containing bytes
|
||||
content = content_response.content
|
||||
assert content == test_content
|
||||
|
||||
# Delete file
|
||||
delete_response = client.files.delete(uploaded_file.id)
|
||||
assert delete_response.deleted is True
|
||||
|
||||
except Exception as e:
|
||||
# Cleanup in case of failure
|
||||
try:
|
||||
client.files.delete(uploaded_file.id)
|
||||
except Exception:
|
||||
pass
|
||||
raise e
|
5
tests/unit/files/__init__.py
Normal file
5
tests/unit/files/__init__.py
Normal file
|
@ -0,0 +1,5 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
334
tests/unit/files/test_files.py
Normal file
334
tests/unit/files/test_files.py
Normal file
|
@ -0,0 +1,334 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from llama_stack.apis.common.responses import Order
|
||||
from llama_stack.apis.files import OpenAIFilePurpose
|
||||
from llama_stack.providers.inline.files.localfs import (
|
||||
LocalfsFilesImpl,
|
||||
LocalfsFilesImplConfig,
|
||||
)
|
||||
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
|
||||
|
||||
|
||||
class MockUploadFile:
|
||||
"""Mock UploadFile for testing file uploads."""
|
||||
|
||||
def __init__(self, content: bytes, filename: str, content_type: str = "text/plain"):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def files_provider(tmp_path):
|
||||
"""Create a files provider with temporary storage for testing."""
|
||||
storage_dir = tmp_path / "files"
|
||||
db_path = tmp_path / "files_metadata.db"
|
||||
|
||||
config = LocalfsFilesImplConfig(
|
||||
storage_dir=storage_dir.as_posix(), metadata_store=SqliteSqlStoreConfig(db_path=db_path.as_posix())
|
||||
)
|
||||
|
||||
provider = LocalfsFilesImpl(config)
|
||||
await provider.initialize()
|
||||
yield provider
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_text_file():
|
||||
"""Sample text file for testing."""
|
||||
content = b"Hello, this is a test file for the OpenAI Files API!"
|
||||
return MockUploadFile(content, "test.txt", "text/plain")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_json_file():
|
||||
"""Sample JSON file for testing."""
|
||||
content = b'{"message": "Hello, World!", "type": "test"}'
|
||||
return MockUploadFile(content, "data.json", "application/json")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_file():
|
||||
"""Large file for testing file size handling."""
|
||||
content = b"x" * 1024 * 1024 # 1MB file
|
||||
return MockUploadFile(content, "large_file.bin", "application/octet-stream")
|
||||
|
||||
|
||||
class TestOpenAIFilesAPI:
|
||||
"""Test suite for OpenAI Files API endpoints."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_success(self, files_provider, sample_text_file):
|
||||
"""Test successful file upload."""
|
||||
# Upload file
|
||||
result = await files_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
|
||||
# Verify response
|
||||
assert result.id.startswith("file-")
|
||||
assert result.filename == "test.txt"
|
||||
assert result.purpose == OpenAIFilePurpose.ASSISTANTS
|
||||
assert result.bytes == len(sample_text_file.content)
|
||||
assert result.created_at > 0
|
||||
assert result.expires_at > result.created_at
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_different_purposes(self, files_provider, sample_text_file):
|
||||
"""Test uploading files with different purposes."""
|
||||
purposes = list(OpenAIFilePurpose)
|
||||
|
||||
uploaded_files = []
|
||||
for purpose in purposes:
|
||||
result = await files_provider.openai_upload_file(file=sample_text_file, purpose=purpose)
|
||||
uploaded_files.append(result)
|
||||
assert result.purpose == purpose
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_different_file_types(self, files_provider, sample_text_file, sample_json_file, large_file):
|
||||
"""Test uploading different types and sizes of files."""
|
||||
files_to_test = [
|
||||
(sample_text_file, "test.txt"),
|
||||
(sample_json_file, "data.json"),
|
||||
(large_file, "large_file.bin"),
|
||||
]
|
||||
|
||||
for file_obj, expected_filename in files_to_test:
|
||||
result = await files_provider.openai_upload_file(file=file_obj, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
assert result.filename == expected_filename
|
||||
assert result.bytes == len(file_obj.content)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_files_empty(self, files_provider):
|
||||
"""Test listing files when no files exist."""
|
||||
result = await files_provider.openai_list_files()
|
||||
|
||||
assert result.data == []
|
||||
assert result.has_more is False
|
||||
assert result.first_id == ""
|
||||
assert result.last_id == ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_files_with_content(self, files_provider, sample_text_file, sample_json_file):
|
||||
"""Test listing files when files exist."""
|
||||
# Upload multiple files
|
||||
file1 = await files_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
file2 = await files_provider.openai_upload_file(file=sample_json_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
|
||||
# List files
|
||||
result = await files_provider.openai_list_files()
|
||||
|
||||
assert len(result.data) == 2
|
||||
file_ids = [f.id for f in result.data]
|
||||
assert file1.id in file_ids
|
||||
assert file2.id in file_ids
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_files_with_purpose_filter(self, files_provider, sample_text_file):
|
||||
"""Test listing files with purpose filtering."""
|
||||
# Upload file with specific purpose
|
||||
uploaded_file = await files_provider.openai_upload_file(
|
||||
file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS
|
||||
)
|
||||
|
||||
# List files with matching purpose
|
||||
result = await files_provider.openai_list_files(purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0].id == uploaded_file.id
|
||||
assert result.data[0].purpose == OpenAIFilePurpose.ASSISTANTS
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_files_with_limit(self, files_provider, sample_text_file):
|
||||
"""Test listing files with limit parameter."""
|
||||
# Upload multiple files
|
||||
for _ in range(5):
|
||||
await files_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
|
||||
# List with limit
|
||||
result = await files_provider.openai_list_files(limit=3)
|
||||
assert len(result.data) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_files_with_order(self, files_provider, sample_text_file):
|
||||
"""Test listing files with different order."""
|
||||
# Upload multiple files
|
||||
files = []
|
||||
for _ in range(3):
|
||||
file = await files_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
files.append(file)
|
||||
|
||||
# Test descending order (default)
|
||||
result_desc = await files_provider.openai_list_files(order=Order.desc)
|
||||
assert len(result_desc.data) == 3
|
||||
# Most recent should be first
|
||||
assert result_desc.data[0].created_at >= result_desc.data[1].created_at >= result_desc.data[2].created_at
|
||||
|
||||
# Test ascending order
|
||||
result_asc = await files_provider.openai_list_files(order=Order.asc)
|
||||
assert len(result_asc.data) == 3
|
||||
# Oldest should be first
|
||||
assert result_asc.data[0].created_at <= result_asc.data[1].created_at <= result_asc.data[2].created_at
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_file_success(self, files_provider, sample_text_file):
|
||||
"""Test successful file retrieval."""
|
||||
# Upload file
|
||||
uploaded_file = await files_provider.openai_upload_file(
|
||||
file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS
|
||||
)
|
||||
|
||||
# Retrieve file
|
||||
retrieved_file = await files_provider.openai_retrieve_file(uploaded_file.id)
|
||||
|
||||
# Verify response
|
||||
assert retrieved_file.id == uploaded_file.id
|
||||
assert retrieved_file.filename == uploaded_file.filename
|
||||
assert retrieved_file.purpose == uploaded_file.purpose
|
||||
assert retrieved_file.bytes == uploaded_file.bytes
|
||||
assert retrieved_file.created_at == uploaded_file.created_at
|
||||
assert retrieved_file.expires_at == uploaded_file.expires_at
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_file_not_found(self, files_provider):
|
||||
"""Test retrieving a non-existent file."""
|
||||
with pytest.raises(ValueError, match="File with id file-nonexistent not found"):
|
||||
await files_provider.openai_retrieve_file("file-nonexistent")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_file_content_success(self, files_provider, sample_text_file):
|
||||
"""Test successful file content retrieval."""
|
||||
# Upload file
|
||||
uploaded_file = await files_provider.openai_upload_file(
|
||||
file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS
|
||||
)
|
||||
|
||||
# Retrieve file content
|
||||
content = await files_provider.openai_retrieve_file_content(uploaded_file.id)
|
||||
|
||||
# Verify content
|
||||
assert content.body == sample_text_file.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_file_content_not_found(self, files_provider):
|
||||
"""Test retrieving content of a non-existent file."""
|
||||
with pytest.raises(ValueError, match="File with id file-nonexistent not found"):
|
||||
await files_provider.openai_retrieve_file_content("file-nonexistent")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_file_success(self, files_provider, sample_text_file):
|
||||
"""Test successful file deletion."""
|
||||
# Upload file
|
||||
uploaded_file = await files_provider.openai_upload_file(
|
||||
file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS
|
||||
)
|
||||
|
||||
# Verify file exists
|
||||
await files_provider.openai_retrieve_file(uploaded_file.id)
|
||||
|
||||
# Delete file
|
||||
delete_response = await files_provider.openai_delete_file(uploaded_file.id)
|
||||
|
||||
# Verify delete response
|
||||
assert delete_response.id == uploaded_file.id
|
||||
assert delete_response.deleted is True
|
||||
|
||||
# Verify file no longer exists
|
||||
with pytest.raises(ValueError, match=f"File with id {uploaded_file.id} not found"):
|
||||
await files_provider.openai_retrieve_file(uploaded_file.id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_file_not_found(self, files_provider):
|
||||
"""Test deleting a non-existent file."""
|
||||
with pytest.raises(ValueError, match="File with id file-nonexistent not found"):
|
||||
await files_provider.openai_delete_file("file-nonexistent")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_persistence_across_operations(self, files_provider, sample_text_file):
|
||||
"""Test that files persist correctly across multiple operations."""
|
||||
# Upload file
|
||||
uploaded_file = await files_provider.openai_upload_file(
|
||||
file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS
|
||||
)
|
||||
|
||||
# Verify it appears in listing
|
||||
files_list = await files_provider.openai_list_files()
|
||||
assert len(files_list.data) == 1
|
||||
assert files_list.data[0].id == uploaded_file.id
|
||||
|
||||
# Retrieve file info
|
||||
retrieved_file = await files_provider.openai_retrieve_file(uploaded_file.id)
|
||||
assert retrieved_file.id == uploaded_file.id
|
||||
|
||||
# Retrieve file content
|
||||
content = await files_provider.openai_retrieve_file_content(uploaded_file.id)
|
||||
assert content.body == sample_text_file.content
|
||||
|
||||
# Delete file
|
||||
await files_provider.openai_delete_file(uploaded_file.id)
|
||||
|
||||
# Verify it's gone from listing
|
||||
files_list = await files_provider.openai_list_files()
|
||||
assert len(files_list.data) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_files_operations(self, files_provider, sample_text_file, sample_json_file):
|
||||
"""Test operations with multiple files."""
|
||||
# Upload multiple files
|
||||
file1 = await files_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
file2 = await files_provider.openai_upload_file(file=sample_json_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||
|
||||
# Verify both exist
|
||||
files_list = await files_provider.openai_list_files()
|
||||
assert len(files_list.data) == 2
|
||||
|
||||
# Delete one file
|
||||
await files_provider.openai_delete_file(file1.id)
|
||||
|
||||
# Verify only one remains
|
||||
files_list = await files_provider.openai_list_files()
|
||||
assert len(files_list.data) == 1
|
||||
assert files_list.data[0].id == file2.id
|
||||
|
||||
# Verify the remaining file is still accessible
|
||||
content = await files_provider.openai_retrieve_file_content(file2.id)
|
||||
assert content.body == sample_json_file.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_id_uniqueness(self, files_provider, sample_text_file):
|
||||
"""Test that each uploaded file gets a unique ID."""
|
||||
file_ids = set()
|
||||
|
||||
# Upload same file multiple times
|
||||
for _ in range(10):
|
||||
uploaded_file = await files_provider.openai_upload_file(
|
||||
file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS
|
||||
)
|
||||
assert uploaded_file.id not in file_ids, f"Duplicate file ID: {uploaded_file.id}"
|
||||
file_ids.add(uploaded_file.id)
|
||||
assert uploaded_file.id.startswith("file-")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_no_filename_handling(self, files_provider):
|
||||
"""Test handling files with no filename."""
|
||||
file_without_name = MockUploadFile(b"content", None) # No filename
|
||||
|
||||
uploaded_file = await files_provider.openai_upload_file(
|
||||
file=file_without_name, purpose=OpenAIFilePurpose.ASSISTANTS
|
||||
)
|
||||
|
||||
assert uploaded_file.filename == "uploaded_file" # Default filename
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_after_pagination_not_implemented(self, files_provider):
|
||||
"""Test that 'after' pagination raises NotImplementedError."""
|
||||
with pytest.raises(NotImplementedError, match="After pagination not yet implemented"):
|
||||
await files_provider.openai_list_files(after="file-some-id")
|
Loading…
Add table
Add a link
Reference in a new issue