arize-phoenix 11.4.0__py3-none-any.whl → 11.6.0__py3-none-any.whl

phoenix/server/api/routers/v1/prompts.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Optional, Union
 
-from fastapi import APIRouter, HTTPException, Path, Query
+from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from pydantic import ValidationError, model_validator
 from sqlalchemy import select
 from sqlalchemy.sql import Select
@@ -33,6 +33,7 @@ from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Prompt import Prompt as PromptNodeType
 from phoenix.server.api.types.PromptVersion import PromptVersion as PromptVersionNodeType
 from phoenix.server.api.types.PromptVersionTag import PromptVersionTag as PromptVersionTagNodeType
+from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
 
 logger = logging.getLogger(__name__)
@@ -393,6 +394,7 @@ async def get_prompt_version_by_latest(
 
 @router.post(
     "/prompts",
+    dependencies=[Depends(is_not_locked)],
     operation_id="postPromptVersion",
     summary="Create a new prompt",
     description="Create a new prompt and its initial version. A prompt can have multiple versions.",
@@ -602,6 +604,7 @@ async def list_prompt_version_tags(
 
 @router.post(
     "/prompt_versions/{prompt_version_id}/tags",
+    dependencies=[Depends(is_not_locked)],
     operation_id="createPromptVersionTag",
     summary="Add tag to prompt version",
     description="Add a new tag to a specific prompt version. Tags help identify and categorize "

phoenix/server/api/routers/v1/spans.py

@@ -8,7 +8,7 @@ from secrets import token_urlsafe
 from typing import Annotated, Any, Literal, Optional, Union
 
 import pandas as pd
-from fastapi import APIRouter, Header, HTTPException, Path, Query
+from fastapi import APIRouter, Depends, Header, HTTPException, Path, Query
 from pydantic import BaseModel, Field
 from sqlalchemy import select
 from starlette.requests import Request
@@ -28,6 +28,7 @@ from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
 from phoenix.db.insertion.types import Precursors
 from phoenix.server.api.routers.utils import df_to_bytes
+from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import SpanAnnotationInsertEvent
 from phoenix.trace.attributes import flatten
@@ -601,7 +602,7 @@ async def span_search_otlpv1(
             models.Trace.trace_id,
         )
         .join(models.Trace, onclause=models.Trace.id == models.Span.trace_rowid)
-        .join(models.Project, onclause=models.Project.id == project_id)
+        .where(models.Trace.project_rowid == project_id)
         .order_by(*order_by)
     )
 
@@ -736,7 +737,7 @@ async def span_search(
             models.Trace.trace_id,
         )
         .join(models.Trace, onclause=models.Trace.id == models.Span.trace_rowid)
-        .join(models.Project, onclause=models.Project.id == project_id)
+        .where(models.Trace.project_rowid == project_id)
         .order_by(*order_by)
     )
 
@@ -907,6 +908,7 @@ class AnnotateSpansResponseBody(ResponseBody[list[InsertedSpanAnnotation]]):
 
 @router.post(
     "/span_annotations",
+    dependencies=[Depends(is_not_locked)],
     operation_id="annotateSpans",
     summary="Create span annotations",
     responses=add_errors_to_responses(
@@ -990,6 +992,7 @@ class CreateSpansResponseBody(V1RoutesBaseModel):
 
 @router.post(
     "/projects/{project_identifier}/spans",
+    dependencies=[Depends(is_not_locked)],
     operation_id="createSpans",
     summary="Create spans",
     description=(

phoenix/server/api/routers/v1/traces.py

@@ -2,7 +2,7 @@ import gzip
 import zlib
 from typing import Any, Literal, Optional
 
-from fastapi import APIRouter, BackgroundTasks, Header, HTTPException, Query
+from fastapi import APIRouter, BackgroundTasks, Depends, Header, HTTPException, Query
 from google.protobuf.message import DecodeError
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest,
@@ -24,6 +24,7 @@ from strawberry.relay import GlobalID
 from phoenix.db import models
 from phoenix.db.insertion.helpers import as_kv
 from phoenix.db.insertion.types import Precursors
+from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import TraceAnnotationInsertEvent
 from phoenix.trace.otel import decode_otlp_span
@@ -37,6 +38,7 @@ router = APIRouter(tags=["traces"])
 
 @router.post(
     "/traces",
+    dependencies=[Depends(is_not_locked)],
     operation_id="addTraces",
     summary="Send traces",
     responses=add_errors_to_responses(
@@ -160,6 +162,7 @@ class AnnotateTracesResponseBody(ResponseBody[list[InsertedTraceAnnotation]]):
 
 @router.post(
     "/trace_annotations",
+    dependencies=[Depends(is_not_locked)],
     operation_id="annotateTraces",
     summary="Create trace annotations",
     responses=add_errors_to_responses(

phoenix/server/api/routers/v1/users.py

@@ -43,7 +43,7 @@ from phoenix.server.api.routers.v1.utils import (
     add_errors_to_responses,
 )
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.authorization import require_admin
+from phoenix.server.authorization import is_not_locked, require_admin
 
 logger = logging.getLogger(__name__)
 
@@ -194,7 +194,7 @@ async def list_users(
             HTTP_422_UNPROCESSABLE_ENTITY,
         ]
     ),
-    dependencies=[Depends(require_admin)],
+    dependencies=[Depends(require_admin), Depends(is_not_locked)],
     response_model_by_alias=True,
     response_model_exclude_unset=True,
     response_model_exclude_defaults=True,

phoenix/server/api/types/Span.py

@@ -19,7 +19,6 @@ from typing_extensions import Annotated, TypeAlias
 import phoenix.trace.schemas as trace_schema
 from phoenix.db import models
 from phoenix.server.api.context import Context
-from phoenix.server.api.dataloaders import types as dataloader_types
 from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
@@ -829,104 +828,6 @@ class Span(Node):
             for entry in entries
         ]
 
-    @strawberry.field
-    async def cumulative_cost_summary(self, info: Info[Context, None]) -> Optional[SpanCostSummary]:
-        max_depth = 0
-        descendant_rowids = await info.context.data_loaders.span_descendants.load(
-            (self.span_rowid, max_depth)
-        )
-        span_costs = await info.context.data_loaders.span_cost_by_span.load_many(
-            (self.span_rowid, *descendant_rowids)
-        )
-        total_cost: Optional[float] = None
-        total_tokens: Optional[float] = None
-        prompt_cost: Optional[float] = None
-        prompt_tokens: Optional[float] = None
-        completion_cost: Optional[float] = None
-        completion_tokens: Optional[float] = None
-        for span_cost in span_costs:
-            if span_cost is None:
-                continue
-            if span_cost.total_cost is not None:
-                total_cost = (total_cost or 0) + span_cost.total_cost
-            if span_cost.total_tokens is not None:
-                total_tokens = (total_tokens or 0) + span_cost.total_tokens
-            if span_cost.prompt_cost is not None:
-                prompt_cost = (prompt_cost or 0) + span_cost.prompt_cost
-            if span_cost.prompt_tokens is not None:
-                prompt_tokens = (prompt_tokens or 0) + span_cost.prompt_tokens
-            if span_cost.completion_cost is not None:
-                completion_cost = (completion_cost or 0) + span_cost.completion_cost
-            if span_cost.completion_tokens is not None:
-                completion_tokens = (completion_tokens or 0) + span_cost.completion_tokens
-        return SpanCostSummary(
-            prompt=CostBreakdown(
-                tokens=prompt_tokens,
-                cost=prompt_cost,
-            ),
-            completion=CostBreakdown(
-                tokens=completion_tokens,
-                cost=completion_cost,
-            ),
-            total=CostBreakdown(
-                tokens=total_tokens,
-                cost=total_cost,
-            ),
-        )
-
-    @strawberry.field
-    async def cumulative_cost_detail_summary_entries(
-        self, info: Info[Context, None]
-    ) -> list[SpanCostDetailSummaryEntry]:
-        max_depth = 0
-        descendant_rowids = await info.context.data_loaders.span_descendants.load(
-            (self.span_rowid, max_depth)
-        )
-        entry_lists = (
-            await info.context.data_loaders.span_cost_detail_summary_entries_by_span.load_many(
-                (self.span_rowid, *descendant_rowids)
-            )
-        )
-
-        TokenType: TypeAlias = str
-        IsPrompt: TypeAlias = bool
-        grouped_entries: dict[
-            IsPrompt, dict[TokenType, list[dataloader_types.SpanCostDetailSummaryEntry]]
-        ] = {}
-
-        for entries in entry_lists:
-            for entry in entries:
-                is_prompt = entry.is_prompt
-                token_type = entry.token_type
-                if is_prompt not in grouped_entries:
-                    grouped_entries[is_prompt] = {}
-                if token_type not in grouped_entries[is_prompt]:
-                    grouped_entries[is_prompt][token_type] = []
-                grouped_entries[is_prompt][token_type].append(entry)
-
-        result: list[SpanCostDetailSummaryEntry] = []
-        for is_prompt in (True, False):
-            entries_by_token_type = grouped_entries[is_prompt]
-            for token_type, entries in sorted(entries_by_token_type.items()):
-                cost: Optional[float] = None
-                tokens: Optional[float] = None
-                for entry in entries:
-                    if entry.value.cost is not None:
-                        cost = (cost or 0) + entry.value.cost
-                    if entry.value.tokens is not None:
-                        tokens = (tokens or 0) + entry.value.tokens
-                result.append(
-                    SpanCostDetailSummaryEntry(
-                        token_type=token_type,
-                        is_prompt=is_prompt,
-                        value=CostBreakdown(
-                            tokens=tokens,
-                            cost=cost,
-                        ),
-                    )
-                )
-        return result
-
 
 def _hide_embedding_vectors(attributes: Mapping[str, Any]) -> Mapping[str, Any]:
     if not (

phoenix/server/app.py

@@ -16,17 +16,20 @@ from typing import (
     Any,
     NamedTuple,
     Optional,
+    Protocol,
     TypedDict,
     Union,
     cast,
 )
 from urllib.parse import urlparse
 
+import grpc
 import strawberry
 from fastapi import APIRouter, Depends, FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.utils import is_body_allowed_for_status_code
 from grpc.aio import ServerInterceptor
+from grpc_interceptor import AsyncServerInterceptor
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
 from starlette.datastructures import URL, Secret
@@ -44,7 +47,7 @@ from starlette.types import Scope, StatefulLifespan
 from strawberry.extensions import SchemaExtension
 from strawberry.fastapi import GraphQLRouter
 from strawberry.subscriptions import GRAPHQL_TRANSPORT_WS_PROTOCOL
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, override
 
 import phoenix.trace.v1 as pb
 from phoenix.config import (
@@ -134,6 +137,7 @@ from phoenix.server.api.routers import (
 from phoenix.server.api.routers.v1 import REST_API_VERSION
 from phoenix.server.api.schema import build_graphql_schema
 from phoenix.server.bearer_auth import BearerTokenAuthBackend, is_authenticated
+from phoenix.server.daemons.db_disk_usage_monitor import DbDiskUsageMonitor
 from phoenix.server.daemons.generative_model_store import GenerativeModelStore
 from phoenix.server.daemons.span_cost_calculator import SpanCostCalculator
 from phoenix.server.dml_event import DmlEvent
@@ -382,19 +386,16 @@ async def version() -> PlainTextResponse:
     return PlainTextResponse(f"{phoenix_version}")
 
 
-DB_MUTEX: Optional[asyncio.Lock] = None
-
-
 def _db(
-    engine: AsyncEngine, bypass_lock: bool = False
-) -> Callable[[], AbstractAsyncContextManager[AsyncSession]]:
+    engine: AsyncEngine,
+) -> Callable[[Optional[asyncio.Lock]], AbstractAsyncContextManager[AsyncSession]]:
     Session = async_sessionmaker(engine, expire_on_commit=False)
 
     @contextlib.asynccontextmanager
-    async def factory() -> AsyncIterator[AsyncSession]:
+    async def factory(lock: Optional[asyncio.Lock] = None) -> AsyncIterator[AsyncSession]:
         async with contextlib.AsyncExitStack() as stack:
-            if not bypass_lock and DB_MUTEX:
-                await stack.enter_async_context(DB_MUTEX)
+            if lock:
+                await stack.enter_async_context(lock)
             yield await stack.enter_async_context(Session.begin())
 
     return factory
@@ -523,6 +524,7 @@ def _lifespan(
     trace_data_sweeper: Optional[TraceDataSweeper],
     span_cost_calculator: SpanCostCalculator,
     generative_model_store: GenerativeModelStore,
+    db_disk_usage_monitor: DbDiskUsageMonitor,
     token_store: Optional[TokenStore] = None,
     tracer_provider: Optional["TracerProvider"] = None,
     enable_prometheus: bool = False,
@@ -530,14 +532,14 @@ def _lifespan(
     shutdown_callbacks: Iterable[_Callback] = (),
     read_only: bool = False,
     scaffolder_config: Optional[ScaffolderConfig] = None,
+    grpc_interceptors: Iterable[AsyncServerInterceptor] = (),
 ) -> StatefulLifespan[FastAPI]:
     @contextlib.asynccontextmanager
     async def lifespan(_: FastAPI) -> AsyncIterator[dict[str, Any]]:
         for callback in startup_callbacks:
             if isinstance((res := callback()), Awaitable):
                 await res
-        global DB_MUTEX
-        DB_MUTEX = asyncio.Lock() if db.dialect is SupportedSQLDialect.SQLITE else None
+        db.lock = asyncio.Lock() if db.dialect is SupportedSQLDialect.SQLITE else None
         async with AsyncExitStack() as stack:
             (
                 enqueue,
@@ -551,7 +553,7 @@ def _lifespan(
                 tracer_provider=tracer_provider,
                 enable_prometheus=enable_prometheus,
                 token_store=token_store,
-                interceptors=user_grpc_interceptors(),
+                interceptors=user_grpc_interceptors() + list(grpc_interceptors),
             )
             await stack.enter_async_context(grpc_server)
             await stack.enter_async_context(dml_event_handler)
@@ -559,6 +561,7 @@ def _lifespan(
                 await stack.enter_async_context(trace_data_sweeper)
             await stack.enter_async_context(span_cost_calculator)
             await stack.enter_async_context(generative_model_store)
+            await stack.enter_async_context(db_disk_usage_monitor)
             if scaffolder_config:
                 scaffolder = Scaffolder(
                     config=scaffolder_config,
@@ -826,6 +829,34 @@ async def plain_text_http_exception_handler(request: Request, exc: HTTPException
     return PlainTextResponse(str(exc.detail), status_code=exc.status_code, headers=headers)
 
 
+class _HasDbStatus(Protocol):
+    @property
+    def should_not_insert_or_update(self) -> bool: ...
+
+
+class DbDiskUsageInterceptor(AsyncServerInterceptor):
+    def __init__(self, db: _HasDbStatus) -> None:
+        self._db = db
+
+    @override
+    async def intercept(
+        self,
+        method: Callable[[Any, grpc.aio.ServicerContext], Awaitable[Any]],
+        request_or_iterator: Any,
+        context: grpc.aio.ServicerContext,
+        method_name: str,
+    ) -> Any:
+        if (
+            method_name.endswith("trace.v1.TraceService/Export")
+            and self._db.should_not_insert_or_update
+        ):
+            await context.abort(
+                grpc.StatusCode.RESOURCE_EXHAUSTED,
+                "Database disk usage threshold exceeded",
+            )
+        return await method(request_or_iterator, context)
+
+
 def create_app(
     db: DbSessionFactory,
     export_path: Path,
@@ -971,6 +1002,8 @@ def create_app(
         from phoenix.server.prometheus import PrometheusMiddleware
 
         middlewares.append(Middleware(PrometheusMiddleware))
+    grpc_interceptors: list[AsyncServerInterceptor] = []
+    grpc_interceptors.append(DbDiskUsageInterceptor(db))
     app = FastAPI(
         title="Arize-Phoenix REST API",
         version=REST_API_VERSION,
@@ -982,6 +1015,8 @@ def create_app(
             trace_data_sweeper=trace_data_sweeper,
             span_cost_calculator=span_cost_calculator,
             generative_model_store=generative_model_store,
+            db_disk_usage_monitor=DbDiskUsageMonitor(db, email_sender),
+            grpc_interceptors=grpc_interceptors,
             token_store=token_store,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,

phoenix/server/authorization.py

@@ -51,3 +51,12 @@ def require_admin(request: Request) -> None:
             status_code=fastapi_status.HTTP_403_FORBIDDEN,
             detail="Only admin or system users can perform this action.",
         )
+
+
+def is_not_locked(request: Request) -> None:
+    if request.app.state.db.should_not_insert_or_update:
+        raise HTTPException(
+            status_code=fastapi_status.HTTP_507_INSUFFICIENT_STORAGE,
+            detail="Operations that insert or update database "
+            "records are currently not allowed.",
+        )

phoenix/server/bearer_auth.py

@@ -7,10 +7,10 @@ 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 (
@@ -100,16 +100,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 (
@@ -119,16 +122,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(

phoenix/server/cost_tracking/cost_model_lookup.py

@@ -109,7 +109,7 @@ class CostModelLookup:
             model
             for model in self._models
             if (not model.start_time or model.start_time <= start_time)
-            and model.name_pattern.match(model_name)
+            and model.name_pattern.search(model_name)
         ]
         if not candidates:
             return None