arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/dataloaders/trace_annotations_by_trace.py

@@ -0,0 +1,27 @@
+from collections import defaultdict
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db.models import TraceAnnotation
+from phoenix.server.types import DbSessionFactory
+
+TraceRowId: TypeAlias = int
+Key: TypeAlias = TraceRowId
+Result: TypeAlias = list[TraceAnnotation]
+
+
+class TraceAnnotationsByTraceDataLoader(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]:
+        annotations_by_id: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for annotation in await session.stream_scalars(
+                select(TraceAnnotation).where(TraceAnnotation.trace_rowid.in_(keys))
+            ):
+                annotations_by_id[annotation.trace_rowid].append(annotation)
+        return [annotations_by_id[key] for key in keys]

phoenix/server/api/dataloaders/types.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class CostBreakdown:
+    tokens: Optional[float] = None
+    cost: Optional[float] = None
+
+    @cached_property
+    def cost_per_token(self) -> Optional[float]:
+        if self.tokens and self.cost:
+            return self.cost / self.tokens
+        return None
+
+
+@dataclass(frozen=True)
+class SpanCostSummary:
+    prompt: CostBreakdown = field(default_factory=CostBreakdown)
+    completion: CostBreakdown = field(default_factory=CostBreakdown)
+    total: CostBreakdown = field(default_factory=CostBreakdown)
+
+
+@dataclass(frozen=True)
+class SpanCostDetailSummaryEntry:
+    token_type: str
+    is_prompt: bool
+    value: CostBreakdown = field(default_factory=CostBreakdown)

phoenix/server/api/exceptions.py

@@ -1,6 +1,8 @@
 from graphql.error import GraphQLError
 from strawberry.extensions import MaskErrors
 
+from phoenix.config import get_env_mask_internal_server_errors
+
 
 class CustomGraphQLError(Exception):
     """
@@ -27,6 +29,12 @@ class Unauthorized(CustomGraphQLError):
     """
 
 
+class InsufficientStorage(CustomGraphQLError):
+    """
+    An error raised when the database has insufficient storage to complete a request.
+    """
+
+
 class Conflict(CustomGraphQLError):
     """
     An error raised when a mutation cannot be completed due to a conflict with
@@ -45,4 +53,6 @@ def _should_mask_error(error: GraphQLError) -> bool:
     """
     Masks unexpected errors raised from GraphQL resolvers.
     """
-    return not isinstance(error.original_error, CustomGraphQLError)
+    return get_env_mask_internal_server_errors() and not isinstance(
+        error.original_error, CustomGraphQLError
+    )

phoenix/server/api/helpers/dataset_helpers.py

@@ -91,7 +91,9 @@ def _get_llm_span_input(
     if not input:
         input = _get_generic_io_value(io_value=input_value, mime_type=input_mime_type, kind="input")
     if prompt_template_variables_data := _safely_json_decode(prompt_template_variables):
-        input["prompt_template_variables"] = prompt_template_variables_data
+        # Hoist template variables to top level as individual key-value pairs
+        input.update(prompt_template_variables_data)
+        # Keep the original nested structure for compatibility
     if tool_definitions_data := [_safely_json_decode(tool_definition) for tool_definition in tools]:
         input["tools"] = tool_definitions_data
     return input
@@ -193,6 +195,8 @@ def _safely_json_decode(value: Any) -> Any:
     """
     Safely decodes a JSON-encoded value.
     """
+    if isinstance(value, dict):
+        return value
     if not isinstance(value, str):
         return None
     try: