arize-phoenix 4.14.1__tar.gz → 4.15.0__tar.gz

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 4.14.1
+Version: 4.15.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -22,6 +22,7 @@ Requires-Dist: aiosqlite
 Requires-Dist: alembic<2,>=1.3.0
 Requires-Dist: arize-phoenix-evals>=0.13.1
 Requires-Dist: cachetools
+Requires-Dist: fastapi
 Requires-Dist: grpcio
 Requires-Dist: hdbscan>=0.8.33
 Requires-Dist: httpx
@@ -40,8 +41,8 @@ Requires-Dist: pandas>=1.0
 Requires-Dist: protobuf<6.0,>=3.20
 Requires-Dist: psutil
 Requires-Dist: pyarrow
+Requires-Dist: pydantic!=2.0.*,<3,>=1.0
 Requires-Dist: python-multipart
-Requires-Dist: pyyaml
 Requires-Dist: scikit-learn
 Requires-Dist: scipy
 Requires-Dist: sqlalchemy[asyncio]<3,>=2.0.4
@@ -56,9 +57,9 @@ Requires-Dist: uvicorn
 Requires-Dist: wrapt
 Provides-Extra: container
 Requires-Dist: opentelemetry-exporter-otlp; extra == 'container'
+Requires-Dist: opentelemetry-instrumentation-fastapi; extra == 'container'
 Requires-Dist: opentelemetry-instrumentation-grpc; extra == 'container'
 Requires-Dist: opentelemetry-instrumentation-sqlalchemy; extra == 'container'
-Requires-Dist: opentelemetry-instrumentation-starlette; extra == 'container'
 Requires-Dist: opentelemetry-proto>=1.12.0; extra == 'container'
 Requires-Dist: opentelemetry-sdk; extra == 'container'
 Requires-Dist: opentelemetry-semantic-conventions; extra == 'container'

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/pyproject.toml

@@ -59,7 +59,8 @@ dependencies = [
   "cachetools",
   "python-multipart",  # see https://www.starlette.io/#dependencies
   "arize-phoenix-evals>=0.13.1",
-  "pyyaml", # for OpenAPI
+  "fastapi",
+  "pydantic>=1.0,!=2.0.*,<3,", # exclude 2.0.* since it does not support the `json_encoders` configuration setting
 ]
 dynamic = ["version"]
 
