arize-phoenix 11.24.1__tar.gz → 11.25.0__tar.gz

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize-phoenix
-Version: 11.24.1
+Version: 11.25.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://arize.com/docs/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/__init__.py

@@ -87,8 +87,9 @@ class PhoenixTraceOpenAILoader(Loader):
             "https://arize.com/docs/phoenix/tracing/integrations-tracing/openai"
             "\n\n"
             "Example usage:\n\n"
+            "pip install openinference-instrumentation-openai\n\n"
             "```python\n"
-            "from phoenix.otel register\n"
+            "from phoenix.otel import register\n"
             "from openinference.instrumentation.openai import OpenAIInstrumentor\n\n"
             "tracer_provider = register()\n"
             "OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)\n"

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/db/engines.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import logging
 from collections.abc import Callable
-from datetime import datetime
 from enum import Enum
 from sqlite3 import Connection
 from typing import Any
 
 import aiosqlite
 import numpy as np
+import orjson
 import sqlalchemy
 import sqlean
 from sqlalchemy import URL, StaticPool, event, make_url
@@ -193,19 +192,14 @@ def aio_postgresql_engine(
 
 
 def _dumps(obj: Any) -> str:
-    return json.dumps(obj, cls=_Encoder)
-
-
-class _Encoder(json.JSONEncoder):
-    def default(self, obj: Any) -> Any:
-        if isinstance(obj, datetime):
-            return obj.isoformat()
-        elif isinstance(obj, Enum):
-            return obj.value
-        elif isinstance(obj, np.ndarray):
-            return list(obj)
-        elif isinstance(obj, np.integer):
-            return int(obj)
-        elif isinstance(obj, np.floating):
-            return float(obj)
-        return super().default(obj)
+    return orjson.dumps(obj, default=_default).decode()
+
+
+def _default(obj: Any) -> Any:
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, (np.integer, np.floating, np.bool_)):
+        return obj.item()
+    if isinstance(obj, Enum):
+        return obj.value
+    raise TypeError(f"Object of type {type(obj).__name__} is not serializable")

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/db/insertion/document_annotation.py

@@ -107,7 +107,6 @@ class DocumentAnnotationQueueInserter(
                             received_at=p.received_at,
                             item=p.item.as_insertable(
                                 span_rowid=anno.span_rowid,
-                                id_=anno.id_,
                             ),
                         )
                     )

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/db/insertion/span_annotation.py

@@ -100,7 +100,6 @@ class SpanAnnotationQueueInserter(
                             received_at=p.received_at,
                             item=p.item.as_insertable(
                                 span_rowid=anno.span_rowid,
-                                id_=anno.id_,
                             ),
                         )
                     )

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/db/insertion/trace_annotation.py

@@ -99,7 +99,6 @@ class TraceAnnotationQueueInserter(
                             received_at=p.received_at,
                             item=p.item.as_insertable(
                                 trace_rowid=anno.trace_rowid,
-                                id_=anno.id_,
                             ),
                         )
                     )

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/db/insertion/types.py

@@ -180,13 +180,11 @@ class Precursors(ABC):
         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)
@@ -197,13 +195,11 @@ class Precursors(ABC):
         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)
@@ -215,14 +211,12 @@ class Precursors(ABC):
         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_,
             )
 
 
@@ -231,40 +225,31 @@ class Insertables(ABC):
     class SpanAnnotation(Precursors.SpanAnnotation):
         span_rowid: int
         identifier: str = ""
-        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
         identifier: str = ""
-        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
         identifier: str = ""
-        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

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/db/models.py

@@ -188,7 +188,7 @@ class JsonDict(TypeDecorator[dict[str, Any]]):
     impl = JSON_
 
     def process_bind_param(self, value: Optional[dict[str, Any]], _: Dialect) -> dict[str, Any]:
-        return orjson.loads(orjson.dumps(value)) if isinstance(value, dict) else {}
+        return value if isinstance(value, dict) else {}
 
     def process_result_value(self, value: Optional[Any], _: Dialect) -> Optional[dict[str, Any]]:
         return orjson.loads(orjson.dumps(value)) if isinstance(value, dict) and value else value
@@ -200,7 +200,7 @@ class JsonList(TypeDecorator[list[Any]]):
     impl = JSON_
 
     def process_bind_param(self, value: Optional[list[Any]], _: Dialect) -> list[Any]:
