arize-phoenix 10.15.0__py3-none-any.whl → 11.0.0__py3-none-any.whl

phoenix/db/facilitator.py

@@ -1,13 +1,20 @@
 from __future__ import annotations
 
 import asyncio
+import json
 import logging
+import re
 import secrets
 from asyncio import gather
+from datetime import datetime, timedelta, timezone
 from functools import partial
-from typing import Optional
+from pathlib import Path
+from typing import NamedTuple, Optional, Union
 
 import sqlalchemy as sa
+from sqlalchemy import select
+from sqlalchemy.orm import InstrumentedAttribute, joinedload
+from sqlalchemy.sql.dml import ReturningDelete
 
 from phoenix import config
 from phoenix.auth import (
@@ -26,7 +33,6 @@ from phoenix.config import (
 from phoenix.db import models
 from phoenix.db.constants import DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID
 from phoenix.db.enums import ENUM_COLUMNS
-from phoenix.db.models import UserRoleName
 from phoenix.db.types.trace_retention import (
     MaxDaysRule,
     TraceRetentionCronExpression,
@@ -62,6 +68,8 @@ class Facilitator:
             _get_system_user_id,
             partial(_ensure_admins, email_sender=self._email_sender),
             _ensure_default_project_trace_retention_policy,
+            _ensure_model_costs,
+            _delete_expired_childless_records,
         ):
             await fn(self._db)
 
@@ -92,13 +100,13 @@ async def _ensure_user_roles(db: DbSessionFactory) -> None:
     the email "admin@localhost".
     """
     async with db() as session:
-        role_ids: dict[UserRoleName, int] = {
+        role_ids: dict[models.UserRoleName, int] = {
             name: id_
             async for name, id_ in await session.stream(
                 sa.select(models.UserRole.name, models.UserRole.id)
             )
         }
-        existing_roles: list[UserRoleName] = [
+        existing_roles: list[models.UserRoleName] = [
             name
             async for name in await session.stream_scalars(
                 sa.select(sa.distinct(models.UserRole.name)).join_from(models.User, models.UserRole)
@@ -216,6 +224,83 @@ async def _ensure_admins(
             logger.error(f"Failed to send welcome email: {exc}")
 
 
+_CHILDLESS_RECORD_DELETION_GRACE_PERIOD_DAYS = 1
+
+
+def _stmt_to_delete_expired_childless_records(
+    table: type[models.Base],
+    foreign_key: Union[InstrumentedAttribute[int], InstrumentedAttribute[Optional[int]]],
+) -> ReturningDelete[tuple[int]]:
+    """
+    Creates a SQLAlchemy DELETE statement to permanently remove childless records.
+
+    Args:
+        table: The table model class that has a deleted_at column
+        foreign_key: The foreign key attribute to check for child relationships
+
+    Returns:
+        A DELETE statement that removes childless records marked for deletion more than
+        _CHILDLESS_RECORD_DELETION_GRACE_PERIOD_DAYS days ago
+    """  # noqa: E501
+    if not hasattr(table, "deleted_at"):
+        raise TypeError("Table must have a 'deleted_at' column")
+    cutoff_time = datetime.now(timezone.utc) - timedelta(
+        days=_CHILDLESS_RECORD_DELETION_GRACE_PERIOD_DAYS
+    )
+    return (
+        sa.delete(table)
+        .where(table.deleted_at.isnot(None))
+        .where(table.deleted_at < cutoff_time)
+        .where(~sa.exists().where(table.id == foreign_key))
+        .returning(table.id)
+    )
+
+
+async def _delete_expired_childless_records_on_generative_models(
+    db: DbSessionFactory,
+) -> None:
+    """
+    Permanently deletes childless GenerativeModel records that have been marked for deletion.
+
+    This function removes GenerativeModel records that:
+    - Have been marked for deletion (deleted_at is not NULL)
+    - Were marked more than 1 day ago (grace period expired)
+    - Have no associated SpanCost records (childless)
+
+    This cleanup is necessary to remove orphaned records that may have been left behind
+    due to previous migrations or deletions.
+    """  # noqa: E501
+    stmt = _stmt_to_delete_expired_childless_records(
+        models.GenerativeModel,
+        models.SpanCost.model_id,
+    )
+    async with db() as session:
+        result = (await session.scalars(stmt)).all()
+    if result:
+        logger.info(f"Permanently deleted {len(result)} expired childless GenerativeModel records")
+    else:
+        logger.debug("No expired childless GenerativeModel records found for permanent deletion")
+
+
+async def _delete_expired_childless_records(
+    db: DbSessionFactory,
+) -> None:
+    """
+    Permanently deletes childless records across all relevant tables.
+
+    This function runs the deletion process for all table types that support soft deletion,
+    handling any exceptions that occur during the process. Only records that have been
+    marked for deletion for more than the grace period (1 day) are permanently removed.
+    """  # noqa: E501
+    exceptions = await gather(
+        _delete_expired_childless_records_on_generative_models(db),
+        return_exceptions=True,
+    )
+    for exc in exceptions:
+        if isinstance(exc, Exception):
+            logger.error(f"Failed to delete childless records: {exc}")
+
+
 async def _ensure_default_project_trace_retention_policy(db: DbSessionFactory) -> None:
     """
     Ensures the default trace retention policy (id=1) exists in the database. Default policy
@@ -261,3 +346,177 @@ async def _ensure_default_project_trace_retention_policy(db: DbSessionFactory) -
                 }
             ],
         )
