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

llama_stack/providers/utils/vector_io/vector_utils.py

@@ -0,0 +1,156 @@
+# 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 hashlib
+import re
+import uuid
+
+
+def generate_chunk_id(document_id: str, chunk_text: str, chunk_window: str | None = None) -> str:
+    """
+    Generate a unique chunk ID using a hash of the document ID and chunk text.
+    Then use the first 32 characters of the hash to create a UUID.
+    """
+    hash_input = f"{document_id}:{chunk_text}".encode()
+    if chunk_window:
+        hash_input += f":{chunk_window}".encode()
+    return str(uuid.UUID(hashlib.sha256(hash_input).hexdigest()[:32]))
+
+
+def proper_case(s: str) -> str:
+    """Convert a string to proper case (first letter uppercase, rest lowercase)."""
+    return s[0].upper() + s[1:].lower() if s else s
+
+
+def sanitize_collection_name(name: str, weaviate_format=False) -> str:
+    """
+    Sanitize collection name to ensure it only contains numbers, letters, and underscores.
+    Any other characters are replaced with underscores.
+    """
+    if not weaviate_format:
+        s = re.sub(r"[^a-zA-Z0-9_]", "_", name)
+    else:
+        s = proper_case(re.sub(r"[^a-zA-Z0-9]", "", name))
+    return s
+
+
+class WeightedInMemoryAggregator:
+    @staticmethod
+    def _normalize_scores(scores: dict[str, float]) -> dict[str, float]:
+        """
+        Normalize scores to 0-1 range using min-max normalization.
+
+        Args:
+            scores: dictionary of scores with document IDs as keys and scores as values
+
+        Returns:
+            Normalized scores with document IDs as keys and normalized scores as values
+        """
+        if not scores:
+            return {}
+        min_score, max_score = min(scores.values()), max(scores.values())
+        score_range = max_score - min_score
+        if score_range > 0:
+            return {doc_id: (score - min_score) / score_range for doc_id, score in scores.items()}
+        return dict.fromkeys(scores, 1.0)
+
+    @staticmethod
+    def weighted_rerank(
+        vector_scores: dict[str, float],
+        keyword_scores: dict[str, float],
+        alpha: float = 0.5,
+    ) -> dict[str, float]:
+        """
+        Rerank via weighted average of scores.
+
+        Args:
+            vector_scores: scores from vector search
+            keyword_scores: scores from keyword search
+            alpha: weight factor between 0 and 1 (default: 0.5)
+                   0 = keyword only, 1 = vector only, 0.5 = equal weight
+
+        Returns:
+            All unique document IDs with weighted combined scores
+        """
+        all_ids = set(vector_scores.keys()) | set(keyword_scores.keys())
+        normalized_vector_scores = WeightedInMemoryAggregator._normalize_scores(vector_scores)
+        normalized_keyword_scores = WeightedInMemoryAggregator._normalize_scores(keyword_scores)
+
+        # Weighted formula: score = (1-alpha) * keyword_score + alpha * vector_score
+        # alpha=0 means keyword only, alpha=1 means vector only
+        return {
+            doc_id: ((1 - alpha) * normalized_keyword_scores.get(doc_id, 0.0))
+            + (alpha * normalized_vector_scores.get(doc_id, 0.0))
+            for doc_id in all_ids
+        }
+
+    @staticmethod
+    def rrf_rerank(
+        vector_scores: dict[str, float],
+        keyword_scores: dict[str, float],
+        impact_factor: float = 60.0,
+    ) -> dict[str, float]:
+        """
+        Rerank via Reciprocal Rank Fusion.
+
+        Args:
+            vector_scores: scores from vector search
+            keyword_scores: scores from keyword search
+            impact_factor: impact factor for RRF (default: 60.0)
+
+        Returns:
+            All unique document IDs with RRF combined scores
+        """
+
+        # Convert scores to ranks
+        vector_ranks = {
+            doc_id: i + 1
+            for i, (doc_id, _) in enumerate(sorted(vector_scores.items(), key=lambda x: x[1], reverse=True))
+        }
+        keyword_ranks = {
+            doc_id: i + 1
+            for i, (doc_id, _) in enumerate(sorted(keyword_scores.items(), key=lambda x: x[1], reverse=True))
+        }
+
+        all_ids = set(vector_scores.keys()) | set(keyword_scores.keys())
+        rrf_scores = {}
+        for doc_id in all_ids:
+            vector_rank = vector_ranks.get(doc_id, float("inf"))
+            keyword_rank = keyword_ranks.get(doc_id, float("inf"))
+
+            # RRF formula: score = 1/(k + r) where k is impact_factor (default: 60.0) and r is the rank
+            rrf_scores[doc_id] = (1.0 / (impact_factor + vector_rank)) + (1.0 / (impact_factor + keyword_rank))
+        return rrf_scores
+
+    @staticmethod
+    def combine_search_results(
+        vector_scores: dict[str, float],
+        keyword_scores: dict[str, float],
+        reranker_type: str = "rrf",
+        reranker_params: dict[str, float] | None = None,
+    ) -> dict[str, float]:
+        """
+        Combine vector and keyword search results using specified reranking strategy.
+
+        Args:
+            vector_scores: scores from vector search
+            keyword_scores: scores from keyword search
+            reranker_type: type of reranker to use (default: RERANKER_TYPE_RRF)
+            reranker_params: parameters for the reranker
+
+        Returns:
+            All unique document IDs with combined scores
+        """
+        if reranker_params is None:
+            reranker_params = {}
+
+        if reranker_type == "weighted":
+            alpha = reranker_params.get("alpha", 0.5)
+            return WeightedInMemoryAggregator.weighted_rerank(vector_scores, keyword_scores, alpha)
+        else:
+            # Default to RRF for None, RRF, or any unknown types
+            impact_factor = reranker_params.get("impact_factor", 60.0)
+            return WeightedInMemoryAggregator.rrf_rerank(vector_scores, keyword_scores, impact_factor)