-        return orjson.loads(orjson.dumps(value)) if isinstance(value, list) else []
+        return value if isinstance(value, list) else []
 
     def process_result_value(self, value: Optional[Any], _: Dialect) -> Optional[list[Any]]:
         return orjson.loads(orjson.dumps(value)) if isinstance(value, list) and value else value

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/server/api/dataloaders/annotation_summaries.py

@@ -13,6 +13,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.AnnotationSummary import AnnotationSummary
+from phoenix.server.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
@@ -20,27 +21,41 @@ Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 AnnotationName: TypeAlias = str
 
-Segment: TypeAlias = tuple[Kind, ProjectRowId, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[
+    Kind,
+    ProjectRowId,
+    TimeInterval,
+    FilterCondition,
+    SessionFilterCondition,
+]
 Param: TypeAlias = AnnotationName
 
-Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition, AnnotationName]
+Key: TypeAlias = tuple[
+    Kind,
+    ProjectRowId,
+    Optional[TimeRange],
+    FilterCondition,
+    SessionFilterCondition,
+    AnnotationName,
+]
 Result: TypeAlias = Optional[AnnotationSummary]
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = None
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    kind, project_rowid, time_range, filter_condition, eval_name = key
+    kind, project_rowid, time_range, filter_condition, session_filter_condition, eval_name = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
     )
-    return (kind, project_rowid, interval, filter_condition), eval_name
+    return (kind, project_rowid, interval, filter_condition, session_filter_condition), eval_name
 
 
 _Section: TypeAlias = tuple[ProjectRowId, AnnotationName, Kind]
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition]
 
 
 class AnnotationSummaryCache(
@@ -61,8 +76,21 @@ class AnnotationSummaryCache(
                 del self._cache[section]
 
     def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
-        (kind, project_rowid, interval, filter_condition), annotation_name = _cache_key_fn(key)
-        return (project_rowid, annotation_name, kind), (interval, filter_condition)
+        (
+            (
+                kind,
+                project_rowid,
+                interval,
+                filter_condition,
+                session_filter_condition,
+            ),
+            annotation_name,
+        ) = _cache_key_fn(key)
+        return (project_rowid, annotation_name, kind), (
+            interval,
+            filter_condition,
+            session_filter_condition,
+        )
 
 
 class AnnotationSummaryDataLoader(DataLoader[Key, Result]):
@@ -102,7 +130,9 @@ def _get_stmt(
     segment: Segment,
     *annotation_names: Param,
 ) -> Select[Any]:
-    kind, project_rowid, (start_time, end_time), filter_condition = segment
+    kind, project_rowid, (start_time, end_time), filter_condition, session_filter_condition = (
+        segment
+    )
 
     annotation_model: Union[Type[models.SpanAnnotation], Type[models.TraceAnnotation]]
     entity_model: Union[Type[models.Span], Type[models.Trace]]
@@ -144,6 +174,19 @@ def _get_stmt(
         entity_count_query = entity_count_query.where(
             cast(Type[models.Trace], entity_model).project_rowid == project_rowid
         )
+    else:
+        assert_never(kind)
+
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=[project_rowid],
+            start_time=start_time,
+            end_time=end_time,
+        )
+        entity_count_query = entity_count_query.where(
+            models.Trace.project_session_rowid.in_(filtered_session_rowids)
+        )
 
     entity_count_query = entity_count_query.where(
         or_(score_column.is_not(None), label_column.is_not(None))
@@ -186,6 +229,15 @@ def _get_stmt(
     else:
         assert_never(kind)
 
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=[project_rowid],
+            start_time=start_time,
+            end_time=end_time,
+        )
+        base_stmt = base_stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
+
     base_stmt = base_stmt.where(or_(score_column.is_not(None), label_column.is_not(None)))
     base_stmt = base_stmt.where(name_column.in_(annotation_names))
 

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/server/api/dataloaders/latency_ms_quantile.py

@@ -25,6 +25,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.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
@@ -32,13 +33,16 @@ Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 Probability: TypeAlias = float
 QuantileValue: TypeAlias = float
 
-Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition, SessionFilterCondition]
 Param: TypeAlias = tuple[ProjectRowId, Probability]
 
-Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition, Probability]
+Key: TypeAlias = tuple[
+    Kind, ProjectRowId, Optional[TimeRange], FilterCondition, SessionFilterCondition, Probability
+]
 Result: TypeAlias = Optional[QuantileValue]
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = None
@@ -47,15 +51,18 @@ FloatCol: TypeAlias = SQLColumnExpression[Float[float]]
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    kind, project_rowid, time_range, filter_condition, probability = key
+    kind, project_rowid, time_range, filter_condition, session_filter_condition, probability = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
     )
