arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/authorization.py

@@ -23,8 +23,8 @@ Usage:
 """
 
 from fastapi import HTTPException, Request
-from fastapi import status as fastapi_status
 
+from phoenix.config import get_env_support_email
 from phoenix.server.bearer_auth import PhoenixUser
 
 
@@ -42,12 +42,47 @@ def require_admin(request: Request) -> None:
     Behavior:
         - Allows access if the authenticated user is an admin or a system user.
         - Raises HTTP 403 Forbidden if the user is not authorized.
-        - Expects authentication to be enabled and request.user to be set by the authentication.
+        - Allows access if authentication is not enabled.
     """
+    if not request.app.state.authentication_enabled:
+        return
     user = getattr(request, "user", None)
     # System users have all privileges
     if not (isinstance(user, PhoenixUser) and user.is_admin):
         raise HTTPException(
-            status_code=fastapi_status.HTTP_403_FORBIDDEN,
+            status_code=403,
             detail="Only admin or system users can perform this action.",
         )
+
+
+def is_not_locked(request: Request) -> None:
+    """
+    FastAPI dependency to ensure database operations are not locked due to insufficient storage.
+
+    This dependency checks if data insertion and update operations are disabled due to
+    storage capacity limits. When storage thresholds are exceeded, it raises an HTTP 507
+    error with actionable guidance for users.
+
+    Usage:
+        Add as a dependency to any route that modifies data:
+
+            @router.post("/create-data", dependencies=[Depends(is_not_locked)])
+            async def create_data(...):
+                ...
+
+    Raises:
+        HTTPException: HTTP 507 Insufficient Storage when database operations are locked.
+            The error includes guidance on resolving storage issues and support contact
+            information if configured.
+    """
+    if request.app.state.db.should_not_insert_or_update:
+        detail = (
+            "Database operations are disabled due to insufficient storage. "
+            "Please delete old data or increase storage."
+        )
+        if support_email := get_env_support_email():
+            detail += f" Need help? Contact us at {support_email}"
+        raise HTTPException(
+            status_code=507,
+            detail=detail,
+        )

phoenix/server/bearer_auth.py

@@ -7,10 +7,9 @@ from typing import Any, Optional, cast
 import grpc
 from fastapi import HTTPException, Request, WebSocket, WebSocketException
 from grpc_interceptor import AsyncServerInterceptor
-from grpc_interceptor.exceptions import Unauthenticated
 from starlette.authentication import AuthCredentials, AuthenticationBackend, BaseUser
 from starlette.requests import HTTPConnection