llama_stack/schema_utils.py

@@ -0,0 +1,118 @@
+# 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 collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, TypeVar
+
+from .strong_typing.schema import json_schema_type, register_schema  # noqa: F401
+
+
+class ExtraBodyField[T]:
+    """
+    Marker annotation for parameters that arrive via extra_body in the client SDK.
+
+    These parameters:
+    - Will NOT appear in the generated client SDK method signature
+    - WILL be documented in OpenAPI spec under x-llama-stack-extra-body-params
+    - MUST be passed via the extra_body parameter in client SDK calls
+    - WILL be available in server-side method signature with proper typing
+
+    Example:
+        ```python
+        async def create_openai_response(
+            self,
+            input: str,
+            model: str,
+            shields: Annotated[
+                list[str] | None, ExtraBodyField("List of shields to apply")
+            ] = None,
+        ) -> ResponseObject:
+            # shields is available here with proper typing
+            if shields:
+                print(f"Using shields: {shields}")
+        ```
+
+        Client usage:
+        ```python
+        client.responses.create(
+            input="hello", model="llama-3", extra_body={"shields": ["shield-1"]}
+        )
+        ```
+    """
+
+    def __init__(self, description: str | None = None):
+        self.description = description
+
+
+@dataclass
+class WebMethod:
+    level: str | None = None
+    route: str | None = None
+    public: bool = False
+    request_examples: list[Any] | None = None
+    response_examples: list[Any] | None = None
+    method: str | None = None
+    raw_bytes_request_body: bool | None = False
+    # A descriptive name of the corresponding span created by tracing
+    descriptive_name: str | None = None
+    required_scope: str | None = None
+    deprecated: bool | None = False
+    require_authentication: bool | None = True
+
+
+CallableT = TypeVar("CallableT", bound=Callable[..., Any])
+
+
+def webmethod(
+    route: str | None = None,
+    method: str | None = None,
+    level: str | None = None,
+    public: bool | None = False,
+    request_examples: list[Any] | None = None,
+    response_examples: list[Any] | None = None,
+    raw_bytes_request_body: bool | None = False,
+    descriptive_name: str | None = None,
+    required_scope: str | None = None,
+    deprecated: bool | None = False,
+    require_authentication: bool | None = True,
+) -> Callable[[CallableT], CallableT]:
+    """
+    Decorator that supplies additional metadata to an endpoint operation function.
+
+    :param route: The URL path pattern associated with this operation which path parameters are substituted into.
+    :param public: True if the operation can be invoked without prior authentication.
+    :param request_examples: Sample requests that the operation might take. Pass a list of objects, not JSON.
+    :param response_examples: Sample responses that the operation might produce. Pass a list of objects, not JSON.
+    :param required_scope: Required scope for this endpoint (e.g., 'monitoring.viewer').
+    :param require_authentication: Whether this endpoint requires authentication (default True).
+    """
+
+    def wrap(func: CallableT) -> CallableT:
+        webmethod_obj = WebMethod(
+            route=route,
+            method=method,
+            level=level,
+            public=public or False,
+            request_examples=request_examples,
+            response_examples=response_examples,
+            raw_bytes_request_body=raw_bytes_request_body,
+            descriptive_name=descriptive_name,
+            required_scope=required_scope,
+            deprecated=deprecated,
+            require_authentication=require_authentication if require_authentication is not None else True,
+        )
+
+        # Store all webmethods in a list to support multiple decorators
+        if not hasattr(func, "__webmethods__"):
+            func.__webmethods__ = []  # type: ignore
+        func.__webmethods__.append(webmethod_obj)  # type: ignore
+
+        # Keep the last one as __webmethod__ for backwards compatibility
+        func.__webmethod__ = webmethod_obj  # type: ignore
+        return func
+
+    return wrap