-    return (kind, interval, filter_condition), (project_rowid, probability)
+    return (kind, interval, filter_condition, session_filter_condition), (
+        project_rowid,
+        probability,
+    )
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, Kind, Probability]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition, Kind, Probability]
 
 
 class LatencyMsQuantileCache(
@@ -71,8 +78,17 @@ class LatencyMsQuantileCache(
         )
 
     def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
-        (kind, interval, filter_condition), (project_rowid, probability) = _cache_key_fn(key)
-        return project_rowid, (interval, filter_condition, kind, probability)
+        (
+            (kind, interval, filter_condition, session_filter_condition),
+            (project_rowid, probability),
+        ) = _cache_key_fn(key)
+        return project_rowid, (
+            interval,
+            filter_condition,
+            session_filter_condition,
+            kind,
+            probability,
+        )
 
 
 class LatencyMsQuantileDataLoader(DataLoader[Key, Result]):
@@ -113,11 +129,18 @@ async def _get_results(
     segment: Segment,
     params: Mapping[Param, list[ResultPosition]],
 ) -> AsyncIterator[tuple[ResultPosition, QuantileValue]]:
-    kind, (start_time, end_time), filter_condition = segment
+    kind, (start_time, end_time), filter_condition, session_filter_condition = segment
     stmt = select(models.Trace.project_rowid)
     if kind == "trace":
         latency_column = cast(FloatCol, models.Trace.latency_ms)
         time_column = models.Trace.start_time
+        if filter_condition:
+            sf = SpanFilter(filter_condition)
+            stmt = stmt.where(
+                models.Trace.id.in_(
+                    sf(select(models.Span.trace_rowid).distinct()).scalar_subquery()
+                )
+            )
     elif kind == "span":
         latency_column = cast(FloatCol, models.Span.latency_ms)
         time_column = models.Span.start_time
@@ -127,6 +150,15 @@ async def _get_results(
             stmt = sf(stmt)
     else:
         assert_never(kind)
+    if session_filter_condition:
+        project_rowids = [project_rowid for project_rowid, _ in params]
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=project_rowids,
+            start_time=start_time,
+            end_time=end_time,
+        )
+        stmt = stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
     if start_time:
         stmt = stmt.where(start_time <= time_column)
     if end_time:

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/server/api/dataloaders/record_counts.py

@@ -3,13 +3,14 @@ from datetime import datetime
 from typing import Any, Literal, Optional
 
 from cachetools import LFUCache, TTLCache
-from sqlalchemy import Select, func, select
+from sqlalchemy import Select, distinct, func, select
 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.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
@@ -17,27 +18,35 @@ Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 SpanCount: TypeAlias = int
 
-Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition, SessionFilterCondition]
 Param: TypeAlias = ProjectRowId
 
-Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
+Key: TypeAlias = tuple[
+    Kind, ProjectRowId, Optional[TimeRange], FilterCondition, SessionFilterCondition
+]
 Result: TypeAlias = SpanCount
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = 0
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    kind, project_rowid, time_range, filter_condition = key
+    kind, project_rowid, time_range, filter_condition, session_filter_condition = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
     )
-    return (kind, interval, filter_condition), project_rowid
+    return (
+        kind,
+        interval,
+        filter_condition,
+        session_filter_condition,
+    ), project_rowid
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, Kind]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition, Kind]
 
 
 class RecordCountCache(
@@ -53,8 +62,10 @@ class RecordCountCache(
         )
 
     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)
+        (kind, interval, filter_condition, session_filter_condition), project_rowid = _cache_key_fn(
+            key
+        )
+        return project_rowid, (interval, filter_condition, session_filter_condition, kind)
 
 
 class RecordCountDataLoader(DataLoader[Key, Result]):
@@ -93,7 +104,7 @@ def _get_stmt(
     segment: Segment,
     *project_rowids: Param,
 ) -> Select[Any]:
