arize-phoenix 5.5.2__py3-none-any.whl → 5.7.0__py3-none-any.whl

phoenix/server/api/dataloaders/record_counts.py

@@ -1,13 +1,6 @@
 from collections import defaultdict
 from datetime import datetime
-from typing import (
-    Any,
-    DefaultDict,
-    List,
-    Literal,
-    Optional,
-    Tuple,
-)
+from typing import Any, Literal, Optional
 
 from cachetools import LFUCache, TTLCache
 from sqlalchemy import Select, func, select
@@ -22,20 +15,20 @@ from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
-TimeInterval: TypeAlias = Tuple[Optional[datetime], Optional[datetime]]
+TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
 SpanCount: TypeAlias = int
 
-Segment: TypeAlias = Tuple[Kind, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition]
 Param: TypeAlias = ProjectRowId
 
-Key: TypeAlias = Tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
+Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
 Result: TypeAlias = SpanCount
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = 0
 
 
-def _cache_key_fn(key: Key) -> Tuple[Segment, Param]:
+def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
     kind, project_rowid, time_range, filter_condition = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
@@ -44,7 +37,7 @@ def _cache_key_fn(key: Key) -> Tuple[Segment, Param]:
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = Tuple[TimeInterval, FilterCondition, Kind]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, Kind]
 
 
 class RecordCountCache(
@@ -59,7 +52,7 @@ class RecordCountCache(
             sub_cache_factory=lambda: LFUCache(maxsize=2 * 2 * 2),
         )
 
-    def _cache_key(self, key: Key) -> Tuple[_Section, _SubKey]:
+    def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
         (kind, interval, filter_condition), project_rowid = _cache_key_fn(key)
         return project_rowid, (interval, filter_condition, kind)
 
@@ -77,11 +70,11 @@ class RecordCountDataLoader(DataLoader[Key, Result]):
         )
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        results: List[Result] = [DEFAULT_VALUE] * len(keys)
-        arguments: DefaultDict[
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        results: list[Result] = [DEFAULT_VALUE] * len(keys)
+        arguments: defaultdict[
             Segment,
-            DefaultDict[Param, List[ResultPosition]],
+            defaultdict[Param, list[ResultPosition]],
         ] = defaultdict(lambda: defaultdict(list))
         for position, key in enumerate(keys):
             segment, param = _cache_key_fn(key)

phoenix/server/api/dataloaders/span_annotations.py

@@ -1,8 +1,4 @@
 from collections import defaultdict
-from typing import (
-    DefaultDict,
-    List,
-)
 
 from sqlalchemy import select
 from strawberry.dataloader import DataLoader
@@ -12,7 +8,7 @@ from phoenix.db.models import SpanAnnotation as ORMSpanAnnotation
 from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
-Result: TypeAlias = List[ORMSpanAnnotation]
+Result: TypeAlias = list[ORMSpanAnnotation]
 
 
 class SpanAnnotationsDataLoader(DataLoader[Key, Result]):
@@ -20,8 +16,8 @@ class SpanAnnotationsDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        span_annotations_by_id: DefaultDict[Key, Result] = defaultdict(list)
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        span_annotations_by_id: defaultdict[Key, Result] = defaultdict(list)
         async with self._db() as session:
             async for span_annotation in await session.stream_scalars(
                 select(ORMSpanAnnotation).where(ORMSpanAnnotation.span_rowid.in_(keys))

phoenix/server/api/dataloaders/span_dataset_examples.py

@@ -1,8 +1,3 @@
-from typing import (
-    Dict,
-    List,
-)
-
 from sqlalchemy import select
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
@@ -12,7 +7,7 @@ from phoenix.server.types import DbSessionFactory
 
 SpanID: TypeAlias = int
 Key: TypeAlias = SpanID
-Result: TypeAlias = List[models.DatasetExample]
+Result: TypeAlias = list[models.DatasetExample]
 
 
 class SpanDatasetExamplesDataLoader(DataLoader[Key, Result]):
@@ -20,10 +15,10 @@ class SpanDatasetExamplesDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+    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]] = {
+            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(

phoenix/server/api/dataloaders/span_descendants.py

@@ -1,8 +1,4 @@
 from random import randint
-from typing import (
-    Dict,
-    List,
-)
 
 from aioitertools.itertools import groupby
 from sqlalchemy import select
@@ -16,7 +12,7 @@ from phoenix.server.types import DbSessionFactory
 SpanId: TypeAlias = str
 
 Key: TypeAlias = SpanId
-Result: TypeAlias = List[models.Span]
+Result: TypeAlias = list[models.Span]
 
 
 class SpanDescendantsDataLoader(DataLoader[Key, Result]):
@@ -24,7 +20,7 @@ class SpanDescendantsDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
         root_ids = set(keys)
         root_id_label = f"root_id_{randint(0, 10**6):06}"
         descendant_ids = (
@@ -53,7 +49,7 @@ class SpanDescendantsDataLoader(DataLoader[Key, Result]):
             .options(joinedload(models.Span.trace, innerjoin=True).load_only(models.Trace.trace_id))
             .order_by(descendant_ids.c[root_id_label])
         )
-        results: Dict[SpanId, Result] = {key: [] for key in keys}
+        results: dict[SpanId, Result] = {key: [] for key in keys}
         async with self._db() as session:
             data = await session.stream(stmt)
             async for root_id, group in groupby(data, key=lambda d: d[0]):

phoenix/server/api/dataloaders/span_projects.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import Union
 
 from sqlalchemy import select
 from strawberry.dataloader import DataLoader
@@ -17,7 +17,7 @@ class SpanProjectsDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Union[Result, ValueError]]:
+    async def _load_fn(self, keys: list[Key]) -> list[Union[Result, ValueError]]:
         span_ids = list(set(keys))
         async with self._db() as session:
             projects = {

phoenix/server/api/dataloaders/token_counts.py

@@ -1,13 +1,6 @@
 from collections import defaultdict
 from datetime import datetime
-from typing import (
-    Any,
-    DefaultDict,
-    List,
-    Literal,
-    Optional,
-    Tuple,
-)
+from typing import Any, Literal, Optional
 
 from cachetools import LFUCache, TTLCache
 from sqlalchemy import Select, func, select
@@ -23,20 +16,20 @@ from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["prompt", "completion", "total"]
 ProjectRowId: TypeAlias = int
-TimeInterval: TypeAlias = Tuple[Optional[datetime], Optional[datetime]]
+TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
 TokenCount: TypeAlias = int
 
-Segment: TypeAlias = Tuple[TimeInterval, FilterCondition]
-Param: TypeAlias = Tuple[ProjectRowId, Kind]
+Segment: TypeAlias = tuple[TimeInterval, FilterCondition]
+Param: TypeAlias = tuple[ProjectRowId, Kind]
 
-Key: TypeAlias = Tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
+Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
 Result: TypeAlias = TokenCount
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = 0
 
 
-def _cache_key_fn(key: Key) -> Tuple[Segment, Param]:
+def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
     kind, project_rowid, time_range, filter_condition = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
@@ -45,7 +38,7 @@ def _cache_key_fn(key: Key) -> Tuple[Segment, Param]:
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = Tuple[TimeInterval, FilterCondition, Kind]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, Kind]
 
 
 class TokenCountCache(
@@ -60,7 +53,7 @@ class TokenCountCache(
             sub_cache_factory=lambda: LFUCache(maxsize=2 * 2 * 3),
         )
 
-    def _cache_key(self, key: Key) -> Tuple[_Section, _SubKey]:
+    def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
         (interval, filter_condition), (project_rowid, kind) = _cache_key_fn(key)
         return project_rowid, (interval, filter_condition, kind)
 
@@ -78,11 +71,11 @@ class TokenCountDataLoader(DataLoader[Key, Result]):
         )
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        results: List[Result] = [DEFAULT_VALUE] * len(keys)
-        arguments: DefaultDict[
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        results: list[Result] = [DEFAULT_VALUE] * len(keys)
+        arguments: defaultdict[
             Segment,
-            DefaultDict[Param, List[ResultPosition]],
+            defaultdict[Param, list[ResultPosition]],
         ] = defaultdict(lambda: defaultdict(list))
         for position, key in enumerate(keys):
             segment, param = _cache_key_fn(key)

phoenix/server/api/dataloaders/trace_row_ids.py

@@ -1,8 +1,4 @@
-from typing import (
-    List,
-    Optional,
-    Tuple,
-)
+from typing import Optional
 
 from sqlalchemy import select
 from strawberry.dataloader import DataLoader
@@ -15,7 +11,7 @@ TraceId: TypeAlias = str
 Key: TypeAlias = TraceId
 TraceRowId: TypeAlias = int
 ProjectRowId: TypeAlias = int
-Result: TypeAlias = Optional[Tuple[TraceRowId, ProjectRowId]]
+Result: TypeAlias = Optional[tuple[TraceRowId, ProjectRowId]]
 
 
 class TraceRowIdsDataLoader(DataLoader[Key, Result]):
@@ -23,7 +19,7 @@ class TraceRowIdsDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
         stmt = select(
             models.Trace.trace_id,
             models.Trace.id,

phoenix/server/api/dataloaders/user_roles.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import DefaultDict, List, Optional
+from typing import Optional
 
 from sqlalchemy import select
 from strawberry.dataloader import DataLoader
@@ -20,8 +20,8 @@ class UserRolesDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        user_roles_by_id: DefaultDict[Key, Result] = defaultdict(None)
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        user_roles_by_id: defaultdict[Key, Result] = defaultdict(None)
         async with self._db() as session:
             data = await session.stream_scalars(select(models.UserRole))
             async for user_role in data:

phoenix/server/api/dataloaders/users.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import DefaultDict, List, Optional
+from typing import Optional
 
 from sqlalchemy import select
 from strawberry.dataloader import DataLoader
@@ -20,9 +20,9 @@ class UsersDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
         user_ids = list(set(keys))
-        users_by_id: DefaultDict[Key, Result] = defaultdict(None)
+        users_by_id: defaultdict[Key, Result] = defaultdict(None)
         async with self._db() as session:
             data = await session.stream_scalars(
                 select(models.User).where(models.User.id.in_(user_ids))

phoenix/server/api/helpers/__init__.py

@@ -1,10 +1,11 @@
-from typing import Iterable, List, Optional, TypeVar
+from collections.abc import Iterable
+from typing import Optional, TypeVar
 
 T = TypeVar("T")
 
 
-def ensure_list(obj: Optional[Iterable[T]]) -> List[T]:
-    if isinstance(obj, List):
+def ensure_list(obj: Optional[Iterable[T]]) -> list[T]:
+    if isinstance(obj, list):
         return obj
     if isinstance(obj, Iterable):
         return list(obj)

phoenix/server/api/helpers/dataset_helpers.py

@@ -1,5 +1,6 @@
 import json
-from typing import Any, Dict, Literal, Mapping, Optional, Protocol
+from collections.abc import Mapping
+from typing import Any, Literal, Optional, Protocol
 
 from openinference.semconv.trace import (
     MessageAttributes,
@@ -28,7 +29,7 @@ class HasSpanIO(Protocol):
     retrieval_documents: Any
 
 
-def get_dataset_example_input(span: HasSpanIO) -> Dict[str, Any]:
+def get_dataset_example_input(span: HasSpanIO) -> dict[str, Any]:
     """
     Extracts the input value from a span and returns it as a dictionary. Input
     values from LLM spans are extracted from the input messages and prompt
@@ -47,7 +48,7 @@ def get_dataset_example_input(span: HasSpanIO) -> Dict[str, Any]:
     return _get_generic_io_value(io_value=input_value, mime_type=input_mime_type, kind="input")
 
 
-def get_dataset_example_output(span: HasSpanIO) -> Dict[str, Any]:
+def get_dataset_example_output(span: HasSpanIO) -> dict[str, Any]:
     """
     Extracts the output value from a span and returns it as a dictionary. Output
     values from LLM spans are extracted from the output messages (if present).
@@ -78,13 +79,13 @@ def _get_llm_span_input(
     input_value: Any,
     input_mime_type: Optional[str],
     prompt_template_variables: Any,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Extracts the input value from an LLM span and returns it as a dictionary.
     The input is extracted from the input messages (if present) and prompt
     template variables (if present).
     """
-    input: Dict[str, Any] = {}
+    input: dict[str, Any] = {}
     if messages := [_get_message(m) for m in input_messages or ()]:
         input["messages"] = messages
     if not input:
@@ -98,7 +99,7 @@ def _get_llm_span_output(
     output_messages: Any,
     output_value: Any,
     output_mime_type: Optional[str],
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Extracts the output value from an LLM span and returns it as a dictionary.
     The output is extracted from the output messages (if present).
@@ -112,7 +113,7 @@ def _get_retriever_span_output(
     retrieval_documents: Any,
     output_value: Any,
     output_mime_type: Optional[str],
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Extracts the output value from a retriever span and returns it as a dictionary.
     The output is extracted from the retrieval documents (if present).
@@ -124,7 +125,7 @@ def _get_retriever_span_output(
 
 def _get_generic_io_value(
     io_value: Any, mime_type: Optional[str], kind: Literal["input", "output"]
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Makes a best-effort attempt to extract the input or output value from a span
     and returns it as a dictionary.
@@ -140,7 +141,7 @@ def _get_generic_io_value(
     return {}
 
 
-def _get_message(message: Mapping[str, Any]) -> Dict[str, Any]:
+def _get_message(message: Mapping[str, Any]) -> dict[str, Any]:
     content = get_attribute_value(message, MESSAGE_CONTENT)
     name = get_attribute_value(message, MESSAGE_NAME)
     function_call_name = get_attribute_value(message, MESSAGE_FUNCTION_CALL_NAME)