@@ -111,7 +112,7 @@ container = [
   "opentelemetry-proto>=1.12.0",
   "opentelemetry-exporter-otlp",
   "opentelemetry-semantic-conventions",
-  "opentelemetry-instrumentation-starlette",
+  "opentelemetry-instrumentation-fastapi",
   "opentelemetry-instrumentation-sqlalchemy",
   "opentelemetry-instrumentation-grpc",
   "py-grpc-prometheus",
@@ -161,8 +162,10 @@ dependencies = [
   "protobuf==3.20",  # version minimum (for tests)
   "responses",
   "tiktoken",
-  "typing-extensions==4.5.0; python_version<'3.12'",
-  "typing-extensions==4.6.0; python_version>='3.12'",
+  "typing-extensions==4.5.0; python_version=='3.8'",
+  "typing-extensions==4.6.0; python_version=='3.9'",
+  "pydantic==1.9.0; python_version<='3.9'",  # minimum version of pydantic compatible with openai
+  "pydantic==2.8.2; python_version=='3.12'",
   "httpx", # For OpenAI testing
   "respx", # For OpenAI testing
   "nest-asyncio", # for executor testing
@@ -190,12 +193,15 @@ dependencies = [
   "opentelemetry-proto>=1.12.0",
   "opentelemetry-exporter-otlp",
   "opentelemetry-semantic-conventions",
-  "opentelemetry-instrumentation-starlette",
+  "opentelemetry-instrumentation-fastapi",
   "opentelemetry-instrumentation-sqlalchemy",
   "opentelemetry-instrumentation-grpc",
   "py-grpc-prometheus",
   "strawberry-graphql[opentelemetry]==0.236.0",  # need to pin version because we're monkey-patching
   "requests",  # this is needed to type-check third-party packages
+  "pydantic==1.10.17; python_version=='3.8'",  # lower minor versions of pydantic break strawberry mypy plugin
+  "pydantic==1.10.17; python_version=='3.9'",  # lower minor versions of pydantic break strawberry mypy plugin
+  "pydantic==2.8.2; python_version=='3.12'",
 ]
 
 [[tool.hatch.envs.type.matrix]]
@@ -310,11 +316,12 @@ build = 'strawberry export-schema phoenix.server.api.schema:schema > app/schema.
 
 [tool.hatch.envs.openapi]
 dependencies = [
-  "pyyaml",
+  "pydantic==2.8.2",
+  "fastapi==0.111.0",
 ]
 
 [tool.hatch.envs.openapi.scripts]
-build = 'python -m phoenix.server.api.openapi.main > schemas/openapi.yaml'
+build = "python -m phoenix.server.api.openapi.main > schemas/openapi.json"
 
 [tool.hatch.envs.proto]
 detached = true
@@ -379,6 +386,7 @@ module = [
   "sqlean",
   "grpc.*",
   "py_grpc_prometheus.*",
+  "orjson",  # suppress fastapi internal type errors
 ]
 ignore_missing_imports = true
 

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/phoenix/db/bulk_inserter.py

@@ -7,7 +7,6 @@ from itertools import islice
 from time import perf_counter
 from typing import (
     Any,
-    AsyncContextManager,
     Awaitable,
     Callable,
     Iterable,
@@ -19,7 +18,6 @@ from typing import (
 )
 
 from cachetools import LRUCache
-from sqlalchemy.ext.asyncio import AsyncSession
 from typing_extensions import TypeAlias
 
 import phoenix.trace.v1 as pb
@@ -31,6 +29,7 @@ from phoenix.db.insertion.evaluation import (
 from phoenix.db.insertion.helpers import DataManipulation, DataManipulationEvent
 from phoenix.db.insertion.span import SpanInsertionEvent, insert_span
 from phoenix.server.api.dataloaders import CacheForDataLoaders
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.schemas import Span
 
 logger = logging.getLogger(__name__)
@@ -46,7 +45,7 @@ class TransactionResult:
 class BulkInserter:
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         *,
         cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
         initial_batch_of_operations: Iterable[DataManipulation] = (),

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/phoenix/db/engines.py

@@ -8,7 +8,7 @@ from typing import Any
 import aiosqlite
 import numpy as np
 import sqlean
-from sqlalchemy import URL, event, make_url
+from sqlalchemy import URL, StaticPool, event, make_url
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from typing_extensions import assert_never
 
@@ -105,6 +105,7 @@ def aio_sqlite_engine(
         echo=echo,
         json_serializer=_dumps,
         async_creator=async_creator,
+        poolclass=StaticPool,
     )
     event.listen(engine.sync_engine, "connect", set_sqlite_pragma)
     if not migrate:

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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,
     )

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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",

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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
 

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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
 

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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__(

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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
 

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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
 

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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__(

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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)

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/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

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/phoenix/server/api/dataloaders/experiment_sequence_number.py

@@ -1,16 +1,14 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
 )
 
 from sqlalchemy import distinct, 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
 Key: TypeAlias = ExperimentId
@@ -18,7 +16,7 @@ Result: TypeAlias = Optional[int]
 
 
 class ExperimentSequenceNumberDataLoader(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
 

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/phoenix/server/api/dataloaders/latency_ms_quantile.py

@@ -2,9 +2,7 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
     AsyncIterator,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -36,6 +34,7 @@ from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["span", "trace"]
@@ -88,7 +87,7 @@ class LatencyMsQuantileCache(
 class LatencyMsQuantileDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/phoenix/server/api/dataloaders/min_start_or_max_end_times.py

@@ -1,8 +1,6 @@
 from collections import defaultdict
 from datetime import datetime
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -12,12 +10,12 @@ from typing import (
 
 from cachetools import LFUCache
 from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias, assert_never
 
 from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
+from phoenix.server.types import DbSessionFactory
 
 Kind: TypeAlias = Literal["start", "end"]
 ProjectRowId: TypeAlias = int
@@ -50,7 +48,7 @@ class MinStartOrMaxEndTimeCache(
 class MinStartOrMaxEndTimeDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

{arize_phoenix-4.14.1 → arize_phoenix-4.15.0}/src/phoenix/server/api/dataloaders/project_by_name.py

@@ -1,12 +1,12 @@
 from collections import defaultdict
-from typing import AsyncContextManager, Callable, DefaultDict, List, Optional
+from typing import DefaultDict, List, Optional
 
 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.types import DbSessionFactory
 
 ProjectName: TypeAlias = str
 Key: TypeAlias = ProjectName
@@ -14,7 +14,7 @@ Result: TypeAlias = Optional[models.Project]
 
 
 class ProjectByNameDataLoader(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