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

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
 

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__(

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__(

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
 

phoenix/server/api/dataloaders/record_counts.py

@@ -2,8 +2,6 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -13,13 +11,13 @@ from typing import (
 
 from cachetools import LFUCache, TTLCache
 from sqlalchemy import Select, 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.api.input_types.TimeRange import TimeRange
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["span", "trace"]
@@ -69,7 +67,7 @@ class RecordCountCache(
 class RecordCountDataLoader(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/span_annotations.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.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
 Result: TypeAlias = List[SpanAnnotation]
 
 
 class SpanAnnotationsDataLoader(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/span_dataset_examples.py

@@ -0,0 +1,36 @@
+from typing import (
+    Dict,
+    List,
+)
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SpanID: TypeAlias = int
+Key: TypeAlias = SpanID
+Result: TypeAlias = List[models.DatasetExample]
+
+
+class SpanDatasetExamplesDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        span_rowids = keys
+        async with self._db() as session:
+            dataset_examples: Dict[Key, List[models.DatasetExample]] = {
+                span_rowid: [] for span_rowid in span_rowids
+            }
+            async for span_rowid, dataset_example in await session.stream(
+                select(models.Span.id, models.DatasetExample)
+                .select_from(models.Span)
+                .join(models.DatasetExample, models.DatasetExample.span_rowid == models.Span.id)
+                .where(models.Span.id.in_(span_rowids))
+            ):
+                dataset_examples[span_rowid].append(dataset_example)
+        return [dataset_examples.get(span_rowid, []) for span_rowid in span_rowids]

phoenix/server/api/dataloaders/span_descendants.py

@@ -1,19 +1,17 @@
 from random import randint
 from typing import (
-    AsyncContextManager,
-    Callable,
     Dict,
     List,
 )
 
 from aioitertools.itertools import groupby
 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
 
 SpanId: TypeAlias = str
 
@@ -22,7 +20,7 @@ Result: TypeAlias = List[models.Span]
 
 
 class SpanDescendantsDataLoader(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/span_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 SpanEvaluation
+from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
 Result: TypeAlias = List[SpanEvaluation]
 
 
 class SpanEvaluationsDataLoader(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/span_projects.py

@@ -1,11 +1,11 @@
-from typing import AsyncContextManager, Callable, List, Union
+from typing import List, Union
 
 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
 
 SpanID: TypeAlias = int
 Key: TypeAlias = SpanID
@@ -13,7 +13,7 @@ Result: TypeAlias = models.Project
 
 
 class SpanProjectsDataLoader(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/token_counts.py

@@ -2,8 +2,6 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -14,7 +12,6 @@ from typing import (
 from cachetools import LFUCache, TTLCache
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import Select, func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.sql.functions import coalesce
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias
@@ -22,6 +19,7 @@ from typing_extensions import TypeAlias
 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.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["prompt", "completion", "total"]
@@ -71,7 +69,7 @@ class TokenCountCache(
 class TokenCountDataLoader(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/trace_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 TraceEvaluation
+from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
 Result: TypeAlias = List[TraceEvaluation]
 
 
 class TraceEvaluationsDataLoader(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/trace_row_ids.py

@@ -1,17 +1,15 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
     Tuple,
 )
 
 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
 
 TraceId: TypeAlias = str
 Key: TypeAlias = TraceId
@@ -21,7 +19,7 @@ Result: TypeAlias = Optional[Tuple[TraceRowId, ProjectRowId]]
 
 
 class TraceRowIdsDataLoader(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/input_types/SpanAnnotationSort.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+import strawberry
+
+from phoenix.server.api.types.SortDir import SortDir
+
+
+@strawberry.enum
+class SpanAnnotationColumn(Enum):
+    createdAt = "created_at"
+    name = "name"
+
+
+@strawberry.input(description="The sort key and direction for SpanAnnotation connections")
+class SpanAnnotationSort:
+    col: SpanAnnotationColumn
+    dir: SortDir

phoenix/server/api/input_types/TraceAnnotationSort.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+import strawberry
+
+from phoenix.server.api.types.SortDir import SortDir
+
+
+@strawberry.enum
+class TraceAnnotationColumn(Enum):
+    createdAt = "created_at"
+    name = "name"
+
+
+@strawberry.input(description="The sort key and direction for TraceAnnotation connections")
+class TraceAnnotationSort:
+    col: TraceAnnotationColumn
+    dir: SortDir

phoenix/server/api/mutations/span_annotations_mutations.py

@@ -11,6 +11,7 @@ from phoenix.server.api.input_types.CreateSpanAnnotationInput import CreateSpanA
 from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
 from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationInput
 from phoenix.server.api.mutations.auth import IsAuthenticated
+from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
 
@@ -18,6 +19,7 @@ from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_
 @strawberry.type
 class SpanAnnotationMutationPayload:
     span_annotations: List[SpanAnnotation]
+    query: Query
 
 
 @strawberry.type
@@ -49,7 +51,8 @@ class SpanAnnotationMutationMixin:
         return SpanAnnotationMutationPayload(
             span_annotations=[
                 to_gql_span_annotation(annotation) for annotation in inserted_annotations
-            ]
+            ],
+            query=Query(),
         )
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -89,7 +92,7 @@ class SpanAnnotationMutationMixin:
                 if span_annotation is not None:
                     patched_annotations.append(to_gql_span_annotation(span_annotation))
 
-        return SpanAnnotationMutationPayload(span_annotations=patched_annotations)
+        return SpanAnnotationMutationPayload(span_annotations=patched_annotations, query=Query())
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
     async def delete_span_annotations(
@@ -111,4 +114,6 @@ class SpanAnnotationMutationMixin:
             deleted_annotations_gql = [
                 to_gql_span_annotation(annotation) for annotation in deleted_annotations
             ]
-        return SpanAnnotationMutationPayload(span_annotations=deleted_annotations_gql)
+        return SpanAnnotationMutationPayload(
+            span_annotations=deleted_annotations_gql, query=Query()
+        )

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -11,6 +11,7 @@ from phoenix.server.api.input_types.CreateTraceAnnotationInput import CreateTrac
 from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
 from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationInput
 from phoenix.server.api.mutations.auth import IsAuthenticated
+from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
 
@@ -18,6 +19,7 @@ from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_tra
 @strawberry.type
 class TraceAnnotationMutationPayload:
     trace_annotations: List[TraceAnnotation]
+    query: Query
 
 
 @strawberry.type
@@ -49,7 +51,8 @@ class TraceAnnotationMutationMixin:
         return TraceAnnotationMutationPayload(
             trace_annotations=[
                 to_gql_trace_annotation(annotation) for annotation in inserted_annotations
-            ]
+            ],
+            query=Query(),
         )
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -89,7 +92,7 @@ class TraceAnnotationMutationMixin:
                 if trace_annotation:
                     patched_annotations.append(to_gql_trace_annotation(trace_annotation))
 
-        return TraceAnnotationMutationPayload(trace_annotations=patched_annotations)
+        return TraceAnnotationMutationPayload(trace_annotations=patched_annotations, query=Query())
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
     async def delete_trace_annotations(
@@ -111,4 +114,6 @@ class TraceAnnotationMutationMixin:
             deleted_annotations_gql = [
                 to_gql_trace_annotation(annotation) for annotation in deleted_annotations
             ]
-        return TraceAnnotationMutationPayload(trace_annotations=deleted_annotations_gql)
+        return TraceAnnotationMutationPayload(
+            trace_annotations=deleted_annotations_gql, query=Query()
+        )

phoenix/server/api/openapi/main.py

@@ -1,6 +1,22 @@
+import json
+from argparse import ArgumentParser
+from typing import Optional, Tuple
+
 from .schema import get_openapi_schema
 
 if __name__ == "__main__":
-    import yaml  # type: ignore
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--compress",
+        action="store_true",
+        help="Whether to output a compressed version of the OpenAPI schema",
+    )
+    args = parser.parse_args()
 
-    print(yaml.dump(get_openapi_schema(), indent=2))
+    indent: Optional[int] = None
+    separator: Optional[Tuple[str, str]] = None
+    if args.compress:
+        separator = (",", ":")
+    else:
+        indent = 2
+    print(json.dumps(get_openapi_schema(), indent=indent, separators=separator))

phoenix/server/api/openapi/schema.py

@@ -1,16 +1,16 @@
-from typing import Any
+from typing import Any, Dict
 
-from starlette.schemas import SchemaGenerator
+from fastapi.openapi.utils import get_openapi
 
-from phoenix.server.api.routers.v1 import V1_ROUTES
+from phoenix.server.api.routers.v1 import REST_API_VERSION
+from phoenix.server.api.routers.v1 import router as v1_router
 
-OPENAPI_SCHEMA_GENERATOR = SchemaGenerator(
-    {"openapi": "3.0.0", "info": {"title": "Arize-Phoenix API", "version": "1.0"}}
-)
 
-
-def get_openapi_schema() -> Any:
-    """
-    Exports an OpenAPI schema for the Phoenix REST API as a JSON object.
-    """
-    return OPENAPI_SCHEMA_GENERATOR.get_schema(V1_ROUTES)  # type: ignore
+def get_openapi_schema() -> Dict[str, Any]:
+    return get_openapi(
+        title="Arize-Phoenix REST API",
+        version=REST_API_VERSION,
+        openapi_version="3.1.0",
+        description="Schema for Arize-Phoenix REST API",
+        routes=v1_router.routes,
+    )