arize-phoenix 4.5.0__py3-none-any.whl → 4.6.2__py3-none-any.whl

phoenix/server/api/dataloaders/experiment_run_counts.py

@@ -0,0 +1,42 @@
+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
+
+ExperimentID: TypeAlias = int
+RunCount: TypeAlias = int
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = RunCount
+
+
+class ExperimentRunCountsDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+    ) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        experiment_ids = keys
+        async with self._db() as session:
+            run_counts = {
+                experiment_id: run_count
+                async for experiment_id, run_count in await session.stream(
+                    select(models.ExperimentRun.experiment_id, func.count())
+                    .where(models.ExperimentRun.experiment_id.in_(set(experiment_ids)))
+                    .group_by(models.ExperimentRun.experiment_id)
+                )
+            }
+        return [
+            run_counts.get(experiment_id, ValueError(f"Unknown experiment: {experiment_id}"))
+            for experiment_id in experiment_ids
+        ]

phoenix/server/api/dataloaders/experiment_sequence_number.py

@@ -0,0 +1,49 @@
+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
+
+ExperimentId: TypeAlias = int
+Key: TypeAlias = ExperimentId
+Result: TypeAlias = Optional[int]
+
+
+class ExperimentSequenceNumberDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        experiment_ids = keys
+        dataset_ids = (
+            select(distinct(models.Experiment.dataset_id))
+            .where(models.Experiment.id.in_(experiment_ids))
+            .scalar_subquery()
+        )
+        row_number = (
+            func.row_number().over(
+                partition_by=models.Experiment.dataset_id,
+                order_by=models.Experiment.id,
+            )
+        ).label("row_number")
+        subq = (
+            select(models.Experiment.id, row_number)
+            .where(models.Experiment.dataset_id.in_(dataset_ids))
+            .subquery()
+        )
+        stmt = select(subq).where(subq.c.id.in_(experiment_ids))
+        async with self._db() as session:
+            result = {
+                experiment_id: sequence_number
+                async for experiment_id, sequence_number in await session.stream(stmt)
+            }
+        return [result.get(experiment_id) for experiment_id in experiment_ids]

phoenix/server/api/dataloaders/project_by_name.py

@@ -0,0 +1,31 @@
+from collections import defaultdict
+from typing import AsyncContextManager, Callable, 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
+
+ProjectName: TypeAlias = str
+Key: TypeAlias = ProjectName
+Result: TypeAlias = Optional[models.Project]
+
+
+class ProjectByNameDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        project_names = list(set(keys))
+        projects_by_name: DefaultDict[Key, Result] = defaultdict(None)
+        async with self._db() as session:
+            data = await session.stream_scalars(
+                select(models.Project).where(models.Project.name.in_(project_names))
+            )
+            async for project in data:
+                projects_by_name[project.name] = project
+
+        return [projects_by_name[project_name] for project_name in project_names]

phoenix/server/api/dataloaders/span_descendants.py

@@ -9,7 +9,7 @@ from typing import (
 from aioitertools.itertools import groupby
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import contains_eager
+from sqlalchemy.orm import joinedload
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -52,8 +52,7 @@ class SpanDescendantsDataLoader(DataLoader[Key, Result]):
         stmt = (
             select(descendant_ids.c[root_id_label], models.Span)
             .join(descendant_ids, models.Span.id == descendant_ids.c.id)
-            .join(models.Trace)
-            .options(contains_eager(models.Span.trace))
+            .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}

phoenix/server/api/dataloaders/span_projects.py

@@ -0,0 +1,33 @@
+from typing import AsyncContextManager, Callable, 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
+
+SpanID: TypeAlias = int
+Key: TypeAlias = SpanID
+Result: TypeAlias = models.Project
+
+
+class SpanProjectsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Union[Result, ValueError]]:
+        span_ids = list(set(keys))
+        async with self._db() as session:
+            projects = {
+                span_id: project
+                async for span_id, project in await session.stream(
+                    select(models.Span.id, models.Project)
+                    .select_from(models.Span)
+                    .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
+                    .join(models.Project, models.Trace.project_rowid == models.Project.id)
+                    .where(models.Span.id.in_(span_ids))
+                )
+            }
+        return [projects.get(span_id) or ValueError("Invalid span ID") for span_id in span_ids]