+
+
+_COST_MODEL_MANIFEST: Path = (
+    Path(__file__).parent.parent / "server" / "cost_tracking" / "model_cost_manifest.json"
+)
+
+
+class _TokenTypeKey(NamedTuple):
+    """
+    Composite key for uniquely identifying token price configurations.
+
+    Token prices are differentiated by both their type (e.g., "input", "output", "audio")
+    and whether they represent prompt tokens (input to the model) or completion tokens
+    (output from the model). Some token types like "audio" can exist in both categories.
+
+    Attributes:
+        token_type: The category of token (e.g., "input", "output", "audio", "cache_write")
+        is_prompt: True if these are prompt/input tokens, False if completion/output tokens
+    """
+
+    token_type: str
+    is_prompt: bool
+
+
+async def _ensure_model_costs(db: DbSessionFactory) -> None:
+    """
+    Ensures that built-in generative models and their token pricing information are up-to-date
+    in the database based on the model cost manifest file.
+
+    This function performs a comprehensive synchronization between the database and the manifest:
+
+    1. **Model Management**: Creates new built-in models from the manifest or updates existing ones
+    2. **Token Price Synchronization**: Ensures all token prices match the manifest rates
+    3. **Cleanup**: Soft-deletes built-in models no longer present in the manifest
+
+    The function handles different token types including:
+    - Input tokens (prompt): Standard input tokens for generation
+    - Cache write tokens (prompt): Tokens written to cache systems
+    - Cache read tokens (prompt): Tokens read from cache systems
+    - Output tokens (non-prompt): Generated response tokens
+    - Audio tokens (both prompt and non-prompt): Audio processing tokens
+
+    Token prices are uniquely identified by (token_type, is_prompt) pairs to handle
+    cases like audio tokens that can be both prompt and non-prompt.
+
+    Args:
+        db (DbSessionFactory): Database session factory for database operations
+
+    Returns:
+        None
+
+    Raises:
+        FileNotFoundError: If the model cost manifest file is not found
+        json.JSONDecodeError: If the manifest file contains invalid JSON
+        ValueError: If manifest data is malformed or missing required fields
+    """
+    # Load the authoritative model cost data from the manifest file
+    with open(_COST_MODEL_MANIFEST) as f:
+        manifest = json.load(f)
+
+    # Define all supported token types with their prompt/non-prompt classification
+    # This determines how tokens are categorized for billing purposes
+    token_types: list[_TokenTypeKey] = [
+        _TokenTypeKey("input", True),  # Standard input tokens
+        _TokenTypeKey("cache_write", True),  # Tokens written to cache
+        _TokenTypeKey("cache_read", True),  # Tokens read from cache
+        _TokenTypeKey("output", False),  # Generated output tokens
+        _TokenTypeKey("audio", True),  # Audio input tokens
+        _TokenTypeKey("audio", False),  # Audio output tokens
+    ]
+
+    async with db() as session:
+        # Fetch all existing built-in models with their token prices eagerly loaded
+        # Using .unique() to deduplicate models when multiple token prices are joined
+        built_in_models = {
+            omodel.name: omodel
+            for omodel in (
+                await session.scalars(
+                    select(models.GenerativeModel)
+                    .where(models.GenerativeModel.deleted_at.is_(None))
+                    .where(models.GenerativeModel.is_built_in.is_(True))
+                    .options(joinedload(models.GenerativeModel.token_prices))
+                )
+            ).unique()
+        }
+
+        seen_names: set[str] = set()
+        seen_patterns: set[tuple[re.Pattern[str], str]] = set()
+
+        # Process each model in the manifest
+        for model_data in manifest:
+            name = str(model_data.get("model") or "").strip()
+            if not name:
+                logger.warning("Skipping model with empty name in manifest")
+                continue
+            if name in seen_names:
+                logger.warning(f"Skipping model '{name}' with duplicate name in manifest")
+                continue
+            seen_names.add(name)
+            regex = str(model_data.get("regex") or "").strip()
+            try:
+                pattern = re.compile(regex)
+            except re.error as e:
+                logger.warning(f"Skipping model '{name}' with invalid regex: {e}")
+                continue
+            provider = str(model_data.get("provider") or "").strip()
+            if (pattern, provider) in seen_patterns:
+                logger.warning(
+                    f"Skipping model '{name}' with duplicate name_pattern/provider combination"
+                )
+                continue
+            seen_patterns.add((pattern, provider))
+            # Remove model from built_in_models dict (for cleanup tracking)
+            # or create new model if not found
+            model = built_in_models.pop(model_data["model"], None)
+            if model is None:
+                # Create new built-in model from manifest data
+                model = models.GenerativeModel(
+                    name=name,
+                    provider=provider,
+                    name_pattern=pattern,
+                    is_built_in=True,
+                )
+                session.add(model)
+            else:
+                # Update existing model's metadata from manifest
+                model.provider = provider
+                model.name_pattern = pattern
+
+            # Create lookup table for existing token prices by (token_type, is_prompt)
+            # Using pop() during iteration allows us to track which prices are no longer needed
+            token_prices = {
+                _TokenTypeKey(token_price.token_type, token_price.is_prompt): token_price
+                for token_price in model.token_prices
+            }
+
+            # Synchronize token prices for all supported token types
+            for token_type, is_prompt in token_types:
+                # Skip if this token type has no rate in the manifest
+                if not (base_rate := model_data.get(token_type)):
+                    continue
+
+                key = _TokenTypeKey(token_type, is_prompt)
+                # Remove from tracking dict and get existing price (if any)
+                if not (token_price := token_prices.pop(key, None)):
+                    # Create new token price if it doesn't exist
+                    token_price = models.TokenPrice(
+                        token_type=token_type,
+                        is_prompt=is_prompt,
+                        base_rate=base_rate,
+                    )
+                    model.token_prices.append(token_price)
+                elif token_price.base_rate != base_rate:
+                    # Update existing price if rate has changed
+                    token_price.base_rate = base_rate
+
+            # Remove any token prices that are no longer in the manifest
+            # These are prices that weren't popped from the token_prices dict above
+            for token_price in token_prices.values():
+                model.token_prices.remove(token_price)
+
+    # Clean up built-in models that are no longer in the manifest
+    # These are models that weren't popped from built_in_models dict above
+    remaining_models = list(built_in_models.values())
+    if not remaining_models:
+        return
+
+    # Soft delete obsolete built-in models
+    async with db() as session:
+        await session.execute(
+            sa.update(models.GenerativeModel)
+            .values(deleted_at=sa.func.now())
+            .where(models.GenerativeModel.id.in_([m.id for m in remaining_models]))
+        )

