pre-commit fixes

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -8,7 +8,6 @@ from typing import AsyncGenerator, AsyncIterator, List, Optional, Union
 
 import litellm
 
-from llama_stack import logcat
 from llama_stack.apis.common.content_types import (
     InterleavedContent,
     InterleavedContentItem,
@@ -33,6 +32,7 @@ from llama_stack.apis.inference import (
 )
 from llama_stack.apis.models.models import Model
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
@@ -47,6 +47,8 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 
+logger = get_logger(name=__name__, category="inference")
+
 
 class LiteLLMOpenAIMixin(
     ModelRegistryHelper,
@@ -74,7 +76,7 @@ class LiteLLMOpenAIMixin(
         self,
         model_id: str,
         content: InterleavedContent,
-        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        sampling_params: Optional[SamplingParams] = None,
         response_format: Optional[ResponseFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
@@ -85,7 +87,7 @@ class LiteLLMOpenAIMixin(
         self,
         model_id: str,
         messages: List[Message],
-        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        sampling_params: Optional[SamplingParams] = None,
         tools: Optional[List[ToolDefinition]] = None,
         tool_choice: Optional[ToolChoice] = ToolChoice.auto,
         tool_prompt_format: Optional[ToolPromptFormat] = None,
@@ -94,6 +96,8 @@ class LiteLLMOpenAIMixin(
         logprobs: Optional[LogProbConfig] = None,
         tool_config: Optional[ToolConfig] = None,
     ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionResponseStreamChunk]]:
+        if sampling_params is None:
+            sampling_params = SamplingParams()
         model = await self.model_store.get_model(model_id)
         request = ChatCompletionRequest(
             model=model.provider_resource_id,
@@ -107,8 +111,7 @@ class LiteLLMOpenAIMixin(
         )
 
         params = await self._get_params(request)
-        logcat.debug("inference", f"params to litellm (openai compat): {params}")
-
+        logger.debug(f"params to litellm (openai compat): {params}")
         # unfortunately, we need to use synchronous litellm.completion here because litellm
         # caches various httpx.client objects in a non-eventloop aware manner
         response = litellm.completion(**params)

llama_stack/providers/utils/inference/openai_compat.py

@@ -615,6 +615,14 @@ def convert_tool_call(
     return valid_tool_call
 
 
+PYTHON_TYPE_TO_LITELLM_TYPE = {
+    "int": "integer",
+    "float": "number",
+    "bool": "boolean",
+    "str": "string",
+}
+
+
 def convert_tooldef_to_openai_tool(tool: ToolDefinition) -> dict:
     """
     Convert a ToolDefinition to an OpenAI API-compatible dictionary.
@@ -675,7 +683,7 @@ def convert_tooldef_to_openai_tool(tool: ToolDefinition) -> dict:
         properties = parameters["properties"]
         required = []
         for param_name, param in tool.parameters.items():
-            properties[param_name] = {"type": param.param_type}
+            properties[param_name] = {"type": PYTHON_TYPE_TO_LITELLM_TYPE.get(param.param_type, param.param_type)}
             if param.description:
                 properties[param_name].update(description=param.description)
             if param.default:

llama_stack/providers/utils/inference/prompt_adapter.py

@@ -8,14 +8,12 @@ import asyncio
 import base64
 import io
 import json
-import logging
 import re
 from typing import List, Optional, Tuple, Union
 
 import httpx
 from PIL import Image as PIL_Image
 
-from llama_stack import logcat
 from llama_stack.apis.common.content_types import (
     ImageContentItem,
     InterleavedContent,
@@ -34,6 +32,7 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     UserMessage,
 )
+from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import (
     ModelFamily,
     RawContent,
@@ -58,7 +57,7 @@ from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.utils.inference import supported_inference_models
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="inference")
 
 
 class ChatCompletionRequestWithRawContent(ChatCompletionRequest):
@@ -464,7 +463,7 @@ def _get_tool_choice_prompt(tool_choice: ToolChoice | str, tools: List[ToolDefin
 def get_default_tool_prompt_format(model: str) -> ToolPromptFormat:
     llama_model = resolve_model(model)
     if llama_model is None:
-        logcat.warning("inference", f"Could not resolve model {model}, defaulting to json tool prompt format")
+        log.warning(f"Could not resolve model {model}, defaulting to json tool prompt format")
         return ToolPromptFormat.json
 
     if llama_model.model_family == ModelFamily.llama3_1 or (

llama_stack/providers/utils/kvstore/config.py

@@ -55,11 +55,11 @@ class SqliteKVStoreConfig(CommonConfig):
     )
 
     @classmethod
-    def sample_run_config(cls, __distro_dir__: str = "runtime", db_name: str = "kvstore.db"):
+    def sample_run_config(cls, __distro_dir__: str, db_name: str = "kvstore.db"):
         return {
             "type": "sqlite",
             "namespace": None,
-            "db_path": "${env.SQLITE_STORE_DIR:~/.llama/" + __distro_dir__ + "}/" + db_name,
+            "db_path": "${env.SQLITE_STORE_DIR:" + __distro_dir__ + "}/" + db_name,
         }
 
 

llama_stack/providers/utils/kvstore/mongodb/mongodb.py

@@ -8,9 +8,11 @@ import logging
 from datetime import datetime
 from typing import List, Optional
 
-from pymongo import MongoClient
+from pymongo import AsyncMongoClient
 
-from llama_stack.providers.utils.kvstore import KVStore, MongoDBKVStoreConfig
+from llama_stack.providers.utils.kvstore import KVStore
+
+from ..config import MongoDBKVStoreConfig
 
 log = logging.getLogger(__name__)
 
@@ -30,7 +32,7 @@ class MongoDBKVStoreImpl(KVStore):
                 "password": self.config.password,
             }
             conn_creds = {k: v for k, v in conn_creds.items() if v is not None}
-            self.conn = MongoClient(**conn_creds)
+            self.conn = AsyncMongoClient(**conn_creds)
             self.collection = self.conn[self.config.db][self.config.collection_name]
         except Exception as e:
             log.exception("Could not connect to MongoDB database server")
@@ -44,17 +46,17 @@ class MongoDBKVStoreImpl(KVStore):
     async def set(self, key: str, value: str, expiration: Optional[datetime] = None) -> None:
         key = self._namespaced_key(key)
         update_query = {"$set": {"value": value, "expiration": expiration}}
-        self.collection.update_one({"key": key}, update_query, upsert=True)
+        await self.collection.update_one({"key": key}, update_query, upsert=True)
 
     async def get(self, key: str) -> Optional[str]:
         key = self._namespaced_key(key)
         query = {"key": key}
-        result = self.collection.find_one(query, {"value": 1, "_id": 0})
+        result = await self.collection.find_one(query, {"value": 1, "_id": 0})
         return result["value"] if result else None
 
     async def delete(self, key: str) -> None:
         key = self._namespaced_key(key)
-        self.collection.delete_one({"key": key})
+        await self.collection.delete_one({"key": key})
 
     async def range(self, start_key: str, end_key: str) -> List[str]:
         start_key = self._namespaced_key(start_key)
@@ -63,4 +65,7 @@ class MongoDBKVStoreImpl(KVStore):
             "key": {"$gte": start_key, "$lt": end_key},
         }
         cursor = self.collection.find(query, {"value": 1, "_id": 0}).sort("key", 1)
-        return [doc["value"] for doc in cursor]
+        result = []
+        async for doc in cursor:
+            result.append(doc["value"])
+        return result

llama_stack/providers/utils/memory/vector_store.py

@@ -12,11 +12,9 @@ from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 from urllib.parse import unquote
 
-import chardet
 import httpx
 import numpy as np
 from numpy.typing import NDArray
-from pypdf import PdfReader
 
 from llama_stack.apis.common.content_types import (
     URL,
@@ -38,6 +36,8 @@ log = logging.getLogger(__name__)
 def parse_pdf(data: bytes) -> str:
     # For PDF and DOC/DOCX files, we can't reliably convert to string
     pdf_bytes = io.BytesIO(data)
+    from pypdf import PdfReader
+
     pdf_reader = PdfReader(pdf_bytes)
     return "\n".join([page.extract_text() for page in pdf_reader.pages])
 
@@ -75,6 +75,8 @@ def content_from_data(data_url: str) -> str:
 
     encoding = parts["encoding"]
     if not encoding:
+        import chardet
+
         detected = chardet.detect(data)
         encoding = detected["encoding"]
 

llama_stack/providers/utils/scoring/base_scoring_fn.py

@@ -73,6 +73,11 @@ class RegisteredBaseScoringFn(BaseScoringFn):
             raise ValueError(f"Scoring function def with identifier {scoring_fn.identifier} already exists.")
         self.supported_fn_defs_registry[scoring_fn.identifier] = scoring_fn
 
+    def unregister_scoring_fn_def(self, scoring_fn_id: str) -> None:
+        if scoring_fn_id not in self.supported_fn_defs_registry:
+            raise ValueError(f"Scoring function def with identifier {scoring_fn_id} does not exist.")
+        del self.supported_fn_defs_registry[scoring_fn_id]
+
     @abstractmethod
     async def score_row(
         self,

llama_stack/providers/utils/telemetry/tracing.py

@@ -6,6 +6,7 @@
 
 import asyncio
 import base64
+import contextvars
 import logging
 import queue
 import threading
@@ -24,9 +25,10 @@ from llama_stack.apis.telemetry import (
     Telemetry,
     UnstructuredLogEvent,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.telemetry.trace_protocol import serialize_value
 
-log = logging.getLogger(__name__)
+logger = get_logger(__name__, category="core")
 
 
 def generate_short_uuid(len: int = 8):
@@ -36,7 +38,7 @@ def generate_short_uuid(len: int = 8):
     return encoded.rstrip(b"=").decode("ascii")[:len]
 
 
-CURRENT_TRACE_CONTEXT = None
+CURRENT_TRACE_CONTEXT = contextvars.ContextVar("trace_context", default=None)
 BACKGROUND_LOGGER = None
 
 
@@ -51,7 +53,7 @@ class BackgroundLogger:
         try:
             self.log_queue.put_nowait(event)
         except queue.Full:
-            log.error("Log queue is full, dropping event")
+            logger.error("Log queue is full, dropping event")
 
     def _process_logs(self):
         while True:
@@ -129,35 +131,36 @@ def setup_logger(api: Telemetry, level: int = logging.INFO):
 
     if BACKGROUND_LOGGER is None:
         BACKGROUND_LOGGER = BackgroundLogger(api)
-    logger = logging.getLogger()
-    logger.setLevel(level)
-    logger.addHandler(TelemetryHandler())
+    root_logger = logging.getLogger()
+    root_logger.setLevel(level)
+    root_logger.addHandler(TelemetryHandler())
 
 
 async def start_trace(name: str, attributes: Dict[str, Any] = None) -> TraceContext:
     global CURRENT_TRACE_CONTEXT, BACKGROUND_LOGGER
 
     if BACKGROUND_LOGGER is None:
-        log.info("No Telemetry implementation set. Skipping trace initialization...")
+        logger.debug("No Telemetry implementation set. Skipping trace initialization...")
         return
 
     trace_id = generate_short_uuid(16)
     context = TraceContext(BACKGROUND_LOGGER, trace_id)
     context.push_span(name, {"__root__": True, **(attributes or {})})
 
-    CURRENT_TRACE_CONTEXT = context
+    CURRENT_TRACE_CONTEXT.set(context)
     return context
 
 
 async def end_trace(status: SpanStatus = SpanStatus.OK):
     global CURRENT_TRACE_CONTEXT
 
-    context = CURRENT_TRACE_CONTEXT
+    context = CURRENT_TRACE_CONTEXT.get()
     if context is None:
+        logger.debug("No trace context to end")
         return
 
     context.pop_span(status)
-    CURRENT_TRACE_CONTEXT = None
+    CURRENT_TRACE_CONTEXT.set(None)
 
 
 def severity(levelname: str) -> LogSeverity:
@@ -188,7 +191,7 @@ class TelemetryHandler(logging.Handler):
         if BACKGROUND_LOGGER is None:
             raise RuntimeError("Telemetry API not initialized")
 
-        context = CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT.get()
         if context is None:
             return
 
@@ -218,16 +221,22 @@ class SpanContextManager:
 
     def __enter__(self):
         global CURRENT_TRACE_CONTEXT
-        context = CURRENT_TRACE_CONTEXT
-        if context:
-            self.span = context.push_span(self.name, self.attributes)
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to push span")
+            return self
+
+        self.span = context.push_span(self.name, self.attributes)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
         global CURRENT_TRACE_CONTEXT
-        context = CURRENT_TRACE_CONTEXT
-        if context:
-            context.pop_span()
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to pop span")
+            return
+
+        context.pop_span()
 
     def set_attribute(self, key: str, value: Any):
         if self.span:
@@ -237,16 +246,22 @@ class SpanContextManager:
 
     async def __aenter__(self):
         global CURRENT_TRACE_CONTEXT
-        context = CURRENT_TRACE_CONTEXT
-        if context:
-            self.span = context.push_span(self.name, self.attributes)
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to push span")
+            return self
+
+        self.span = context.push_span(self.name, self.attributes)
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
         global CURRENT_TRACE_CONTEXT
-        context = CURRENT_TRACE_CONTEXT
-        if context:
-            context.pop_span()
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to pop span")
+            return
+
+        context.pop_span()
 
     def __call__(self, func: Callable):
         @wraps(func)
@@ -275,7 +290,11 @@ def span(name: str, attributes: Dict[str, Any] = None):
 
 def get_current_span() -> Optional[Span]:
     global CURRENT_TRACE_CONTEXT
-    context = CURRENT_TRACE_CONTEXT
+    if CURRENT_TRACE_CONTEXT is None:
+        logger.debug("No trace context to get current span")
+        return None
+
+    context = CURRENT_TRACE_CONTEXT.get()
     if context:
         return context.get_current_span()
     return None