arize-phoenix 4.14.1__py3-none-any.whl → 4.16.0__py3-none-any.whl

phoenix/db/insertion/types.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from abc import ABC, abstractmethod
+from copy import copy
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import (
+    Any,
+    Generic,
+    List,
+    Mapping,
+    Optional,
+    Protocol,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    cast,
+)
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.dml import ReturningInsert
+
+from phoenix.db import models
+from phoenix.db.insertion.constants import DEFAULT_RETRY_ALLOWANCE, DEFAULT_RETRY_DELAY_SEC
+from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.server.types import DbSessionFactory
+
+logger = logging.getLogger("__name__")
+
+
+class Insertable(Protocol):
+    @property
+    def row(self) -> models.Base: ...
+
+
+_AnyT = TypeVar("_AnyT")
+_PrecursorT = TypeVar("_PrecursorT")
+_InsertableT = TypeVar("_InsertableT", bound=Insertable)
+_RowT = TypeVar("_RowT", bound=models.Base)
+
+
+@dataclass(frozen=True)
+class Received(Generic[_AnyT]):
+    item: _AnyT
+    received_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def postpone(self, retries_left: int = DEFAULT_RETRY_ALLOWANCE) -> Postponed[_AnyT]:
+        return Postponed(item=self.item, received_at=self.received_at, retries_left=retries_left)
+
+
+@dataclass(frozen=True)
+class Postponed(Received[_AnyT]):
+    retries_left: int = field(default=DEFAULT_RETRY_ALLOWANCE)
+
+
+class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
+    table: Type[_RowT]
+    unique_by: Sequence[str]
+
+    def __init_subclass__(
+        cls,
+        table: Type[_RowT],
+        unique_by: Sequence[str],
+    ) -> None:
+        cls.table = table
+        cls.unique_by = unique_by
+
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
+        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
+    ) -> None:
+        self._queue: List[Received[_PrecursorT]] = []
+        self._db = db
+        self._retry_delay_sec = retry_delay_sec
+        self._retry_allowance = retry_allowance
+
+    @property
+    def empty(self) -> bool:
+        return not bool(self._queue)
+
+    async def enqueue(self, *items: _PrecursorT) -> None:
+        self._queue.extend([Received(item) for item in items])
+
+    @abstractmethod
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[_PrecursorT],
+    ) -> Tuple[
+        List[Received[_InsertableT]],
+        List[Postponed[_PrecursorT]],
+        List[Received[_PrecursorT]],
+    ]: ...
+
+    async def insert(self) -> Tuple[Type[_RowT], List[int]]:
+        if not self._queue:
+            return self.table, []
+        parcels = self._queue
+        self._queue = []
+        inserted_ids: List[int] = []
+        async with self._db() as session:
+            to_insert, to_postpone, _ = await self._partition(session, *parcels)
+            if to_insert:
+                inserted_ids, to_retry, _ = await self._insert(session, *to_insert)
+                to_postpone.extend(to_retry)
+        if to_postpone:
+            loop = asyncio.get_running_loop()
+            loop.call_later(self._retry_delay_sec, self._queue.extend, to_postpone)
+        return self.table, inserted_ids
+
+    def _stmt(self, *records: Mapping[str, Any]) -> ReturningInsert[Tuple[int]]:
+        pk = next(c for c in self.table.__table__.c if c.primary_key)
+        return insert_on_conflict(
+            *records,
+            table=self.table,
+            unique_by=self.unique_by,
+            dialect=self._db.dialect,
+        ).returning(pk)
+
+    async def _insert(
+        self,
+        session: AsyncSession,
+        *insertions: Received[_InsertableT],
+    ) -> Tuple[List[int], List[Postponed[_PrecursorT]], List[Received[_InsertableT]]]:
+        records = [dict(as_kv(ins.item.row)) for ins in insertions]
+        inserted_ids: List[int] = []
+        to_retry: List[Postponed[_PrecursorT]] = []
+        failures: List[Received[_InsertableT]] = []
+        stmt = self._stmt(*records)
+        try:
+            async with session.begin_nested():
+                ids = [id_ async for id_ in await session.stream_scalars(stmt)]
+                inserted_ids.extend(ids)
+        except BaseException:
+            logger.exception(
+                f"Failed to bulk insert for {self.table.__name__}. "
+                f"Will try to insert ({len(records)} records) individually instead."
+            )
+            for i, record in enumerate(records):
+                stmt = self._stmt(record)
+                try:
+                    async with session.begin_nested():
+                        ids = [id_ async for id_ in await session.stream_scalars(stmt)]
+                        inserted_ids.extend(ids)
+                except BaseException:
+                    logger.exception(f"Failed to insert for {self.table.__name__}.")
+                    p = insertions[i]
+                    if isinstance(p, Postponed) and p.retries_left == 1:
+                        failures.append(p)
+                    else:
+                        to_retry.append(
+                            Postponed(
+                                item=cast(_PrecursorT, p.item),
+                                received_at=p.received_at,
+                                retries_left=(p.retries_left - 1)
+                                if isinstance(p, Postponed)
+                                else self._retry_allowance,
+                            )
+                        )
+        return inserted_ids, to_retry, failures
+
+
+class Precursors(ABC):
+    @dataclass(frozen=True)
+    class SpanAnnotation:
+        span_id: str
+        obj: models.SpanAnnotation
+
+        def as_insertable(
+            self,
+            span_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.SpanAnnotation:
+            return Insertables.SpanAnnotation(
+                span_id=self.span_id,
+                obj=self.obj,
+                span_rowid=span_rowid,
+                id_=id_,
+            )
+
+    @dataclass(frozen=True)
+    class TraceAnnotation:
+        trace_id: str
+        obj: models.TraceAnnotation
+
+        def as_insertable(
+            self,
+            trace_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.TraceAnnotation:
+            return Insertables.TraceAnnotation(
+                trace_id=self.trace_id,
+                obj=self.obj,
+                trace_rowid=trace_rowid,
+                id_=id_,
+            )
+
+    @dataclass(frozen=True)
+    class DocumentAnnotation:
+        span_id: str
+        document_position: int
+        obj: models.DocumentAnnotation
+
+        def as_insertable(
+            self,
+            span_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.DocumentAnnotation:
+            return Insertables.DocumentAnnotation(
+                span_id=self.span_id,
+                document_position=self.document_position,
+                obj=self.obj,
+                span_rowid=span_rowid,
+                id_=id_,
+            )
+
+
+class Insertables(ABC):
+    @dataclass(frozen=True)
+    class SpanAnnotation(Precursors.SpanAnnotation):
+        span_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.SpanAnnotation:
+            obj = copy(self.obj)
+            obj.span_rowid = self.span_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj
+
+    @dataclass(frozen=True)
+    class TraceAnnotation(Precursors.TraceAnnotation):
+        trace_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.TraceAnnotation:
+            obj = copy(self.obj)
+            obj.trace_rowid = self.trace_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj
+
+    @dataclass(frozen=True)
+    class DocumentAnnotation(Precursors.DocumentAnnotation):
+        span_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.DocumentAnnotation:
+            obj = copy(self.obj)
+            obj.span_rowid = self.span_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj

phoenix/experiments/functions.py

@@ -72,15 +72,16 @@ from phoenix.experiments.types import (
 )
 from phoenix.experiments.utils import get_dataset_experiments_url, get_experiment_url, get_func_name
 from phoenix.trace.attributes import flatten
+from phoenix.utilities.client import VersionedAsyncClient, VersionedClient
 from phoenix.utilities.json import jsonify
 
 
 def _phoenix_clients() -> Tuple[httpx.Client, httpx.AsyncClient]:
     headers = get_env_client_headers()
-    return httpx.Client(
+    return VersionedClient(
         base_url=get_base_url(),
         headers=headers,
-    ), httpx.AsyncClient(
+    ), VersionedAsyncClient(
         base_url=get_base_url(),
         headers=headers,
     )

phoenix/experiments/types.py

@@ -226,7 +226,7 @@ class ExperimentRun:
 
     def __post_init__(self) -> None:
         if bool(self.output) == bool(self.error):
-            ValueError("Must specify exactly one of experiment_run_output or error")
+            raise ValueError("Must specify exactly one of experiment_run_output or error")
 
 
 @dataclass(frozen=True)
@@ -249,7 +249,7 @@ class EvaluationResult:
 
     def __post_init__(self) -> None:
         if self.score is None and not self.label:
-            ValueError("Must specify score or label, or both")
+            raise ValueError("Must specify score or label, or both")
         if self.score is None and not self.label:
             object.__setattr__(self, "score", 0)
         for k in ("label", "explanation"):
@@ -285,7 +285,7 @@ class ExperimentEvaluationRun:
 
     def __post_init__(self) -> None:
         if bool(self.result) == bool(self.error):
-            ValueError("Must specify either result or error")
+            raise ValueError("Must specify either result or error")
 
 
 ExperimentTask: TypeAlias = Union[

phoenix/server/api/context.py

@@ -1,12 +1,9 @@
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import AsyncContextManager, Callable, Optional, Union
+from typing import Callable, Optional
 
-from sqlalchemy.ext.asyncio import AsyncSession
-from starlette.requests import Request
-from starlette.responses import Response
-from starlette.websockets import WebSocket
+from strawberry.fastapi import BaseContext
 from typing_extensions import TypeAlias
 
 from phoenix.core.model_schema import Model
@@ -28,6 +25,7 @@ from phoenix.server.api.dataloaders import (
     ProjectByNameDataLoader,
     RecordCountDataLoader,
     SpanAnnotationsDataLoader,
+    SpanDatasetExamplesDataLoader,
     SpanDescendantsDataLoader,
     SpanEvaluationsDataLoader,
     SpanProjectsDataLoader,
@@ -35,6 +33,7 @@ from phoenix.server.api.dataloaders import (
     TraceEvaluationsDataLoader,
     TraceRowIdsDataLoader,
 )
+from phoenix.server.types import DbSessionFactory
 
 
 @dataclass
@@ -53,6 +52,7 @@ class DataLoaders:
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
     record_counts: RecordCountDataLoader
+    span_dataset_examples: SpanDatasetExamplesDataLoader
     span_descendants: SpanDescendantsDataLoader
     span_evaluations: SpanEvaluationsDataLoader
     span_projects: SpanProjectsDataLoader
@@ -67,10 +67,8 @@ ProjectRowId: TypeAlias = int
 
 
 @dataclass
-class Context:
-    request: Union[Request, WebSocket]
-    response: Optional[Response]
-    db: Callable[[], AsyncContextManager[AsyncSession]]
+class Context(BaseContext):
+    db: DbSessionFactory
     data_loaders: DataLoaders
     cache_for_dataloaders: Optional[CacheForDataLoaders]
     model: Model

phoenix/server/api/dataloaders/__init__.py

@@ -27,6 +27,7 @@ from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMax
 from .project_by_name import ProjectByNameDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
 from .span_annotations import SpanAnnotationsDataLoader
+from .span_dataset_examples import SpanDatasetExamplesDataLoader
 from .span_descendants import SpanDescendantsDataLoader
 from .span_evaluations import SpanEvaluationsDataLoader
 from .span_projects import SpanProjectsDataLoader
@@ -50,6 +51,7 @@ __all__ = [
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",
     "RecordCountDataLoader",
+    "SpanDatasetExamplesDataLoader",
     "SpanDescendantsDataLoader",
     "SpanEvaluationsDataLoader",
     "SpanProjectsDataLoader",

phoenix/server/api/dataloaders/average_experiment_run_latency.py

@@ -1,11 +1,11 @@
-from typing import AsyncContextManager, Callable, List, Optional
+from typing import List, Optional
 
 from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ExperimentID: TypeAlias = int
 RunLatency: TypeAlias = Optional[float]
@@ -16,7 +16,7 @@ Result: TypeAlias = RunLatency
 class AverageExperimentRunLatencyDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
     ) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db

phoenix/server/api/dataloaders/dataset_example_revisions.py

@@ -1,6 +1,4 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
     Tuple,
@@ -8,12 +6,12 @@ from typing import (
 )
 
 from sqlalchemy import Integer, case, func, literal, or_, select, union
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
+from phoenix.server.types import DbSessionFactory
 
 ExampleID: TypeAlias = int
 VersionID: TypeAlias = Optional[int]
@@ -22,7 +20,7 @@ Result: TypeAlias = DatasetExampleRevision
 
 
 class DatasetExampleRevisionsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/dataset_example_spans.py

@@ -1,17 +1,15 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
 )
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ExampleID: TypeAlias = int
 Key: TypeAlias = ExampleID
@@ -19,7 +17,7 @@ Result: TypeAlias = Optional[models.Span]
 
 
 class DatasetExampleSpansDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/document_evaluation_summaries.py

@@ -2,8 +2,6 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Optional,
@@ -14,7 +12,6 @@ import numpy as np
 from aioitertools.itertools import groupby
 from cachetools import LFUCache, TTLCache
 from sqlalchemy import Select, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias
 
@@ -24,6 +21,7 @@ from phoenix.metrics.retrieval_metrics import RetrievalMetrics
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
 from phoenix.server.api.types.DocumentEvaluationSummary import DocumentEvaluationSummary
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 ProjectRowId: TypeAlias = int
@@ -77,7 +75,7 @@ class DocumentEvaluationSummaryCache(
 class DocumentEvaluationSummaryDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

phoenix/server/api/dataloaders/document_evaluations.py

@@ -1,25 +1,23 @@
 from collections import defaultdict
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
 )
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.server.api.types.Evaluation import DocumentEvaluation
+from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
 Result: TypeAlias = List[DocumentEvaluation]
 
 
 class DocumentEvaluationsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/document_retrieval_metrics.py

@@ -1,7 +1,5 @@
 from collections import defaultdict
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     Dict,
     List,
@@ -13,13 +11,13 @@ from typing import (
 import numpy as np
 from aioitertools.itertools import groupby
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.metrics.retrieval_metrics import RetrievalMetrics
 from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
+from phoenix.server.types import DbSessionFactory
 
 RowId: TypeAlias = int
 NumDocs: TypeAlias = int
@@ -30,7 +28,7 @@ Result: TypeAlias = List[DocumentRetrievalMetrics]
 
 
 class DocumentRetrievalMetricsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/evaluation_summaries.py

@@ -2,8 +2,6 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -15,7 +13,6 @@ import pandas as pd
 from aioitertools.itertools import groupby
 from cachetools import LFUCache, TTLCache
 from sqlalchemy import Select, func, or_, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias, assert_never
 
@@ -23,6 +20,7 @@ from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
 from phoenix.server.api.types.EvaluationSummary import EvaluationSummary
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["span", "trace"]
@@ -77,7 +75,7 @@ class EvaluationSummaryCache(
 class EvaluationSummaryDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

phoenix/server/api/dataloaders/experiment_annotation_summaries.py

@@ -1,19 +1,17 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Optional,
 )
 
 from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 
 @dataclass
@@ -34,7 +32,7 @@ Result: TypeAlias = List[ExperimentAnnotationSummary]
 class ExperimentAnnotationSummaryDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(load_fn=self._load_fn)

phoenix/server/api/dataloaders/experiment_error_rates.py

@@ -1,16 +1,14 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
 )
 
 from sqlalchemy import case, func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ExperimentID: TypeAlias = int
 ErrorRate: TypeAlias = float
@@ -21,7 +19,7 @@ Result: TypeAlias = Optional[ErrorRate]
 class ExperimentErrorRatesDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
     ) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db

phoenix/server/api/dataloaders/experiment_run_counts.py

@@ -1,15 +1,13 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
 )
 
 from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ExperimentID: TypeAlias = int
 RunCount: TypeAlias = int
@@ -20,7 +18,7 @@ Result: TypeAlias = RunCount
 class ExperimentRunCountsDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
     ) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db