phoenix/db/insertion/helpers.py

@@ -3,6 +3,7 @@ from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Se
 from enum import Enum, auto
 from typing import Any, Optional
 
+from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes
 from sqlalchemy import Insert
 from sqlalchemy.dialects.postgresql import insert as insert_postgresql
 from sqlalchemy.dialects.sqlite import insert as insert_sqlite
@@ -13,6 +14,7 @@ from typing_extensions import TypeAlias, assert_never
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.models import Base
+from phoenix.trace.attributes import get_attribute_value
 
 
 class DataManipulationEvent(ABC):
@@ -97,3 +99,16 @@ def as_kv(obj: models.Base) -> Iterator[tuple[str, Any]]:
             # postgresql disallows None for primary key
             continue
         yield k, v
+
+
+def should_calculate_span_cost(
+    attributes: Optional[Mapping[str, Any]],
+) -> bool:
+    return bool(
+        (span_kind := get_attribute_value(attributes, SpanAttributes.OPENINFERENCE_SPAN_KIND))
+        and isinstance(span_kind, str)
+        and span_kind == OpenInferenceSpanKindValues.LLM.value
+        and (llm_name := get_attribute_value(attributes, SpanAttributes.LLM_MODEL_NAME))
+        and isinstance(llm_name, str)
+        and llm_name.strip()
+    )

