Merge branch 'main' into HuggingfacePostTrainingConfig-branch

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -84,7 +84,7 @@ MEMORY_QUERY_TOOL = "knowledge_search"
 WEB_SEARCH_TOOL = "web_search"
 RAG_TOOL_GROUP = "builtin::rag"
 
-logger = get_logger(name=__name__, category="agents")
+logger = get_logger(name=__name__, category="agents::meta_reference")
 
 
 class ChatAgent(ShieldRunnerMixin):

llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import uuid
 from collections.abc import AsyncGenerator
 from datetime import UTC, datetime
@@ -42,6 +41,7 @@ from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
 from llama_stack.core.datatypes import AccessRule
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
 from llama_stack.providers.utils.pagination import paginate_records
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
@@ -51,7 +51,7 @@ from .config import MetaReferenceAgentsImplConfig
 from .persistence import AgentInfo
 from .responses.openai_responses import OpenAIResponsesImpl
 
-logger = logging.getLogger()
+logger = get_logger(name=__name__, category="agents::meta_reference")
 
 
 class MetaReferenceAgentsImpl(Agents):

llama_stack/providers/inline/agents/meta_reference/persistence.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import json
-import logging
 import uuid
 from datetime import UTC, datetime
 
@@ -15,9 +14,10 @@ from llama_stack.core.access_control.access_control import AccessDeniedError, is
 from llama_stack.core.access_control.datatypes import AccessRule
 from llama_stack.core.datatypes import User
 from llama_stack.core.request_headers import get_authenticated_user
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.kvstore import KVStore
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="agents::meta_reference")
 
 
 class AgentSessionInfo(Session):

llama_stack/providers/inline/agents/meta_reference/responses/openai_responses.py

@@ -41,7 +41,7 @@ from .utils import (
     convert_response_text_to_chat_response_format,
 )
 
-logger = get_logger(name=__name__, category="responses")
+logger = get_logger(name=__name__, category="openai::responses")
 
 
 class OpenAIResponsePreviousResponseWithInputItems(BaseModel):

llama_stack/providers/inline/agents/meta_reference/responses/streaming.py

@@ -47,7 +47,7 @@ from llama_stack.log import get_logger
 from .types import ChatCompletionContext, ChatCompletionResult
 from .utils import convert_chat_choice_to_response_message, is_function_tool_call
 
-logger = get_logger(name=__name__, category="responses")
+logger = get_logger(name=__name__, category="agents::meta_reference")
 
 
 class StreamingResponseOrchestrator:

llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py

@@ -38,7 +38,7 @@ from llama_stack.log import get_logger
 
 from .types import ChatCompletionContext, ToolExecutionResult
 
-logger = get_logger(name=__name__, category="responses")
+logger = get_logger(name=__name__, category="agents::meta_reference")
 
 
 class ToolExecutor:

llama_stack/providers/inline/agents/meta_reference/responses/utils.py

