llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/inline/files/localfs/files.py

@@ -0,0 +1,219 @@
+# 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 Depends, File, Form, Response, UploadFile
+
+from llama_stack.apis.common.errors import ResourceNotFoundError
+from llama_stack.apis.common.responses import Order
+from llama_stack.apis.files import (
+    ExpiresAfter,
+    Files,
+    ListOpenAIFileResponse,
+    OpenAIFileDeleteResponse,
+    OpenAIFileObject,
+    OpenAIFilePurpose,
+)
+from llama_stack.core.datatypes import AccessRule
+from llama_stack.core.id_generation import generate_object_id
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.files.form_data import parse_expires_after
+from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
+from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
+
+from .config import LocalfsFilesImplConfig
+
+logger = get_logger(name=__name__, category="files")
+
+
+class LocalfsFilesImpl(Files):
+    def __init__(self, config: LocalfsFilesImplConfig, policy: list[AccessRule]) -> None:
+        self.config = config
+        self.policy = policy
+        self.sql_store: AuthorizedSqlStore | 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 = AuthorizedSqlStore(sqlstore_impl(self.config.metadata_store), self.policy)
+        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
+            },
+        )
+
+    async def shutdown(self) -> None:
+        pass
+
+    def _generate_file_id(self) -> str:
+        """Generate a unique file ID for OpenAI API."""
+        return generate_object_id("file", lambda: 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
+
+    async def _lookup_file_id(self, file_id: str) -> tuple[OpenAIFileObject, Path]:
+        """Look up a OpenAIFileObject and filesystem path from its ID."""
+        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 ResourceNotFoundError(file_id, "File", "client.files.list()")
+
+        file_path = Path(row.pop("file_path"))
+        return OpenAIFileObject(**row), file_path
+
+    # OpenAI Files API Implementation
+    async def openai_upload_file(
+        self,
+        file: Annotated[UploadFile, File()],
+        purpose: Annotated[OpenAIFilePurpose, Form()],
+        expires_after: Annotated[ExpiresAfter | None, Depends(parse_expires_after)] = None,
+    ) -> OpenAIFileObject:
+        """Upload a file that can be used across various endpoints."""
+        if not self.sql_store:
+            raise RuntimeError("Files provider not initialized")
+
+        if expires_after is not None:
+            logger.warning(
+                f"File expiration is not supported by this provider, ignoring expires_after: {expires_after}"
+            )
+
+        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")
+
+        if not order:
+            order = Order.desc
+
+        where_conditions = {}
+        if purpose:
+            where_conditions["purpose"] = purpose.value
+
+        paginated_result = await self.sql_store.fetch_all(
+            table="openai_files",
+            where=where_conditions if where_conditions else None,
+            order_by=[("created_at", order.value)],
+            cursor=("id", after) if after else None,
+            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 paginated_result.data
+        ]
+
+        return ListOpenAIFileResponse(
+            data=files,
+            has_more=paginated_result.has_more,
+            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."""
+        file_obj, _ = await self._lookup_file_id(file_id)
+
+        return file_obj
+
+    async def openai_delete_file(self, file_id: str) -> OpenAIFileDeleteResponse:
+        """Delete a file."""
+        # Delete physical file
+        _, file_path = await self._lookup_file_id(file_id)
+        if file_path.exists():
+            file_path.unlink()
+
+        # Delete metadata from database
+        assert self.sql_store is not None, "Files provider not initialized"
+        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."""
+        # Read file content
+        file_obj, file_path = await self._lookup_file_id(file_id)
+
+        if not file_path.exists():
+            logger.warning(f"File '{file_id}'s underlying '{file_path}' is missing, deleting metadata.")
+            await self.openai_delete_file(file_id)
+            raise ResourceNotFoundError(file_id, "File", "client.files.list()")
+
+        # Return as binary response with appropriate content type
+        return Response(
+            content=file_path.read_bytes(),
+            media_type="application/octet-stream",
+            headers={"Content-Disposition": f'attachment; filename="{file_obj.filename}"'},
+        )

llama_stack/providers/inline/inference/__init__.py

@@ -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.

llama_stack/providers/{impls/meta_reference/inference → inline/inference/meta_reference}/__init__.py

@@ -4,14 +4,14 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Union
+from typing import Any
 
-from .config import MetaReferenceInferenceConfig, MetaReferenceQuantizedInferenceConfig
+from .config import MetaReferenceInferenceConfig
 
 
 async def get_provider_impl(
-    config: Union[MetaReferenceInferenceConfig, MetaReferenceQuantizedInferenceConfig],
-    _deps,
+    config: MetaReferenceInferenceConfig,
+    _deps: dict[str, Any],
 ):
     from .inference import MetaReferenceInferenceImpl
 

llama_stack/providers/inline/inference/meta_reference/common.py

@@ -0,0 +1,24 @@
+# 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 pathlib import Path
+
+from llama_stack.core.utils.model_utils import model_local_dir
+
+
+def model_checkpoint_dir(model_id) -> str:
+    checkpoint_dir = Path(model_local_dir(model_id))
+
+    paths = [Path(checkpoint_dir / f"consolidated.{ext}") for ext in ["pth", "00.pth"]]
+    if not any(p.exists() for p in paths):
+        checkpoint_dir = checkpoint_dir / "original"
+
+    assert checkpoint_dir.exists(), (
+        f"Could not find checkpoints in: {model_local_dir(model_id)}. "
+        f"If you try to use the native llama model, please download the model using `llama-model download --source meta --model-id {model_id}` (see https://github.com/meta-llama/llama-models). "
+        f"Otherwise, please save your model checkpoint under {model_local_dir(model_id)}"
+    )
+    return str(checkpoint_dir)

llama_stack/providers/inline/inference/meta_reference/config.py

@@ -0,0 +1,68 @@
+# 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_validator
+
+from llama_stack.apis.inference import QuantizationConfig
+from llama_stack.providers.utils.inference import supported_inference_models
+
+
+class MetaReferenceInferenceConfig(BaseModel):
+    # this is a placeholder to indicate inference model id
+    # the actual inference model id is dtermined by the moddel id in the request
+    # Note: you need to register the model before using it for inference
+    # models in the resouce list in the run.yaml config will be registered automatically
+    model: str | None = None
+    torch_seed: int | None = None
+    max_seq_len: int = 4096
+    max_batch_size: int = 1
+    model_parallel_size: int | None = None
+
+    # when this is False, we assume that the distributed process group is setup by someone
+    # outside of this code (e.g., when run inside `torchrun`). that is useful for clients
+    # (including our testing code) who might be using llama-stack as a library.
+    create_distributed_process_group: bool = True
+
+    # By default, the implementation will look at ~/.llama/checkpoints/<model> but you
+    # can override by specifying the directory explicitly
+    checkpoint_dir: str | None = None
+
+    quantization: QuantizationConfig | None = None
+
+    @field_validator("model")
+    @classmethod
+    def validate_model(cls, model: str) -> str:
+        permitted_models = supported_inference_models()
+        descriptors = [m.descriptor() for m in permitted_models]
+        repos = [m.huggingface_repo for m in permitted_models if m.huggingface_repo is not None]
+        if model not in (descriptors + repos):
+            model_list = "\n\t".join(repos)
+            raise ValueError(f"Unknown model: `{model}`. Choose from [\n\t{model_list}\n]")
+        return model
+
+    @classmethod
+    def sample_run_config(
+        cls,
+        model: str = "Llama3.2-3B-Instruct",
+        checkpoint_dir: str = "${env.CHECKPOINT_DIR:=null}",
+        quantization_type: str = "${env.QUANTIZATION_TYPE:=bf16}",
+        model_parallel_size: str = "${env.MODEL_PARALLEL_SIZE:=0}",
+        max_batch_size: str = "${env.MAX_BATCH_SIZE:=1}",
+        max_seq_len: str = "${env.MAX_SEQ_LEN:=4096}",
+        **kwargs,
+    ) -> dict[str, Any]:
+        return {
+            "model": model,
+            "checkpoint_dir": checkpoint_dir,
+            "quantization": {
+                "type": quantization_type,
+            },
+            "model_parallel_size": model_parallel_size,
+            "max_batch_size": max_batch_size,
+            "max_seq_len": max_seq_len,
+        }

llama_stack/providers/inline/inference/meta_reference/generators.py

@@ -0,0 +1,211 @@
+# 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 math
+from collections.abc import Generator
+from typing import Optional
+
+import torch
+from lmformatenforcer import JsonSchemaParser, TokenEnforcer, TokenEnforcerTokenizerData
+
+from llama_stack.apis.inference import (
+    GreedySamplingStrategy,
+    JsonSchemaResponseFormat,
+    ResponseFormat,
+    SamplingParams,
+    TopPSamplingStrategy,
+)
+from llama_stack.models.llama.datatypes import QuantizationMode
+from llama_stack.models.llama.llama3.generation import Llama3
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer as Llama3Tokenizer
+from llama_stack.models.llama.llama4.generation import Llama4
+from llama_stack.models.llama.llama4.tokenizer import Tokenizer as Llama4Tokenizer
+from llama_stack.models.llama.sku_types import Model, ModelFamily
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    ChatCompletionRequestWithRawContent,
+    CompletionRequestWithRawContent,
+    get_default_tool_prompt_format,
+)
+
+from .common import model_checkpoint_dir
+from .config import MetaReferenceInferenceConfig
+from .inference import resolve_model
+
+Tokenizer = Llama4Tokenizer | Llama3Tokenizer
+
+
+class LogitsProcessor:
+    def __init__(self, token_enforcer: TokenEnforcer):
+        self.token_enforcer = token_enforcer
+        self.mask: torch.Tensor | None = None
+
+    def __call__(self, tokens: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
+        token_sequence = tokens[0, :].tolist()
+        allowed_tokens = self.token_enforcer.get_allowed_tokens(token_sequence)
+
+        if self.mask is not None:
+            self.mask.fill_(-math.inf)
+        else:
+            self.mask = torch.full_like(scores, -math.inf)
+
+        self.mask[:, :, allowed_tokens] = 0
+        scores = scores + self.mask
+        return scores
+
+
+def get_logits_processor(
+    tokenizer: Tokenizer,
+    vocab_size: int,
+    response_format: ResponseFormat | None,
+) -> Optional["LogitsProcessor"]:
+    if response_format is None:
+        return None
+
+    if not isinstance(response_format, JsonSchemaResponseFormat):
+        raise ValueError(f"Unsupported response format type {response_format.type}")
+
+    parser = JsonSchemaParser(response_format.json_schema)
+    data = TokenEnforcerTokenizerData(
+        _build_regular_tokens_list(tokenizer, vocab_size),
+        tokenizer.decode,
+        tokenizer.stop_tokens,
+    )
+    token_enforcer = TokenEnforcer(data, parser)
+    return LogitsProcessor(token_enforcer)
+
+
+def _build_regular_tokens_list(tokenizer: Tokenizer, vocab_size: int) -> list[tuple[int, str, bool]]:
+    token_0 = tokenizer.encode("0", bos=False, eos=False)[-1]
+    regular_tokens = []
+
+    special_token_ids = set(tokenizer.special_tokens.values())
+    for token_idx in range(vocab_size):
+        if token_idx in special_token_ids:
+            continue
+
+        # We prepend token 0 and skip the first letter of the result to get a space if the token is a start word.
+        decoded_after_0 = tokenizer.decode([token_0, token_idx])[1:]
+        decoded_regular = tokenizer.decode([token_idx])
+        is_word_start_token = len(decoded_after_0) > len(decoded_regular)
+        regular_tokens.append((token_idx, decoded_after_0, is_word_start_token))
+    return regular_tokens
+
+
+def _infer_sampling_params(sampling_params: SamplingParams):
+    if isinstance(sampling_params.strategy, GreedySamplingStrategy):
+        temperature = 0.0
+        top_p = 1.0
+    elif isinstance(sampling_params.strategy, TopPSamplingStrategy):
+        temperature = sampling_params.strategy.temperature or 1.0
+        top_p = sampling_params.strategy.top_p or 1.0
+    else:
+        raise ValueError(f"Unsupported sampling strategy {sampling_params.strategy}")
+    return temperature, top_p
+
+
+def _infer_tool_prompt_format(request: ChatCompletionRequestWithRawContent):
+    tool_config = request.tool_config
+    if tool_config is not None and tool_config.tool_prompt_format is not None:
+        return tool_config.tool_prompt_format
+    else:
+        return get_default_tool_prompt_format(request.model)
+
+
+class LlamaGenerator:
+    def __init__(
+        self,
+        config: MetaReferenceInferenceConfig,
+        model_id: str,
+        llama_model: Model,
+    ):
+        if config.checkpoint_dir and config.checkpoint_dir != "null":
+            ckpt_dir = config.checkpoint_dir
+        else:
+            resolved_model = resolve_model(model_id)
+            if resolved_model is None:
+                # if the model is not a native llama model, get the default checkpoint_dir based on model id
+                ckpt_dir = model_checkpoint_dir(model_id)
+            else:
+                # if the model is a native llama model, get the default checkpoint_dir based on model core_model_id value
+                ckpt_dir = model_checkpoint_dir(resolved_model.descriptor())
+
+        if config.quantization:
+            if config.quantization.type == "fp8_mixed":
+                quantization_mode = QuantizationMode.fp8_mixed
+            elif config.quantization.type == "int4_mixed":
+                quantization_mode = QuantizationMode.int4_mixed
+            elif config.quantization.type == "bf16":
+                quantization_mode = None
+            else:
+                raise ValueError(f"Unsupported quantization mode {config.quantization}")
+        else:
+            quantization_mode = None
+
+        cls = Llama4 if llama_model.model_family == ModelFamily.llama4 else Llama3
+        self.inner_generator = cls.build(
+            ckpt_dir=ckpt_dir,
+            max_seq_len=config.max_seq_len,
+            max_batch_size=config.max_batch_size,
+            world_size=config.model_parallel_size or llama_model.pth_file_count,
+            quantization_mode=quantization_mode,
+        )
+
+        self.tokenizer = self.inner_generator.tokenizer
+        self.args = self.inner_generator.args
+        self.formatter = self.inner_generator.formatter
+
+    def completion(
+        self,
+        request_batch: list[CompletionRequestWithRawContent],
+    ) -> Generator:
+        first_request = request_batch[0]
+        sampling_params = first_request.sampling_params or SamplingParams()
+        max_gen_len = sampling_params.max_tokens
+        if max_gen_len is None or max_gen_len == 0 or max_gen_len >= self.args.max_seq_len:
+            max_gen_len = self.args.max_seq_len - 1
+
+        temperature, top_p = _infer_sampling_params(sampling_params)
+        yield from self.inner_generator.generate(
+            llm_inputs=[self.formatter.encode_content(request.content) for request in request_batch],
+            max_gen_len=max_gen_len,
+            temperature=temperature,
+            top_p=top_p,
+            logprobs=bool(first_request.logprobs),
+            echo=False,
+            logits_processor=get_logits_processor(
+                self.tokenizer,
+                self.args.vocab_size,
+                first_request.response_format,
+            ),
+        )
+
+    def chat_completion(
+        self,
+        request_batch: list[ChatCompletionRequestWithRawContent],
+    ) -> Generator:
+        first_request = request_batch[0]
+        sampling_params = first_request.sampling_params or SamplingParams()
+        max_gen_len = sampling_params.max_tokens
+        if max_gen_len is None or max_gen_len == 0 or max_gen_len >= self.args.max_seq_len:
+            max_gen_len = self.args.max_seq_len - 1
+
+        temperature, top_p = _infer_sampling_params(sampling_params)
+        yield from self.inner_generator.generate(
+            llm_inputs=[
+                self.formatter.encode_dialog_prompt(request.messages, _infer_tool_prompt_format(request))
+                for request in request_batch
+            ],
+            max_gen_len=max_gen_len,
+            temperature=temperature,
+            top_p=top_p,
+            logprobs=bool(first_request.logprobs),
+            echo=False,
+            logits_processor=get_logits_processor(
+                self.tokenizer,
+                self.args.vocab_size,
+                first_request.response_format,
+            ),
+        )