phoenix/server/api/dataloaders/trace_row_ids.py

@@ -0,0 +1,39 @@
+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
+
+TraceId: TypeAlias = str
+Key: TypeAlias = TraceId
+TraceRowId: TypeAlias = int
+ProjectRowId: TypeAlias = int
+Result: TypeAlias = Optional[Tuple[TraceRowId, ProjectRowId]]
+
+
+class TraceRowIdsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        stmt = select(
+            models.Trace.trace_id,
+            models.Trace.id,
+            models.Trace.project_rowid,
+        ).where(models.Trace.trace_id.in_(keys))
+        async with self._db() as session:
+            result = {
+                trace_id: (id_, project_rowid)
+                async for trace_id, id_, project_rowid in await session.stream(stmt)
+            }
+        return list(map(result.get, keys))

phoenix/server/api/helpers/dataset_helpers.py

@@ -0,0 +1,179 @@
+import json
+from typing import Any, Dict, Literal, Mapping, Optional, Protocol
+
+from openinference.semconv.trace import (
+    MessageAttributes,
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    ToolCallAttributes,
+)
+
+from phoenix.trace.attributes import get_attribute_value
+
+
+class HasSpanIO(Protocol):
+    """
+    An interface that contains the information needed to extract dataset example
+    input and output values from a span.
+    """
+
+    span_kind: Optional[str]
+    input_value: Any
+    input_mime_type: Optional[str]
+    output_value: Any
+    output_mime_type: Optional[str]
+    llm_prompt_template_variables: Any
+    llm_input_messages: Any
+    llm_output_messages: Any
+    retrieval_documents: 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
+    template variables (if present). For other span kinds, the input is
+    extracted from the input value and input mime type attributes.
+    """
+    input_value = span.input_value
+    input_mime_type = span.input_mime_type
+    if span.span_kind == OpenInferenceSpanKindValues.LLM.value:
+        return _get_llm_span_input(
+            input_messages=span.llm_input_messages,
+            input_value=input_value,
+            input_mime_type=input_mime_type,
+            prompt_template_variables=span.llm_prompt_template_variables,
+        )
+    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]:
+    """
+    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).
+    Output from retriever spans are extracted from the retrieval documents (if
+    present). For other span kinds, the output is extracted from the output
+    value and output mime type attributes.
+    """
+
+    output_value = span.output_value
+    output_mime_type = span.output_mime_type
+    if (span_kind := span.span_kind) == OpenInferenceSpanKindValues.LLM.value:
+        return _get_llm_span_output(
+            output_messages=span.llm_output_messages,
+            output_value=output_value,
+            output_mime_type=output_mime_type,
+        )
+    if span_kind == OpenInferenceSpanKindValues.RETRIEVER.value:
+        return _get_retriever_span_output(
+            retrieval_documents=span.retrieval_documents,
+            output_value=output_value,
+            output_mime_type=output_mime_type,
+        )
+    return _get_generic_io_value(io_value=output_value, mime_type=output_mime_type, kind="output")
+
+
+def _get_llm_span_input(
+    input_messages: Any,
+    input_value: Any,
+    input_mime_type: Optional[str],
+    prompt_template_variables: 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] = {}
+    if messages := [_get_message(m) for m in input_messages or ()]:
+        input["messages"] = messages
+    if not input:
+        input = _get_generic_io_value(io_value=input_value, mime_type=input_mime_type, kind="input")
+    if prompt_template_variables:
+        input = {**input, "prompt_template_variables": prompt_template_variables}
+    return input
+
+
+def _get_llm_span_output(
+    output_messages: Any,
+    output_value: Any,
+    output_mime_type: Optional[str],
+) -> 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).
+    """
+    if messages := [_get_message(m) for m in output_messages or ()]:
+        return {"messages": messages}
+    return _get_generic_io_value(io_value=output_value, mime_type=output_mime_type, kind="output")
+
+
+def _get_retriever_span_output(
+    retrieval_documents: Any,
+    output_value: Any,
+    output_mime_type: Optional[str],
+) -> 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).
+    """
+    if retrieval_documents is not None:
+        return {"documents": retrieval_documents}
+    return _get_generic_io_value(io_value=output_value, mime_type=output_mime_type, kind="output")
+
+
+def _get_generic_io_value(
+    io_value: Any, mime_type: Optional[str], kind: Literal["input", "output"]
+) -> Dict[str, Any]:
+    """
+    Makes a best-effort attempt to extract the input or output value from a span
+    and returns it as a dictionary.
+    """
+    if mime_type == OpenInferenceMimeTypeValues.JSON.value:
+        parsed_value = json.loads(io_value)
+        if isinstance(parsed_value, dict):
+            return parsed_value
+        else:
+            return {kind: parsed_value}
+    if isinstance(io_value, str):
+        return {kind: io_value}
+    return {}
+
+
+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)
+    function_call_arguments = get_attribute_value(message, MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON)
+    function_call = (
+        {"name": function_call_name, "arguments": function_call_arguments}
+        if function_call_name is not None or function_call_arguments is not None
+        else None
+    )
+    tool_calls = [
+        {
+            "function": {
+                "name": get_attribute_value(tool_call, TOOL_CALL_FUNCTION_NAME),
+                "arguments": get_attribute_value(tool_call, TOOL_CALL_FUNCTION_ARGUMENTS_JSON),
+            }
+        }
+        for tool_call in get_attribute_value(message, MESSAGE_TOOL_CALLS) or ()
+    ]
+    return {
+        "role": get_attribute_value(message, MESSAGE_ROLE),
+        **({"content": content} if content is not None else {}),
+        **({"name": name} if name is not None else {}),
+        **({"function_call": function_call} if function_call is not None else {}),
+        **({"tool_calls": tool_calls} if tool_calls else {}),
+    }
+
+
+MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
+MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON = MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON
+MESSAGE_FUNCTION_CALL_NAME = MessageAttributes.MESSAGE_FUNCTION_CALL_NAME
+MESSAGE_NAME = MessageAttributes.MESSAGE_NAME
+MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
+MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS
+
+TOOL_CALL_FUNCTION_NAME = ToolCallAttributes.TOOL_CALL_FUNCTION_NAME
+TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON

phoenix/server/api/input_types/AddExamplesToDatasetInput.py

@@ -0,0 +1,16 @@
+from typing import List, Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+from strawberry.scalars import JSON
+
+from .DatasetExampleInput import DatasetExampleInput
+
+
+@strawberry.input
+class AddExamplesToDatasetInput:
+    dataset_id: GlobalID
+    examples: List[DatasetExampleInput]
+    dataset_version_description: Optional[str] = UNSET
+    dataset_version_metadata: Optional[JSON] = UNSET

phoenix/server/api/input_types/AddSpansToDatasetInput.py

@@ -0,0 +1,14 @@
+from typing import List, Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+from strawberry.scalars import JSON
+
+
+@strawberry.input
+class AddSpansToDatasetInput:
+    dataset_id: GlobalID
+    span_ids: List[GlobalID]
+    dataset_version_description: Optional[str] = UNSET
+    dataset_version_metadata: Optional[JSON] = UNSET

phoenix/server/api/input_types/ClearProjectInput.py

@@ -0,0 +1,15 @@
+from datetime import datetime
+from typing import Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+
+
+@strawberry.input
+class ClearProjectInput:
+    id: GlobalID
+    end_time: Optional[datetime] = strawberry.field(
+        default=UNSET,
+        description="The time up to which to purge data. Time is right-open /non-inclusive.",
+    )

phoenix/server/api/input_types/CreateDatasetInput.py

@@ -0,0 +1,12 @@
+from typing import Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.scalars import JSON
+
+
+@strawberry.input
+class CreateDatasetInput:
+    name: str
+    description: Optional[str] = UNSET
+    metadata: Optional[JSON] = UNSET

phoenix/server/api/input_types/DatasetExampleInput.py

@@ -0,0 +1,14 @@
+from typing import Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+from strawberry.scalars import JSON
+
+
+@strawberry.input
+class DatasetExampleInput:
+    input: JSON
+    output: JSON
+    metadata: JSON
+    span_id: Optional[GlobalID] = UNSET

phoenix/server/api/input_types/DatasetSort.py

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

phoenix/server/api/input_types/DatasetVersionSort.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+import strawberry
+
+from phoenix.server.api.types.SortDir import SortDir
+
+
+@strawberry.enum
+class DatasetVersionColumn(Enum):
+    createdAt = "created_at"
+
+
+@strawberry.input(description="The sort key and direction for dataset version connections")
+class DatasetVersionSort:
+    col: DatasetVersionColumn
+    dir: SortDir

phoenix/server/api/input_types/DeleteDatasetExamplesInput.py

@@ -0,0 +1,13 @@
+from typing import List, Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+from strawberry.scalars import JSON
+
+
+@strawberry.input
+class DeleteDatasetExamplesInput:
+    example_ids: List[GlobalID]
+    dataset_version_description: Optional[str] = UNSET
+    dataset_version_metadata: Optional[JSON] = UNSET

phoenix/server/api/input_types/DeleteDatasetInput.py

@@ -0,0 +1,7 @@
+import strawberry
+from strawberry.relay import GlobalID
+
+
+@strawberry.input
+class DeleteDatasetInput:
+    dataset_id: GlobalID

phoenix/server/api/input_types/DeleteExperimentsInput.py

@@ -0,0 +1,9 @@
+from typing import List
+
+import strawberry
+from strawberry.relay import GlobalID
+
+
+@strawberry.input
+class DeleteExperimentsInput:
+    experiment_ids: List[GlobalID]

phoenix/server/api/input_types/PatchDatasetExamplesInput.py

@@ -0,0 +1,35 @@
+from typing import List, Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+from strawberry.scalars import JSON
+
+
+@strawberry.input
+class DatasetExamplePatch:
+    """
+    Contains the information needed to apply a patch revision to a dataset example.
+    """
+
+    example_id: GlobalID
+    input: Optional[JSON] = UNSET
+    output: Optional[JSON] = UNSET
+    metadata: Optional[JSON] = UNSET
+
+    def is_empty(self) -> bool:
+        """
+        Non-empty patches have at least one field set.
+        """
+        return all(field is UNSET for field in (self.input, self.output, self.metadata))
+
+
+@strawberry.input
+class PatchDatasetExamplesInput:
+    """
+    Input type to the patchDatasetExamples mutation.
+    """
+
+    patches: List[DatasetExamplePatch]
+    version_description: Optional[str] = UNSET
+    version_metadata: Optional[JSON] = UNSET

phoenix/server/api/input_types/PatchDatasetInput.py

@@ -0,0 +1,14 @@
+from typing import Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+from strawberry.scalars import JSON
+
+
+@strawberry.input
+class PatchDatasetInput:
+    dataset_id: GlobalID
+    name: Optional[str] = UNSET
+    description: Optional[str] = UNSET
+    metadata: Optional[JSON] = UNSET

phoenix/server/api/mutations/__init__.py

@@ -0,0 +1,13 @@
+import strawberry
+
+from phoenix.server.api.mutations.dataset_mutations import DatasetMutationMixin
+from phoenix.server.api.mutations.experiment_mutations import ExperimentMutationMixin
+from phoenix.server.api.mutations.export_events_mutations import ExportEventsMutationMixin
+from phoenix.server.api.mutations.project_mutations import ProjectMutationMixin
+
+
+@strawberry.type
+class Mutation(
+    ProjectMutationMixin, DatasetMutationMixin, ExperimentMutationMixin, ExportEventsMutationMixin
+):
+    pass

phoenix/server/api/mutations/auth.py

@@ -0,0 +1,11 @@
+from typing import Any
+
+from strawberry import Info
+from strawberry.permission import BasePermission
+
+
+class IsAuthenticated(BasePermission):
+    message = "User is not authenticated"
+
+    async def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        return not info.context.read_only