@@ -17,6 +17,8 @@ from llama_stack.apis.agents.openai_responses import (
     OpenAIResponseOutputMessageContent,
     OpenAIResponseOutputMessageContentOutputText,
     OpenAIResponseOutputMessageFunctionToolCall,
+    OpenAIResponseOutputMessageMCPCall,
+    OpenAIResponseOutputMessageMCPListTools,
     OpenAIResponseText,
 )
 from llama_stack.apis.inference import (
@@ -99,14 +101,22 @@ async def convert_response_input_to_chat_messages(
     """
     messages: list[OpenAIMessageParam] = []
     if isinstance(input, list):
+        # extract all OpenAIResponseInputFunctionToolCallOutput items
+        # so their corresponding OpenAIToolMessageParam instances can
+        # be added immediately following the corresponding
+        # OpenAIAssistantMessageParam
+        tool_call_results = {}
         for input_item in input:
             if isinstance(input_item, OpenAIResponseInputFunctionToolCallOutput):
-                messages.append(
-                    OpenAIToolMessageParam(
-                        content=input_item.output,
-                        tool_call_id=input_item.call_id,
-                    )
+                tool_call_results[input_item.call_id] = OpenAIToolMessageParam(
+                    content=input_item.output,
+                    tool_call_id=input_item.call_id,
                 )
+
+        for input_item in input:
+            if isinstance(input_item, OpenAIResponseInputFunctionToolCallOutput):
+                # skip as these have been extracted and inserted in order
+                pass
             elif isinstance(input_item, OpenAIResponseOutputMessageFunctionToolCall):
                 tool_call = OpenAIChatCompletionToolCall(
                     index=0,
@@ -117,6 +127,28 @@ async def convert_response_input_to_chat_messages(
                     ),
                 )
                 messages.append(OpenAIAssistantMessageParam(tool_calls=[tool_call]))
+                if input_item.call_id in tool_call_results:
+                    messages.append(tool_call_results[input_item.call_id])
+                    del tool_call_results[input_item.call_id]
+            elif isinstance(input_item, OpenAIResponseOutputMessageMCPCall):
+                tool_call = OpenAIChatCompletionToolCall(
+                    index=0,
+                    id=input_item.id,
+                    function=OpenAIChatCompletionToolCallFunction(
+                        name=input_item.name,
+                        arguments=input_item.arguments,
+                    ),
+                )
+                messages.append(OpenAIAssistantMessageParam(tool_calls=[tool_call]))
+                messages.append(
+                    OpenAIToolMessageParam(
+                        content=input_item.output,
+                        tool_call_id=input_item.id,
+                    )
+                )
+            elif isinstance(input_item, OpenAIResponseOutputMessageMCPListTools):
+                # the tool list will be handled separately
+                pass
             else:
                 content = await convert_response_content_to_chat_content(input_item.content)
                 message_type = await get_message_type_by_role(input_item.role)
@@ -125,6 +157,10 @@ async def convert_response_input_to_chat_messages(
                         f"Llama Stack OpenAI Responses does not yet support message role '{input_item.role}' in this context"
                     )
                 messages.append(message_type(content=content))
+        if len(tool_call_results):
+            raise ValueError(
+                f"Received function_call_output(s) with call_id(s) {tool_call_results.keys()}, but no corresponding function_call"
+            )
     else:
         messages.append(OpenAIUserMessageParam(content=input))
     return messages

llama_stack/providers/inline/agents/meta_reference/safety.py

@@ -5,13 +5,13 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 
 from llama_stack.apis.inference import Message
 from llama_stack.apis.safety import Safety, SafetyViolation, ViolationLevel
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.telemetry import tracing
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="agents::meta_reference")
 
 
 class SafetyException(Exception):  # noqa: N818

llama_stack/providers/inline/batches/reference/batches.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import asyncio
+import hashlib
 import itertools
 import json
 import time
@@ -136,28 +137,45 @@ class ReferenceBatchesImpl(Batches):
         endpoint: str,
         completion_window: Literal["24h"],
         metadata: dict[str, str] | None = None,
+        idempotency_key: str | None = None,
     ) -> BatchObject:
         """
         Create a new batch for processing multiple API requests.
 
-        Error handling by levels -
-         0. Input param handling, results in 40x errors before processing, e.g.
-           - Wrong completion_window
-           - Invalid metadata types
-           - Unknown endpoint
-          -> no batch created
-         1. Errors preventing processing, result in BatchErrors aggregated in process_batch, e.g.
-           - input_file_id missing
-           - invalid json in file
-           - missing custom_id, method, url, body
-           - invalid model
-           - streaming
-          -> batch created, validation sends to failed status
-         2. Processing errors, result in error_file_id entries, e.g.
-           - Any error returned from inference endpoint
-          -> batch created, goes to completed status
+        This implementation provides optional idempotency: when an idempotency key
+        (idempotency_key) is provided, a deterministic ID is generated based on the input
+        parameters. If a batch with the same parameters already exists, it will be
+        returned instead of creating a duplicate. Without an idempotency key,
+        each request creates a new batch with a unique ID.
+
+        Args:
+            input_file_id: The ID of an uploaded file containing requests for the batch.
+            endpoint: The endpoint to be used for all requests in the batch.
+            completion_window: The time window within which the batch should be processed.
+            metadata: Optional metadata for the batch.
+            idempotency_key: Optional idempotency key for enabling idempotent behavior.
+
+        Returns:
+            The created or existing batch object.
         """
 
+        # Error handling by levels -
+        #  0. Input param handling, results in 40x errors before processing, e.g.
+        #    - Wrong completion_window
+        #    - Invalid metadata types
+        #    - Unknown endpoint
+        #   -> no batch created
+        #  1. Errors preventing processing, result in BatchErrors aggregated in process_batch, e.g.
+        #    - input_file_id missing
+        #    - invalid json in file
+        #    - missing custom_id, method, url, body
+        #    - invalid model
+        #    - streaming
+        #   -> batch created, validation sends to failed status
+        #  2. Processing errors, result in error_file_id entries, e.g.
+        #    - Any error returned from inference endpoint
+        #   -> batch created, goes to completed status
+
         # TODO: set expiration time for garbage collection
 
         if endpoint not in ["/v1/chat/completions"]:
@@ -171,6 +189,35 @@ class ReferenceBatchesImpl(Batches):
             )
 
         batch_id = f"batch_{uuid.uuid4().hex[:16]}"
+
+        # For idempotent requests, use the idempotency key for the batch ID
+        # This ensures the same key always maps to the same batch ID,
+        # allowing us to detect parameter conflicts
+        if idempotency_key is not None:
+            hash_input = idempotency_key.encode("utf-8")
+            hash_digest = hashlib.sha256(hash_input).hexdigest()[:24]
+            batch_id = f"batch_{hash_digest}"
+
+            try:
+                existing_batch = await self.retrieve_batch(batch_id)
+
+                if (
+                    existing_batch.input_file_id != input_file_id
+                    or existing_batch.endpoint != endpoint
+                    or existing_batch.completion_window != completion_window
+                    or existing_batch.metadata != metadata
+                ):
+                    raise ConflictError(
+                        f"Idempotency key '{idempotency_key}' was previously used with different parameters. "
+                        "Either use a new idempotency key or ensure all parameters match the original request."
+                    )
+
+                logger.info(f"Returning existing batch with ID: {batch_id}")
+                return existing_batch
+            except ResourceNotFoundError:
+                # Batch doesn't exist, continue with creation
+                pass
+
         current_time = int(time.time())
 
         batch = BatchObject(
@@ -185,6 +232,7 @@ class ReferenceBatchesImpl(Batches):
         )
 
         await self.kvstore.set(f"batch:{batch_id}", batch.to_json())
+        logger.info(f"Created new batch with ID: {batch_id}")
 
         if self.process_batches:
             task = asyncio.create_task(self._process_batch(batch_id))

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

@@ -11,6 +11,7 @@ from typing import Annotated
 
 from fastapi import 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 (
     Files,
@@ -20,12 +21,15 @@ from llama_stack.apis.files import (
     OpenAIFilePurpose,
 )
 from llama_stack.core.datatypes import AccessRule
+from llama_stack.log import get_logger
 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:
@@ -65,6 +69,18 @@ class LocalfsFilesImpl(Files):
         """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", policy=self.policy, 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,
@@ -157,37 +173,19 @@ class LocalfsFilesImpl(Files):
 
     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")
+        file_obj, _ = await self._lookup_file_id(file_id)
 
-        row = await self.sql_store.fetch_one("openai_files", policy=self.policy, 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"],
-        )
+        return file_obj
 
     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", policy=self.policy, 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"])