phoenix/db/insertion/span.py

@@ -14,6 +14,8 @@ from phoenix.trace.schemas import Span, SpanStatusCode
 
 class SpanInsertionEvent(NamedTuple):
     project_rowid: int
+    span_rowid: int
+    trace_rowid: int
 
 
 class ClearProjectSpansEvent(NamedTuple):
@@ -190,4 +192,4 @@ async def insert_span(
             + cumulative_llm_token_count_completion,
         )
     )
-    return SpanInsertionEvent(project_rowid)
+    return SpanInsertionEvent(project_rowid, span_rowid, trace.id)

phoenix/db/migrations/versions/a20694b15f82_cost.py

@@ -0,0 +1,196 @@
+"""Cost-related tables
+
+Revision ID: a20694b15f82
+Revises: migrations/versions/6a88424799fe_update_users_with_auth_method.py
+Create Date: 2025-05-30 17:15:12.663565
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "a20694b15f82"
+down_revision: Union[str, None] = "6a88424799fe"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+_Integer = sa.Integer().with_variant(
+    sa.BigInteger(),
+    "postgresql",
+)
+
+
+def upgrade() -> None:
+    op.create_table(
+        "generative_models",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "name",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "provider",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "name_pattern",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "is_built_in",
+            sa.Boolean,
+            nullable=False,
+        ),
+        sa.Column(
+            "start_time",
+            sa.TIMESTAMP(timezone=True),
+        ),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+            onupdate=sa.func.now(),
+        ),
+        sa.Column(
+            "deleted_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=True,
+        ),
+    )
+    op.create_index(
+        "ix_generative_models_match_criteria",
+        "generative_models",
+        ["name_pattern", "provider", "is_built_in"],
+        unique=True,
+        postgresql_where=sa.text("deleted_at IS NULL"),
+        sqlite_where=sa.text("deleted_at IS NULL"),
+    )
+    op.create_index(
+        "ix_generative_models_name_is_built_in",
+        "generative_models",
+        ["name", "is_built_in"],
+        unique=True,
+        postgresql_where=sa.text("deleted_at IS NULL"),
+        sqlite_where=sa.text("deleted_at IS NULL"),
+    )
+    op.create_table(
+        "token_prices",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "model_id",
+            _Integer,
+            sa.ForeignKey("generative_models.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("token_type", sa.String, nullable=False),
+        sa.Column("is_prompt", sa.Boolean, nullable=False),
+        sa.Column("base_rate", sa.Float, nullable=False),
+        sa.Column("customization", sa.JSON),
+        sa.UniqueConstraint(
+            "model_id",
+            "token_type",
+            "is_prompt",
+        ),
+    )
+    op.create_table(
+        "span_costs",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "span_rowid",
+            _Integer,
+            sa.ForeignKey("spans.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "trace_rowid",
+            _Integer,
+            sa.ForeignKey("traces.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "model_id",
+            _Integer,
+            sa.ForeignKey(
+                "generative_models.id",
+                ondelete="RESTRICT",
+            ),
+            nullable=True,
+        ),
+        sa.Column(
+            "span_start_time",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("total_cost", sa.Float),
+        sa.Column("total_tokens", sa.Float),
+        sa.Column("prompt_cost", sa.Float),
+        sa.Column("prompt_tokens", sa.Float),
+        sa.Column("completion_cost", sa.Float),
+        sa.Column("completion_tokens", sa.Float),
+    )
+    op.create_index(
+        "ix_span_costs_model_id_span_start_time",
+        "span_costs",
+        ["model_id", "span_start_time"],
+    )
+    op.create_table(
+        "span_cost_details",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "span_cost_id",
+            _Integer,
+            sa.ForeignKey("span_costs.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("token_type", sa.String, nullable=False, index=True),
+        sa.Column("is_prompt", sa.Boolean, nullable=False),
+        sa.Column("cost", sa.Float),
+        sa.Column("tokens", sa.Float),
+        sa.Column("cost_per_token", sa.Float),
+        sa.UniqueConstraint(
+            "span_cost_id",
+            "token_type",
+            "is_prompt",
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("span_cost_details")
+    op.drop_table("span_costs")
+    op.drop_table("token_prices")
+    op.drop_table("generative_models")