llama_stack/strong_typing/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+"""
+Type-safe data interchange for Python data classes.
+
+Provides auxiliary services for working with Python type annotations, converting typed data to and from JSON,
+and generating a JSON schema for a complex type.
+"""
+
+__version__ = "0.3.4"
+__author__ = "Levente Hunyadi"
+__copyright__ = "Copyright 2021-2024, Levente Hunyadi"
+__license__ = "MIT"
+__maintainer__ = "Levente Hunyadi"
+__status__ = "Production"

llama_stack/strong_typing/auxiliary.py

@@ -0,0 +1,228 @@
+# 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.
+
+"""
+Type-safe data interchange for Python data classes.
+
+:see: https://github.com/hunyadi/strong_typing
+"""
+
+import dataclasses
+import sys
+from dataclasses import is_dataclass
+from typing import Callable, Dict, Optional, Type, TypeVar, Union, overload
+
+if sys.version_info >= (3, 9):
+    from typing import Annotated as Annotated
+else:
+    from typing_extensions import Annotated as Annotated
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias as TypeAlias
+else:
+    from typing_extensions import TypeAlias as TypeAlias
+
+if sys.version_info >= (3, 11):
+    from typing import dataclass_transform as dataclass_transform
+else:
+    from typing_extensions import dataclass_transform as dataclass_transform
+
+T = TypeVar("T")
+
+
+def _compact_dataclass_repr(obj: object) -> str:
+    """
+    Compact data-class representation where positional arguments are used instead of keyword arguments.
+
+    :param obj: A data-class object.
+    :returns: A string that matches the pattern `Class(arg1, arg2, ...)`.
+    """
+
+    if is_dataclass(obj):
+        arglist = ", ".join(repr(getattr(obj, field.name)) for field in dataclasses.fields(obj))
+        return f"{obj.__class__.__name__}({arglist})"
+    else:
+        return obj.__class__.__name__
+
+
+class CompactDataClass:
+    "A data class whose repr() uses positional rather than keyword arguments."
+
+    def __repr__(self) -> str:
+        return _compact_dataclass_repr(self)
+
+
+@overload
+def typeannotation(cls: Type[T], /) -> Type[T]: ...
+
+
+@overload
+def typeannotation(cls: None, *, eq: bool = True, order: bool = False) -> Callable[[Type[T]], Type[T]]: ...
+
+
+@dataclass_transform(eq_default=True, order_default=False)
+def typeannotation(
+    cls: Optional[Type[T]] = None, *, eq: bool = True, order: bool = False
+) -> Union[Type[T], Callable[[Type[T]], Type[T]]]:
+    """
+    Returns the same class as was passed in, with dunder methods added based on the fields defined in the class.
+
+    :param cls: The data-class type to transform into a type annotation.
+    :param eq: Whether to generate functions to support equality comparison.
+    :param order: Whether to generate functions to support ordering.
+    :returns: A data-class type, or a wrapper for data-class types.
+    """
+
+    def wrap(cls: Type[T]) -> Type[T]:
+        # mypy fails to equate bound-y functions (first argument interpreted as
+        # the bound object) with class methods, hence the `ignore` directive.
+        cls.__repr__ = _compact_dataclass_repr  # type: ignore[method-assign]
+        if not dataclasses.is_dataclass(cls):
+            cls = dataclasses.dataclass(  # type: ignore[call-overload]
+                cls,
+                init=True,
+                repr=False,
+                eq=eq,
+                order=order,
+                unsafe_hash=False,
+                frozen=True,
+            )
+        return cls
+
+    # see if decorator is used as @typeannotation or @typeannotation()
+    if cls is None:
+        # called with parentheses
+        return wrap
+    else:
+        # called without parentheses
+        return wrap(cls)
+
+
+@typeannotation
+class Alias:
+    "Alternative name of a property, typically used in JSON serialization."
+
+    name: str
+
+
+@typeannotation
+class Signed:
+    "Signedness of an integer type."
+
+    is_signed: bool
+
+
+@typeannotation
+class Storage:
+    "Number of bytes the binary representation of an integer type takes, e.g. 4 bytes for an int32."
+
+    bytes: int
+
+
+@typeannotation
+class IntegerRange:
+    "Minimum and maximum value of an integer. The range is inclusive."
+
+    minimum: int
+    maximum: int
+
+
+@typeannotation
+class Precision:
+    "Precision of a floating-point value."
+
+    significant_digits: int
+    decimal_digits: int = 0
+
+    @property
+    def integer_digits(self) -> int:
+        return self.significant_digits - self.decimal_digits
+
+
+@typeannotation
+class TimePrecision:
+    """
+    Precision of a timestamp or time interval.
+
+    :param decimal_digits: Number of fractional digits retained in the sub-seconds field for a timestamp.
+    """
+
+    decimal_digits: int = 0
+
+
+@typeannotation
+class Length:
+    "Exact length of a string."
+
+    value: int
+
+
+@typeannotation
+class MinLength:
+    "Minimum length of a string."
+
+    value: int
+
+
+@typeannotation
+class MaxLength:
+    "Maximum length of a string."
+
+    value: int
+
+
+@typeannotation
+class SpecialConversion:
+    "Indicates that the annotated type is subject to custom conversion rules."
+
+
+int8: TypeAlias = Annotated[int, Signed(True), Storage(1), IntegerRange(-128, 127)]
+int16: TypeAlias = Annotated[int, Signed(True), Storage(2), IntegerRange(-32768, 32767)]
+int32: TypeAlias = Annotated[
+    int,
+    Signed(True),
+    Storage(4),
+    IntegerRange(-2147483648, 2147483647),
+]
+int64: TypeAlias = Annotated[
+    int,
+    Signed(True),
+    Storage(8),
+    IntegerRange(-9223372036854775808, 9223372036854775807),
+]
+
+uint8: TypeAlias = Annotated[int, Signed(False), Storage(1), IntegerRange(0, 255)]
+uint16: TypeAlias = Annotated[int, Signed(False), Storage(2), IntegerRange(0, 65535)]
+uint32: TypeAlias = Annotated[
+    int,
+    Signed(False),
+    Storage(4),
+    IntegerRange(0, 4294967295),
+]
+uint64: TypeAlias = Annotated[
+    int,
+    Signed(False),
+    Storage(8),
+    IntegerRange(0, 18446744073709551615),
+]
+
+float32: TypeAlias = Annotated[float, Storage(4)]
+float64: TypeAlias = Annotated[float, Storage(8)]
+
+# maps globals of type Annotated[T, ...] defined in this module to their string names
+_auxiliary_types: Dict[object, str] = {}
+module = sys.modules[__name__]
+for var in dir(module):
+    typ = getattr(module, var)
+    if getattr(typ, "__metadata__", None) is not None:
+        # type is Annotated[T, ...]
+        _auxiliary_types[typ] = var
+
+
+def get_auxiliary_format(data_type: object) -> Optional[str]:
+    "Returns the JSON format string corresponding to an auxiliary type."
+
+    return _auxiliary_types.get(data_type)