+        _, 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(
@@ -197,25 +195,17 @@ class LocalfsFilesImpl(Files):
 
     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", policy=self.policy, 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}")
+        file_obj, file_path = await self._lookup_file_id(file_id)
 
-        with open(file_path, "rb") as f:
-            content = f.read()
+        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=content,
+            content=file_path.read_bytes(),
             media_type="application/octet-stream",
-            headers={"Content-Disposition": f'attachment; filename="{row["filename"]}"'},
+            headers={"Content-Disposition": f'attachment; filename="{file_obj.filename}"'},
         )

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

@@ -12,7 +12,6 @@
 
 import copy
 import json
-import logging
 import multiprocessing
 import os
 import tempfile
@@ -32,13 +31,14 @@ from fairscale.nn.model_parallel.initialize import (
 from pydantic import BaseModel, Field
 from torch.distributed.launcher.api import LaunchConfig, elastic_launch
 
+from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import GenerationResult
 from llama_stack.providers.utils.inference.prompt_adapter import (
     ChatCompletionRequestWithRawContent,
     CompletionRequestWithRawContent,
 )
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="inference")
 
 
 class ProcessingMessageName(str, Enum):

llama_stack/providers/inline/inference/sentence_transformers/sentence_transformers.py

@@ -4,13 +4,11 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from collections.abc import AsyncGenerator
 
 from llama_stack.apis.inference import (
     CompletionResponse,
     InferenceProvider,
-    InterleavedContent,
     LogProbConfig,
     Message,
     ResponseFormat,
@@ -21,6 +19,7 @@ from llama_stack.apis.inference import (
     ToolPromptFormat,
 )
 from llama_stack.apis.models import ModelType
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.embedding_mixin import (
     SentenceTransformerEmbeddingMixin,
@@ -32,7 +31,7 @@ from llama_stack.providers.utils.inference.openai_compat import (
 
 from .config import SentenceTransformersInferenceConfig
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="inference")
 
 
 class SentenceTransformersInferenceImpl(
@@ -100,25 +99,3 @@ class SentenceTransformersInferenceImpl(
         tool_config: ToolConfig | None = None,
     ) -> AsyncGenerator:
         raise ValueError("Sentence transformers don't support chat completion")
-
-    async def batch_completion(
-        self,
-        model_id: str,
-        content_batch: list[InterleavedContent],
-        sampling_params: SamplingParams | None = None,
-        response_format: ResponseFormat | None = None,
-        logprobs: LogProbConfig | None = None,
-    ):
-        raise NotImplementedError("Batch completion is not supported for Sentence Transformers")
-
-    async def batch_chat_completion(
-        self,
-        model_id: str,
-        messages_batch: list[list[Message]],
-        sampling_params: SamplingParams | None = None,
-        tools: list[ToolDefinition] | None = None,
-        tool_config: ToolConfig | None = None,
-        response_format: ResponseFormat | None = None,
-        logprobs: LogProbConfig | None = None,
-    ):
-        raise NotImplementedError("Batch chat completion is not supported for Sentence Transformers")

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device.py

@@ -6,7 +6,6 @@
 
 import gc
 import json
-import logging
 import multiprocessing
 from pathlib import Path
 from typing import Any
@@ -28,6 +27,7 @@ from llama_stack.apis.post_training import (
     LoraFinetuningConfig,
     TrainingConfig,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 
 from ..config import HuggingFacePostTrainingConfig
@@ -44,7 +44,7 @@ from ..utils import (
     split_dataset,
 )
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="post_training")
 
 
 class HFFinetuningSingleDevice:

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device_dpo.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import gc
-import logging
 import multiprocessing
 from pathlib import Path
 from typing import Any
@@ -24,6 +23,7 @@ from llama_stack.apis.post_training import (
     DPOAlignmentConfig,
     TrainingConfig,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 
 from ..config import HuggingFacePostTrainingConfig
@@ -40,7 +40,7 @@ from ..utils import (
     split_dataset,
 )
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="post_training")
 
 
 class HFDPOAlignmentSingleDevice:

llama_stack/providers/inline/post_training/huggingface/utils.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import os
 import signal
 import sys
@@ -19,10 +18,11 @@ from transformers import AutoConfig, AutoModelForCausalLM
 
 from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.post_training import Checkpoint, TrainingConfig
+from llama_stack.log import get_logger
 
 from .config import HuggingFacePostTrainingConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="post_training")
 
 
 def setup_environment():

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import os
 import time
 from datetime import UTC, datetime
@@ -19,6 +18,7 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torchtune import modules, training
 from torchtune import utils as torchtune_utils
 from torchtune.data import padded_collate_sft
+from torchtune.models.llama3._tokenizer import Llama3Tokenizer
 from torchtune.modules.loss import CEWithChunkedOutputLoss
 from torchtune.modules.peft import (
     get_adapter_params,
@@ -45,6 +45,7 @@ from llama_stack.apis.post_training import (
 )
 from llama_stack.core.utils.config_dirs import DEFAULT_CHECKPOINT_DIR
 from llama_stack.core.utils.model_utils import model_local_dir
+from llama_stack.log import get_logger
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 from llama_stack.providers.inline.post_training.torchtune.common import utils
@@ -56,9 +57,7 @@ from llama_stack.providers.inline.post_training.torchtune.config import (
 )
 from llama_stack.providers.inline.post_training.torchtune.datasets.sft import SFTDataset
 
-log = logging.getLogger(__name__)
-
-from torchtune.models.llama3._tokenizer import Llama3Tokenizer
+log = get_logger(name=__name__, category="post_training")
 
 
 class LoraFinetuningSingleDevice:

llama_stack/providers/inline/safety/code_scanner/code_scanner.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import uuid
 from typing import TYPE_CHECKING, Any
 
@@ -20,13 +19,14 @@ from llama_stack.apis.safety import (
 )
 from llama_stack.apis.safety.safety import ModerationObject, ModerationObjectResults
 from llama_stack.apis.shields import Shield
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 
 from .config import CodeScannerConfig
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="safety")
 
 ALLOWED_CODE_SCANNER_MODEL_IDS = [
     "code-scanner",

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import re
 import uuid
 from string import Template
@@ -21,6 +20,7 @@ from llama_stack.apis.safety import (
 from llama_stack.apis.safety.safety import ModerationObject, ModerationObjectResults
 from llama_stack.apis.shields import Shield
 from llama_stack.core.datatypes import Api
+from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import Role
 from llama_stack.models.llama.sku_types import CoreModelId
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
@@ -132,6 +132,8 @@ Provide your safety assessment for ONLY THE LAST $agent_type message in the abov
 
 PROMPT_TEMPLATE = Template(f"{PROMPT_TASK}{SAFETY_CATEGORIES}{PROMPT_CONVERSATION}{PROMPT_INSTRUCTIONS}")
 
+logger = get_logger(name=__name__, category="safety")
+
 
 class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
     def __init__(self, config: LlamaGuardConfig, deps) -> None:
@@ -407,7 +409,7 @@ class LlamaGuardShield:
             unsafe_code_list = [code.strip() for code in unsafe_code.split(",")]
             invalid_codes = [code for code in unsafe_code_list if code not in SAFETY_CODE_TO_CATEGORIES_MAP]
             if invalid_codes:
-                logging.warning(f"Invalid safety codes returned: {invalid_codes}")
+                logger.warning(f"Invalid safety codes returned: {invalid_codes}")
                 # just returning safe object, as we don't know what the invalid codes can map to
                 return ModerationObject(
                     id=f"modr-{uuid.uuid4()}",

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from typing import Any
 
 import torch
@@ -21,6 +20,7 @@ from llama_stack.apis.safety import (
 from llama_stack.apis.safety.safety import ModerationObject
 from llama_stack.apis.shields import Shield
 from llama_stack.core.utils.model_utils import model_local_dir
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
@@ -28,7 +28,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 
 from .config import PromptGuardConfig, PromptGuardType
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="safety")
 
 PROMPT_GUARD_MODEL = "Prompt-Guard-86M"
 

llama_stack/providers/inline/scoring/basic/utils/ifeval_utils.py

@@ -7,7 +7,6 @@
 import collections
 import functools
 import json
-import logging
 import random
 import re
 import string
@@ -20,7 +19,9 @@ import nltk
 from pythainlp.tokenize import sent_tokenize as sent_tokenize_thai
 from pythainlp.tokenize import word_tokenize as word_tokenize_thai
 
-logger = logging.getLogger()
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="scoring")
 
 WORD_LIST = [
     "western",

llama_stack/providers/inline/telemetry/meta_reference/telemetry.py

@@ -4,13 +4,11 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
+import datetime
 import threading
 from typing import Any
 
 from opentelemetry import metrics, trace
-
-logger = logging.getLogger(__name__)
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.metrics import MeterProvider
@@ -40,6 +38,7 @@ from llama_stack.apis.telemetry import (
     UnstructuredLogEvent,
 )
 from llama_stack.core.datatypes import Api
+from llama_stack.log import get_logger
 from llama_stack.providers.inline.telemetry.meta_reference.console_span_processor import (
     ConsoleSpanProcessor,
 )
@@ -61,6 +60,8 @@ _GLOBAL_STORAGE: dict[str, dict[str | int, Any]] = {
 _global_lock = threading.Lock()
 _TRACER_PROVIDER = None
 
+logger = get_logger(name=__name__, category="telemetry")
+
 
 def is_tracing_enabled(tracer):
     with tracer.start_as_current_span("check_tracing") as span:
@@ -145,11 +146,41 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
         metric_name: str,
         start_time: int,
         end_time: int | None = None,
-        granularity: str | None = "1d",
+        granularity: str | None = None,
         query_type: MetricQueryType = MetricQueryType.RANGE,
         label_matchers: list[MetricLabelMatcher] | None = None,
     ) -> QueryMetricsResponse:
-        raise NotImplementedError("Querying metrics is not implemented")
+        """Query metrics from the telemetry store.
+
+        Args:
+            metric_name: The name of the metric to query (e.g., "prompt_tokens")
+            start_time: Start time as Unix timestamp
+            end_time: End time as Unix timestamp (defaults to now if None)
+            granularity: Time granularity for aggregation
+            query_type: Type of query (RANGE or INSTANT)
+            label_matchers: Label filters to apply
+
+        Returns:
+            QueryMetricsResponse with metric time series data
+        """
+        # Convert timestamps to datetime objects
+        start_dt = datetime.datetime.fromtimestamp(start_time, datetime.UTC)
+        end_dt = datetime.datetime.fromtimestamp(end_time, datetime.UTC) if end_time else None
+
+        # Use SQLite trace store if available
+        if hasattr(self, "trace_store") and self.trace_store:
+            return await self.trace_store.query_metrics(
+                metric_name=metric_name,
+                start_time=start_dt,
+                end_time=end_dt,
+                granularity=granularity,
+                query_type=query_type,
+                label_matchers=label_matchers,
+            )
+        else:
+            raise ValueError(
+                f"In order to query_metrics, you must have {TelemetrySink.SQLITE} set in your telemetry sinks"
+            )
 
     def _log_unstructured(self, event: UnstructuredLogEvent, ttl_seconds: int) -> None:
         with self._lock:

llama_stack/providers/inline/tool_runtime/rag/memory.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 import secrets
 import string
 from typing import Any
@@ -32,6 +31,7 @@ from llama_stack.apis.tools import (
     ToolRuntime,
 )
 from llama_stack.apis.vector_io import QueryChunksResponse, VectorIO
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import ToolGroupsProtocolPrivate
 from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
 from llama_stack.providers.utils.memory.vector_store import (
@@ -42,7 +42,7 @@ from llama_stack.providers.utils.memory.vector_store import (
 from .config import RagToolRuntimeConfig
 from .context_retriever import generate_rag_query
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="tool_runtime")
 
 
 def make_random_string(length: int = 8):

llama_stack/providers/inline/vector_io/faiss/faiss.py

@@ -8,7 +8,6 @@ import asyncio
 import base64
 import io
 import json
-import logging
 from typing import Any
 
 import faiss
@@ -24,6 +23,7 @@ from llama_stack.apis.vector_io import (
     QueryChunksResponse,
     VectorIO,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import (
     HealthResponse,
     HealthStatus,
@@ -40,7 +40,7 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import FaissVectorIOConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="vector_io")
 
 VERSION = "v3"
 VECTOR_DBS_PREFIX = f"vector_dbs:{VERSION}::"

llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 import re
 import sqlite3
 import struct
@@ -24,6 +23,7 @@ from llama_stack.apis.vector_io import (
     QueryChunksResponse,
     VectorIO,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import VectorDBsProtocolPrivate
 from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.kvstore.api import KVStore
@@ -36,7 +36,7 @@ from llama_stack.providers.utils.memory.vector_store import (
     VectorDBWithIndex,
 )
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="vector_io")
 
 # Specifying search mode is dependent on the VectorIO provider.
 VECTOR_SEARCH = "vector"