-from starlette.status import HTTP_401_UNAUTHORIZED
+from typing_extensions import override
 
 from phoenix import config
 from phoenix.auth import (
@@ -20,7 +19,7 @@ from phoenix.auth import (
     Token,
 )
 from phoenix.config import get_env_phoenix_admin_secret
-from phoenix.db import enums, models
+from phoenix.db import models
 from phoenix.server.types import (
     AccessToken,
     AccessTokenAttributes,
@@ -74,14 +73,20 @@ class PhoenixUser(BaseUser):
         self.claims = claims
         assert claims.attributes
         self._is_admin = (
-            claims.status is ClaimSetStatus.VALID
-            and claims.attributes.user_role == enums.UserRole.ADMIN
+            claims.status is ClaimSetStatus.VALID and claims.attributes.user_role == "ADMIN"
+        )
+        self._is_viewer = (
+            claims.status is ClaimSetStatus.VALID and claims.attributes.user_role == "VIEWER"
         )
 
     @cached_property
     def is_admin(self) -> bool:
         return self._is_admin
 
+    @cached_property
+    def is_viewer(self) -> bool:
+        return self._is_viewer
+
     @cached_property
     def identity(self) -> UserId:
         return self._user_id
@@ -94,6 +99,8 @@ class PhoenixUser(BaseUser):
 class PhoenixSystemUser(PhoenixUser):
     def __init__(self, user_id: UserId) -> None:
         self._user_id = user_id
+        self._is_admin = True  # System users have admin privileges
+        self._is_viewer = False  # System users are not viewers
 
     @property
     def is_admin(self) -> bool:
@@ -101,16 +108,19 @@ class PhoenixSystemUser(PhoenixUser):
 
 
 class ApiKeyInterceptor(HasTokenStore, AsyncServerInterceptor):
+    @override
     async def intercept(
         self,
-        method: Callable[[Any, grpc.ServicerContext], Awaitable[Any]],
+        method: Callable[[Any, grpc.aio.ServicerContext], Awaitable[Any]],
         request_or_iterator: Any,
-        context: grpc.ServicerContext,
+        context: grpc.aio.ServicerContext,
         method_name: str,
     ) -> Any:
-        for datum in context.invocation_metadata():
-            if datum.key.lower() == "authorization":
-                scheme, _, token = datum.value.partition(" ")
+        for key, value in context.invocation_metadata() or ():
+            if key.lower() == "authorization":
+                if isinstance(value, bytes):
+                    value = value.decode("utf-8")
+                scheme, _, token = value.partition(" ")
                 if scheme.lower() != "bearer" or not token:
                     break
                 if (
@@ -120,16 +130,16 @@ class ApiKeyInterceptor(HasTokenStore, AsyncServerInterceptor):
                 ):
                     return await method(request_or_iterator, context)
                 claims = await self._token_store.read(Token(token))
-                if not (isinstance(claims, UserClaimSet) and isinstance(claims.subject, UserId)):
+                if (
+                    not (
+                        isinstance(claims, (ApiKeyClaims, AccessTokenClaims))
+                        and isinstance(claims.subject, UserId)
+                    )
+                    or claims.status is not ClaimSetStatus.VALID
+                ):
                     break
-                if not isinstance(claims, (ApiKeyClaims, AccessTokenClaims)):
-                    raise Unauthenticated(details="Invalid token")
-                if claims.status is ClaimSetStatus.EXPIRED:
-                    raise Unauthenticated(details="Expired token")
-                if claims.status is ClaimSetStatus.VALID:
-                    return await method(request_or_iterator, context)
-                raise Unauthenticated()
-        raise Unauthenticated()
+                return await method(request_or_iterator, context)
+        await context.abort(grpc.StatusCode.UNAUTHENTICATED)
 
 
 async def is_authenticated(
@@ -142,16 +152,16 @@ async def is_authenticated(
     """
     assert request or websocket
     if request and not isinstance((user := request.user), PhoenixUser):
-        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+        raise HTTPException(status_code=401, detail="Invalid token")
     if websocket and not isinstance((user := websocket.user), PhoenixUser):
-        raise WebSocketException(code=HTTP_401_UNAUTHORIZED, reason="Invalid token")
+        raise WebSocketException(code=401, reason="Invalid token")
     if isinstance(user, PhoenixSystemUser):
         return
     claims = user.claims
     if claims.status is ClaimSetStatus.EXPIRED:
-        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Expired token")
+        raise HTTPException(status_code=401, detail="Expired token")
     if claims.status is not ClaimSetStatus.VALID:
-        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+        raise HTTPException(status_code=401, detail="Invalid token")
 
 
 async def create_access_and_refresh_tokens(
@@ -163,7 +173,7 @@ async def create_access_and_refresh_tokens(
 ) -> tuple[AccessToken, RefreshToken]:
     issued_at = datetime.now(timezone.utc)
     user_id = UserId(user.id)
-    user_role = enums.UserRole(user.role.name)
+    user_role = user.role.name
     refresh_token_claims = RefreshTokenClaims(
         subject=user_id,
         issued_at=issued_at,

phoenix/server/cost_tracking/cost_details_calculator.py

@@ -0,0 +1,196 @@
+from itertools import chain
+from typing import Any, Iterable, Mapping, Optional
+
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.cost_tracking.helpers import get_aggregated_tokens
+from phoenix.server.cost_tracking.token_cost_calculator import (
+    TokenCostCalculator,
+    create_token_cost_calculator,
+)
+from phoenix.trace.attributes import get_attribute_value
+
+_TokenType: TypeAlias = str
+
+
+class SpanCostDetailsCalculator:
+    """
+    Calculates detailed cost breakdowns for LLM spans based on token usage and pricing.
+
+    This calculator processes both detailed token counts (from span attributes) and
+    aggregated token totals to provide comprehensive cost analysis for prompt and
+    completion tokens. It handles multiple token types (e.g., "input", "output",
+    "image", "audio", "video", "document", "reasoning", etc.) and calculates costs
+    using configured pricing models with fallback behavior.
+
+    **Fallback Behavior:**
+    - If a specific token type has a configured calculator, it uses that calculator
+    - If no specific calculator exists, it falls back to the default calculator:
+      - Prompt tokens (is_prompt=True) fall back to "input" calculator
+      - Completion tokens (is_prompt=False) fall back to "output" calculator
+
+    This ensures all token types get cost calculations even if not explicitly configured.
+
+    The calculator expects token prices to include at least:
+    - An "input" token type for prompt tokens (used as fallback for unconfigured prompt token types)
+    - An "output" token type for completion tokens (used as fallback for unconfigured completion token types)
+
+    Additional token types can be configured for more granular cost tracking.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        prices: Iterable[models.TokenPrice],
+    ) -> None:
+        """
+        Initialize the cost calculator with token pricing configuration.
+
+        Args:
+            prices: Collection of token price configurations defining rates for
+                   different token types and whether they're prompt or completion tokens.
+
+        Raises:
+            ValueError: If required "input" (prompt) or "output" (completion)
+                       token types are missing from the pricing configuration.
+        """
+        # Create calculators for prompt token types (is_prompt=True)
+        self._prompt: Mapping[_TokenType, TokenCostCalculator] = {
+            p.token_type: create_token_cost_calculator(p.base_rate, p.customization)
+            for p in prices
+            if p.is_prompt
+        }
+        if self._prompt and "input" not in self._prompt:
+            raise ValueError("Token prices for prompt must include an 'input' token type")
+
+        # Create calculators for completion token types (is_prompt=False)
+        self._completion: Mapping[_TokenType, TokenCostCalculator] = {
+            p.token_type: create_token_cost_calculator(p.base_rate, p.customization)
+            for p in prices
+            if not p.is_prompt
+        }
+        if self._completion and "output" not in self._completion:
+            raise ValueError("Token prices for completion must include an 'output' token type")
+
+    def calculate_details(
+        self,
+        attributes: Mapping[str, Any],
+    ) -> list[models.SpanCostDetail]:
+        """
+        Calculate detailed cost breakdown for a given span.
+
+        This method processes token usage in two phases:
+        1. **Detailed token processing**: Extracts specific token counts from span attributes
+           (e.g., "llm.token_count.prompt_details", "llm.token_count.completion_details")
+           and calculates costs for each token type found. Uses fallback behavior for
+           token types without specific calculators.
+
+        2. **Aggregated token processing**: For default token types ("input"/"output") that
+           weren't found in detailed processing, calculates remaining tokens by subtracting
+           detailed counts from total aggregated tokens.
+
+        **Fallback Calculation Logic:**
+        - For each token type in detailed processing:
+          - If a specific calculator exists for the token type, use it
+          - Otherwise, fall back to the default calculator ("input" for prompt tokens,
+            "output" for completion tokens)
+        - This ensures all token types receive cost calculations regardless of
+          specific calculator configuration
+
+        Args:
+            attributes: Dictionary containing span attributes with token usage data.
+
+        Returns:
+            List of SpanCostDetail objects containing token counts, costs, and cost-per-token
+            for each token type found in the span.
+
+        Note:
+            - Token counts are validated and converted to non-negative integers
+            - All token types receive cost calculations via fallback mechanism
+            - Cost-per-token is calculated only when both cost and token count are positive
+            - If cost is 0.0, cost-per-token will be None (not 0.0) due to falsy evaluation
+        """
+        prompt_details: dict[_TokenType, models.SpanCostDetail] = {}
+        completion_details: dict[_TokenType, models.SpanCostDetail] = {}
+        calculator: Optional[TokenCostCalculator]
+        cost: Optional[float]
+        cost_per_token: Optional[float]
+
+        # Phase 1: Process detailed token counts from span attributes
+        for is_prompt, prefix, calculators, results in (
+            (True, "prompt", self._prompt, prompt_details),
+            (False, "completion", self._completion, completion_details),
+        ):
+            # Extract detailed token counts from span attributes
+            details = get_attribute_value(attributes, f"llm.token_count.{prefix}_details")
+            if isinstance(details, dict) and details:
+                for token_type, token_count in details.items():
+                    # Validate token count is numeric
+                    if not isinstance(token_count, (int, float)):
+                        continue
+                    tokens = max(0, int(token_count))
+
+                    # Calculate cost using specific calculator or fallback to default
+                    calculator = None
+                    calculator_key = "input" if is_prompt else "output"
+                    if token_type in calculators:
+                        # Use specific calculator for this token type
+                        calculator = calculators[token_type]
+                    elif calculator_key in calculators:
+                        calculator = calculators[calculator_key]
+
+                    cost = None
+                    cost_per_token = None
+                    if calculator:
+                        cost = calculator.calculate_cost(attributes, tokens)
+                        cost_per_token = cost / tokens if tokens else None
+
+                    detail = models.SpanCostDetail(
+                        token_type=token_type,
+                        is_prompt=is_prompt,
+                        tokens=tokens,
+                        cost=cost,
+                        cost_per_token=cost_per_token,
+                    )
+                    results[token_type] = detail
+
+        # Get aggregated token totals for fallback calculations
+        prompt_tokens, completion_tokens, _ = get_aggregated_tokens(attributes)
+
+        # Phase 2: Process remaining tokens for default token types
+        for is_prompt, token_type, total, calculators, results in (
+            (True, "input", prompt_tokens, self._prompt, prompt_details),
+            (False, "output", completion_tokens, self._completion, completion_details),
+        ):
+            # Skip if this token type was already processed in detailed phase
+            if token_type in results:
+                continue
+
+            # Calculate remaining tokens by subtracting detailed counts from total
+            tokens = total - sum(
+                int(d.tokens or 0) for d in results.values() if d.is_prompt == is_prompt
+            )
+
+            # Skip if no remaining tokens or negative (shouldn't happen with valid data)
+            if tokens <= 0:
+                continue
+
+            # Calculate cost using calculator if available
+            cost = None
+            cost_per_token = None
+            if token_type in calculators:
+                calculator = calculators[token_type]
+                cost = calculator.calculate_cost(attributes, tokens)
+                cost_per_token = cost / tokens if tokens else None
+
+            detail = models.SpanCostDetail(
+                token_type=token_type,
+                is_prompt=is_prompt,
+                tokens=tokens,
+                cost=cost,
+                cost_per_token=cost_per_token,
+            )
+            results[token_type] = detail
+
+        # Return combined results from both prompt and completion processing
+        return list(chain(prompt_details.values(), completion_details.values()))

phoenix/server/cost_tracking/cost_model_lookup.py

@@ -0,0 +1,179 @@
+from datetime import datetime
+from typing import Any, Iterable, Mapping, Optional
+
+from openinference.semconv.trace import SpanAttributes
+from typing_extensions import TypeAlias
+
+from phoenix.datetime_utils import is_timezone_aware
+from phoenix.db import models
+from phoenix.server.cost_tracking import regex_specificity
+from phoenix.trace.attributes import get_attribute_value
+
+_RegexPatternStr: TypeAlias = str
+_RegexSpecificityScore: TypeAlias = int
+_TieBreakerId: TypeAlias = int
+
+
+class CostModelLookup:
+    def __init__(
+        self,
+        generative_models: Iterable[models.GenerativeModel] = (),
+    ) -> None:
+        self._models_by_id: dict[int, models.GenerativeModel] = {}
+        self._model_priority: dict[
+            int, tuple[_RegexSpecificityScore, float, _TieBreakerId]
+        ] = {}  # higher is better
+
+        for m in generative_models:
+            self._add_or_update_model(m)
+
+    def _add_or_update_model(self, model: models.GenerativeModel) -> None:
+        """Add or update a single model in the lookup."""
+        self._models_by_id[model.id] = model
+
+        specificity_score = regex_specificity.score(model.name_pattern)
+
+        # For built-in models, use negative ID so that earlier IDs win
+        # For user-defined models, use positive ID so later IDs win
+        tie_breaker = -model.id if model.is_built_in else model.id
+
+        self._model_priority[model.id] = (
+            specificity_score,
+            model.start_time.timestamp() if model.start_time else 0.0,
+            tie_breaker,
+        )
+
+    def _remove_model(self, model_id: int) -> None:
+        """Remove a model from the lookup."""
+        if model_id in self._models_by_id:
+            del self._models_by_id[model_id]
+        if model_id in self._model_priority:
+            del self._model_priority[model_id]
+
+    def merge(self, models: Iterable[models.GenerativeModel]) -> None:
+        """
+        Merge a collection of models into the existing lookup.
+
+        For each model:
+        - If deleted_at is set, remove it from the lookup
+        - Otherwise, add or update it in the lookup
+
+        Args:
+            models: An iterable of GenerativeModel objects to merge
+        """
+        for model in models:
+            if model.deleted_at is not None:
+                self._remove_model(model.id)
+            else:
+                self._add_or_update_model(model)
+
+    def find_model(
+        self,
+        start_time: datetime,
+        attributes: Mapping[str, Any],
+    ) -> Optional[models.GenerativeModel]:
+        """
+        Find the most appropriate generative model for cost tracking based on attributes and time.
+
+        This method implements a sophisticated model lookup system that filters and prioritizes
+        generative models based on the provided attributes and timestamp. The lookup follows
+        a specific priority hierarchy to ensure consistent and predictable model selection.
+
+        Args:
+            start_time: The timestamp for which to find a model. Must be timezone-aware.
+                Models with start_time greater than this value will be excluded.
+            attributes: A mapping containing span attributes. Must include:
+                - SpanAttributes.LLM_MODEL_NAME: The name of the LLM model to match
+                - SpanAttributes.LLM_PROVIDER: (Optional) The provider of the LLM model
+
+        Raises:
+            TypeError: If start_time is not timezone-aware (tzinfo is None)
+
+        Returns:
+            The most appropriate GenerativeModel that matches the criteria, or None if no
+            suitable model is found.
+
+        Model Selection Logic:
+            1. **Input Validation**: Returns None if model name is empty or whitespace-only
+            2. **Time and Regex Filtering**: Only models that satisfy both conditions:
+               - start_time <= start_time or start_time=None (active models)
+               - name_pattern regex matches the model name from attributes
+            3. **Early Return Optimization**: If only one candidate remains, return it immediately
+            4. **Two-Tier Priority System**: Models are processed in tiers:
+               - User-defined models (is_built_in=False) are processed first
+               - Built-in models (is_built_in=True) are processed second
+               - If a tier has only one model, return it immediately
+            5. **Provider Filtering**: Within each tier, if provider is specified:
+               - Prefer models with matching provider
+               - Fall back to provider-agnostic models if no provider-specific matches exist
+            6. **Priority Selection**: Select the model with the highest priority tuple:
+               (regex_specificity_score, start_time.timestamp, tie_breaker)
+
+        Priority Tuple Components:
+            - regex_specificity_score: More specific regex patterns have higher priority
+            - start_time.timestamp: Models with later start times have higher priority
+            - tie_breaker: For built-in models, uses negative ID (lower IDs win);
+              for user-defined models, uses positive ID (higher IDs win)
+
+        Examples:
+            >>> lookup = CostModelLookup([model1, model2, model3])
+            >>> model = lookup.find_model(
+            ...     start_time=datetime(2024, 1, 1, tzinfo=timezone.utc),
+            ...     attributes={"llm": {"model_name": "gpt-3.5-turbo", "provider": "openai"}}
+            ... )
+        """  # noqa: E501
+        # 1. extract and validate inputs
+        if not is_timezone_aware(start_time):
+            raise TypeError("start_time must be timezone-aware")
+
+        model_name = str(
+            get_attribute_value(attributes, SpanAttributes.LLM_MODEL_NAME) or ""
+        ).strip()
+        if not model_name:
+            return None
+
+        # 2. only include models that are active and match the regex pattern
+        candidates = [
+            model
+            for model in self._models_by_id.values()
+            if (not model.start_time or model.start_time <= start_time)
+            and model.name_pattern.search(model_name)
+        ]
+        if not candidates:
+            return None
+
+        # 3. early return: if only one candidate remains, return it
+        if len(candidates) == 1:
+            return candidates[0]
+
+        provider = str(get_attribute_value(attributes, SpanAttributes.LLM_PROVIDER) or "").strip()
+
+        # 4. priority-based selection: user-defined models first, then built-in models
+        for is_built_in in (False, True):  # False = user-defined, True = built-in
+            # get candidates for current tier (user-defined or built-in)
+            tier_candidates = [model for model in candidates if model.is_built_in == is_built_in]
+
+            if not tier_candidates:
+                continue  # try next tier
+
+            # early return: if only one candidate in this tier, return it
+            if len(tier_candidates) == 1:
+                return tier_candidates[0]
+
+            # 5. provider filtering: if provider specified, prefer provider-specific models
+            if provider:
+                provider_specific_models = [
+                    model
+                    for model in tier_candidates
+                    if model.provider and model.provider == provider
+                ]
+                # only use provider-specific models if any exist
+                # this allows fallback to provider-agnostic models when no match
+                if provider_specific_models:
+                    tier_candidates = provider_specific_models
+
+            # 6. select best model in this tier
+            return max(tier_candidates, key=lambda model: self._model_priority[model.id])
+
+        # 7. no suitable model found
+        return None

phoenix/server/cost_tracking/helpers.py

@@ -0,0 +1,68 @@
+import logging
+from typing import Any, Mapping
+
+from openinference.semconv.trace import SpanAttributes
+from typing_extensions import TypeAlias
+
+from phoenix.trace.attributes import get_attribute_value
+
+logger = logging.getLogger(__name__)
+
+_PromptTokens: TypeAlias = int
+_CompletionTokens: TypeAlias = int
+_TotalTokens: TypeAlias = int
+
+
+def get_aggregated_tokens(
+    attributes: Mapping[str, Any],
+) -> tuple[_PromptTokens, _CompletionTokens, _TotalTokens]:
+    """Return the total, prompt, and completion token counts from the span attributes."""
+    try:
+        prompt_tokens_value = get_attribute_value(
+            attributes,
+            SpanAttributes.LLM_TOKEN_COUNT_PROMPT,
+        )
+        prompt_tokens: int = (
+            0
+            if not isinstance(prompt_tokens_value, (int, float))
+            else max(0, int(prompt_tokens_value))
+        )
+
+        completion_tokens_value = get_attribute_value(
+            attributes,
+            SpanAttributes.LLM_TOKEN_COUNT_COMPLETION,
+        )
+        completion_tokens: int = (
+            0
+            if not isinstance(completion_tokens_value, (int, float))
+            else max(0, int(completion_tokens_value))
+        )
+
+        total_tokens_value = get_attribute_value(
+            attributes,
+            SpanAttributes.LLM_TOKEN_COUNT_TOTAL,
+        )
+        total_tokens: int = (
+            0
+            if not isinstance(total_tokens_value, (int, float))
+            else max(0, int(total_tokens_value))
+        )
+
+        assert prompt_tokens >= 0
+        assert completion_tokens >= 0
+        assert total_tokens >= 0
+
+        calculated_total = prompt_tokens + completion_tokens
+
+        if total_tokens > calculated_total:
+            if not prompt_tokens:
+                prompt_tokens = total_tokens - completion_tokens
+            else:
+                completion_tokens = total_tokens - prompt_tokens
+        else:
+            total_tokens = calculated_total
+
+        return prompt_tokens, completion_tokens, total_tokens
+    except Exception as e:
+        logger.error(f"Error getting aggregated tokens: {e}")
+        return 0, 0, 0