-    kind, (start_time, end_time), filter_condition = segment
+    kind, (start_time, end_time), filter_condition, session_filter_condition = segment
     pid = models.Trace.project_rowid
     stmt = select(pid)
     if kind == "span":
@@ -102,12 +113,28 @@ def _get_stmt(
         if filter_condition:
             sf = SpanFilter(filter_condition)
             stmt = sf(stmt)
+        stmt = stmt.add_columns(func.count().label("count"))
     elif kind == "trace":
         time_column = models.Trace.start_time
+        if filter_condition:
+            stmt = stmt.join(models.Span, models.Trace.id == models.Span.trace_rowid)
+            stmt = stmt.add_columns(func.count(distinct(models.Trace.id)).label("count"))
+            sf = SpanFilter(filter_condition)
+            stmt = sf(stmt)
+        else:
+            stmt = stmt.add_columns(func.count().label("count"))
     else:
         assert_never(kind)
-    stmt = stmt.add_columns(func.count().label("count"))
     stmt = stmt.where(pid.in_(project_rowids))
+
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=project_rowids,
+            start_time=start_time,
+            end_time=end_time,
+        )
+        stmt = stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
     stmt = stmt.group_by(pid)
     if start_time:
         stmt = stmt.where(start_time <= time_column)

{arize_phoenix-11.24.1 → arize_phoenix-11.25.0}/src/phoenix/server/api/dataloaders/span_cost_summary_by_project.py

@@ -12,32 +12,38 @@ from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.dataloaders.types import CostBreakdown, SpanCostSummary
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 
-Segment: TypeAlias = tuple[TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[
+    TimeInterval,
+    FilterCondition,
+    SessionFilterCondition,
+]
 Param: TypeAlias = ProjectRowId
 
-Key: TypeAlias = tuple[ProjectRowId, Optional[TimeRange], FilterCondition]
+Key: TypeAlias = tuple[ProjectRowId, Optional[TimeRange], FilterCondition, SessionFilterCondition]
 Result: TypeAlias = SpanCostSummary
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = SpanCostSummary()
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    project_rowid, time_range, filter_condition = key
+    project_rowid, time_range, filter_condition, session_filter_condition = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
     )
-    return (interval, filter_condition), project_rowid
+    return (interval, filter_condition, session_filter_condition), project_rowid
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition]
 
 
 class SpanCostSummaryCache(
@@ -53,8 +59,8 @@ class SpanCostSummaryCache(
         )
 
     def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
-        (interval, filter_condition), project_rowid = _cache_key_fn(key)
-        return project_rowid, (interval, filter_condition)
+        (interval, filter_condition, session_filter_condition), project_rowid = _cache_key_fn(key)
+        return project_rowid, (interval, filter_condition, session_filter_condition)
 
 
 class SpanCostSummaryByProjectDataLoader(DataLoader[Key, Result]):
@@ -106,12 +112,12 @@ def _get_stmt(
     segment: Segment,
     *params: Param,
 ) -> Select[Any]:
-    (start_time, end_time), filter_condition = segment
-    pid = models.Trace.project_rowid
+    project_rowids = params
+    (start_time, end_time), filter_condition, session_filter_condition = segment
 
     stmt: Select[Any] = (
         select(
-            pid,
+            models.Trace.project_rowid,
             coalesce(func.sum(models.SpanCost.prompt_cost), 0).label("prompt_cost"),
             coalesce(func.sum(models.SpanCost.completion_cost), 0).label("completion_cost"),
             coalesce(func.sum(models.SpanCost.total_cost), 0).label("total_cost"),
@@ -119,8 +125,10 @@ def _get_stmt(
             coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
             coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
         )
-        .join_from(models.SpanCost, models.Trace)
-        .group_by(pid)
+        .select_from(models.Trace)
+        .join(models.SpanCost, models.Trace.id == models.SpanCost.trace_rowid)
+        .where(models.Trace.project_rowid.in_(project_rowids))
+        .group_by(models.Trace.project_rowid)
     )
 
     if start_time:
@@ -132,7 +140,13 @@ def _get_stmt(
         sf = SpanFilter(filter_condition)
         stmt = sf(stmt.join_from(models.SpanCost, models.Span))
 
-    project_ids = [rowid for rowid in params]
-    stmt = stmt.where(pid.in_(project_ids))
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=project_rowids,
+            start_time=start_time,
+            end_time=end_time,
+        )
+        stmt = stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
 
     return stmt