arize-phoenix 5.11.0__py3-none-any.whl → 6.0.0__py3-none-any.whl

phoenix/server/api/mutations/chat_mutations.py

@@ -25,8 +25,10 @@ from typing_extensions import assert_never
 from phoenix.datetime_utils import local_now, normalize_datetime
 from phoenix.db import models
 from phoenix.db.helpers import get_dataset_example_revisions
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, CustomGraphQLError, NotFound
+from phoenix.server.api.helpers.dataset_helpers import get_dataset_example_output
 from phoenix.server.api.helpers.playground_clients import (
     PlaygroundStreamingClient,
     initialize_playground_clients,
@@ -61,12 +63,13 @@ from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Span import Span, to_gql_span
 from phoenix.server.api.types.TemplateLanguage import TemplateLanguage
 from phoenix.server.dml_event import SpanInsertEvent
-from phoenix.trace.attributes import get_attribute_value, unflatten
+from phoenix.trace.attributes import unflatten
 from phoenix.trace.schemas import SpanException
 from phoenix.utilities.json import jsonify
 from phoenix.utilities.template_formatters import (
     FStringTemplateFormatter,
     MustacheTemplateFormatter,
+    NoOpFormatter,
     TemplateFormatter,
 )
 
@@ -117,7 +120,7 @@ class ChatCompletionOverDatasetMutationPayload:
 
 @strawberry.type
 class ChatCompletionMutationMixin:
-    @strawberry.mutation
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     @classmethod
     async def chat_completion_over_dataset(
         cls,
@@ -242,7 +245,7 @@ class ChatCompletionMutationMixin:
                     dataset_example_id=revision.dataset_example_id,
                     trace_id=str(result.span.context.trace_id),
                     output=models.ExperimentRunOutput(
-                        task_output=get_attribute_value(db_span.attributes, LLM_OUTPUT_MESSAGES),
+                        task_output=get_dataset_example_output(db_span),
                     ),
                     prompt_token_count=db_span.cumulative_llm_token_count_prompt,
                     completion_token_count=db_span.cumulative_llm_token_count_completion,
@@ -272,7 +275,7 @@ class ChatCompletionMutationMixin:
             payload.examples.append(example_payload)
         return payload
 
-    @strawberry.mutation
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     @classmethod
     async def chat_completion(
         cls, info: Info[Context, None], input: ChatCompletionInput
@@ -481,6 +484,8 @@ def _template_formatter(template_language: TemplateLanguage) -> TemplateFormatte
         return MustacheTemplateFormatter()
     if template_language is TemplateLanguage.F_STRING:
         return FStringTemplateFormatter()
+    if template_language is TemplateLanguage.NONE:
+        return NoOpFormatter()
     assert_never(template_language)
 
 

phoenix/server/api/mutations/dataset_mutations.py

@@ -136,23 +136,7 @@ class DatasetMutationMixin:
                 .returning(models.DatasetVersion.id)
             )
             spans = (
-                await session.execute(
-                    select(
-                        models.Span.id,
-                        models.Span.span_kind,
-                        models.Span.attributes,
-                        _span_attribute(INPUT_MIME_TYPE),
-                        _span_attribute(INPUT_VALUE),
-                        _span_attribute(OUTPUT_MIME_TYPE),
-                        _span_attribute(OUTPUT_VALUE),
-                        _span_attribute(LLM_PROMPT_TEMPLATE_VARIABLES),
-                        _span_attribute(LLM_INPUT_MESSAGES),
-                        _span_attribute(LLM_OUTPUT_MESSAGES),
-                        _span_attribute(RETRIEVAL_DOCUMENTS),
-                    )
-                    .select_from(models.Span)
-                    .where(models.Span.id.in_(span_rowids))
-                )
+                await session.scalars(select(models.Span).where(models.Span.id.in_(span_rowids)))
             ).all()
             if missing_span_rowids := span_rowids - {span.id for span in spans}:
                 raise ValueError(
@@ -160,18 +144,10 @@ class DatasetMutationMixin:
                 )  # todo: implement error handling types https://github.com/Arize-ai/phoenix/issues/3221
 
             span_annotations = (
-                await session.execute(
-                    select(
-                        models.SpanAnnotation.span_rowid,
-                        models.SpanAnnotation.name,
-                        models.SpanAnnotation.label,
-                        models.SpanAnnotation.score,
-                        models.SpanAnnotation.explanation,
-                        models.SpanAnnotation.metadata_,
-                        models.SpanAnnotation.annotator_kind,
+                await session.scalars(
+                    select(models.SpanAnnotation).where(
+                        models.SpanAnnotation.span_rowid.in_(span_rowids)
                     )
-                    .select_from(models.SpanAnnotation)
-                    .where(models.SpanAnnotation.span_rowid.in_(span_rowids))
                 )
             ).all()
 
@@ -214,8 +190,12 @@ class DatasetMutationMixin:
                         DatasetExampleRevision.input.key: get_dataset_example_input(span),
                         DatasetExampleRevision.output.key: get_dataset_example_output(span),
                         DatasetExampleRevision.metadata_.key: {
-                            **span.attributes,
-                            "annotations": span_annotations_by_span[span.id],
+                            "span_kind": span.span_kind,
+                            **(
+                                {"annotations": annotations}
+                                if (annotations := span_annotations_by_span[span.id])
+                                else {}
+                            ),
                         },
                         DatasetExampleRevision.revision_kind.key: "CREATE",
                     }

phoenix/server/api/routers/v1/evaluations.py

@@ -40,7 +40,7 @@ from .utils import add_errors_to_responses
 
 EvaluationName: TypeAlias = str
 
-router = APIRouter(tags=["traces"], include_in_schema=False)
+router = APIRouter(tags=["traces"], include_in_schema=True)
 
 
 @router.post(

phoenix/server/api/subscriptions.py

@@ -1,11 +1,11 @@
 import asyncio
 import logging
-from asyncio import FIRST_COMPLETED, Queue, QueueEmpty, Task, create_task, wait, wait_for
 from collections.abc import AsyncIterator, Iterator
 from datetime import datetime, timedelta, timezone
 from typing import (
     Any,
     AsyncGenerator,
+    Coroutine,
     Iterable,
     Mapping,
     Optional,
@@ -24,6 +24,7 @@ from typing_extensions import TypeAlias, assert_never
 
 from phoenix.datetime_utils import local_now, normalize_datetime
 from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, CustomGraphQLError, NotFound
 from phoenix.server.api.helpers.playground_clients import (
@@ -63,6 +64,7 @@ from phoenix.server.types import DbSessionFactory
 from phoenix.utilities.template_formatters import (
     FStringTemplateFormatter,
     MustacheTemplateFormatter,
+    NoOpFormatter,
     TemplateFormatter,
     TemplateFormatterError,
 )
@@ -86,7 +88,7 @@ PLAYGROUND_PROJECT_NAME = "playground"
 
 @strawberry.type
 class Subscription:
-    @strawberry.subscription
+    @strawberry.subscription(permission_classes=[IsNotReadOnly])  # type: ignore
     async def chat_completion(
         self, info: Info[Context, None], input: ChatCompletionInput
     ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
@@ -161,7 +163,7 @@ class Subscription:
         info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
         yield ChatCompletionSubscriptionResult(span=to_gql_span(db_span))
 
-    @strawberry.subscription
+    @strawberry.subscription(permission_classes=[IsNotReadOnly])  # type: ignore
     async def chat_completion_over_dataset(
         self, info: Info[Context, None], input: ChatCompletionOverDatasetInput
     ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
@@ -286,7 +288,7 @@ class Subscription:
             experiment=to_gql_experiment(experiment)
         )  # eagerly yields experiment so it can be linked by consumers of the subscription
 
-        results: Queue[ChatCompletionResult] = Queue()
+        results: asyncio.Queue[ChatCompletionResult] = asyncio.Queue()
         not_started: list[tuple[DatasetExampleID, ChatStream]] = [
             (
                 GlobalID(DatasetExample.__name__, str(revision.dataset_example_id)),
@@ -302,7 +304,11 @@ class Subscription:
             for revision in revisions
         ]
         in_progress: list[
-            tuple[Optional[DatasetExampleID], ChatStream, Task[ChatCompletionSubscriptionPayload]]
+            tuple[
+                Optional[DatasetExampleID],
+                ChatStream,
+                asyncio.Task[ChatCompletionSubscriptionPayload],
+            ]
         ] = []
         max_in_progress = 3
         write_batch_size = 10
@@ -314,7 +320,9 @@ class Subscription:
                 task = _create_task_with_timeout(stream)
                 in_progress.append((ex_id, stream, task))
             async_tasks_to_run = [task for _, _, task in in_progress]
-            completed_tasks, _ = await wait(async_tasks_to_run, return_when=FIRST_COMPLETED)
+            completed_tasks, _ = await asyncio.wait(
+                async_tasks_to_run, return_when=asyncio.FIRST_COMPLETED
+            )
             for completed_task in completed_tasks:
                 idx = [task for _, _, task in in_progress].index(completed_task)
                 example_id, stream, _ = in_progress[idx]
@@ -326,7 +334,7 @@ class Subscription:
                     del in_progress[idx]  # removes timed-out stream
                     if example_id is not None:
                         yield ChatCompletionSubscriptionError(
-                            message="Timed out", dataset_example_id=example_id
+                            message="Playground task timed out", dataset_example_id=example_id
                         )
                 except Exception as error:
                     del in_progress[idx]  # removes failed stream
@@ -367,7 +375,7 @@ async def _stream_chat_completion_over_dataset_example(
     input: ChatCompletionOverDatasetInput,
     llm_client: PlaygroundStreamingClient,
     revision: models.DatasetExampleRevision,
-    results: Queue[ChatCompletionResult],
+    results: asyncio.Queue[ChatCompletionResult],
     experiment_id: int,
     project_id: int,
 ) -> ChatStream:
@@ -469,23 +477,52 @@ def _is_result_payloads_stream(
 
 def _create_task_with_timeout(
     iterable: AsyncIterator[GenericType], timeout_in_seconds: int = 90
-) -> Task[GenericType]:
-    return create_task(wait_for(_as_coroutine(iterable), timeout=timeout_in_seconds))
+) -> asyncio.Task[GenericType]:
+    return asyncio.create_task(
+        _wait_for(
+            _as_coroutine(iterable),
+            timeout=timeout_in_seconds,
+            timeout_message="Playground task timed out",
+        )
+    )
+
+
+async def _wait_for(
+    coro: Coroutine[None, None, GenericType],
+    timeout: float,
+    timeout_message: Optional[str] = None,
+) -> GenericType:
+    """
+    A function that imitates asyncio.wait_for, but allows the task to be
+    cancelled with a custom message.
+    """
+    task = asyncio.create_task(coro)
+    done, pending = await asyncio.wait([task], timeout=timeout)
+    assert len(done) + len(pending) == 1
+    if done:
+        task = done.pop()
+        return task.result()
+    task = pending.pop()
+    task.cancel(msg=timeout_message)
+    try:
+        return await task
+    except asyncio.CancelledError:
+        raise asyncio.TimeoutError()
 
 
-async def _drain(queue: Queue[GenericType]) -> list[GenericType]:
+async def _drain(queue: asyncio.Queue[GenericType]) -> list[GenericType]:
     values: list[GenericType] = []
     while not queue.empty():
         values.append(await queue.get())
     return values
 
 
-def _drain_no_wait(queue: Queue[GenericType]) -> list[GenericType]:
+def _drain_no_wait(queue: asyncio.Queue[GenericType]) -> list[GenericType]:
     values: list[GenericType] = []
     while True:
         try:
             values.append(queue.get_nowait())
-        except QueueEmpty:
+        except asyncio.QueueEmpty:
             break
     return values
 
@@ -526,6 +563,8 @@ def _template_formatter(template_language: TemplateLanguage) -> TemplateFormatte
         return MustacheTemplateFormatter()
     if template_language is TemplateLanguage.F_STRING:
         return FStringTemplateFormatter()
+    if template_language is TemplateLanguage.NONE:
+        return NoOpFormatter()
     assert_never(template_language)
 
 

phoenix/server/api/types/Span.py

@@ -1,13 +1,12 @@
 import json
 from collections.abc import Mapping, Sized
-from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Optional, cast
 
 import numpy as np
 import strawberry
-from openinference.semconv.trace import EmbeddingAttributes, SpanAttributes
+from openinference.semconv.trace import SpanAttributes
 from strawberry import ID, UNSET
 from strawberry.relay import Node, NodeID
 from strawberry.types import Info
@@ -39,18 +38,6 @@ from .SpanAnnotation import SpanAnnotation
 if TYPE_CHECKING:
     from phoenix.server.api.types.Project import Project
 
-EMBEDDING_EMBEDDINGS = SpanAttributes.EMBEDDING_EMBEDDINGS
-EMBEDDING_VECTOR = EmbeddingAttributes.EMBEDDING_VECTOR
-INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
-INPUT_VALUE = SpanAttributes.INPUT_VALUE
-LLM_PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
-LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
-LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
-METADATA = SpanAttributes.METADATA
-OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
-OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
-RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS
-
 
 @strawberry.enum
 class SpanKind(Enum):
@@ -236,21 +223,7 @@ class Span(Node):
         description="The span's attributes translated into an example revision for a dataset",
     )  # type: ignore
     async def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleRevision:
-        db_span = self.db_span
-        attributes = db_span.attributes
-        span_io = _SpanIO(
-            span_kind=db_span.span_kind,
-            input_value=get_attribute_value(attributes, INPUT_VALUE),
-            input_mime_type=get_attribute_value(attributes, INPUT_MIME_TYPE),
-            output_value=get_attribute_value(attributes, OUTPUT_VALUE),
-            output_mime_type=get_attribute_value(attributes, OUTPUT_MIME_TYPE),
-            llm_prompt_template_variables=get_attribute_value(
-                attributes, LLM_PROMPT_TEMPLATE_VARIABLES
-            ),
-            llm_input_messages=get_attribute_value(attributes, LLM_INPUT_MESSAGES),
-            llm_output_messages=get_attribute_value(attributes, LLM_OUTPUT_MESSAGES),
-            retrieval_documents=get_attribute_value(attributes, RETRIEVAL_DOCUMENTS),
-        )
+        span = self.db_span
 
         # Fetch annotations associated with this span
         span_annotations = await self.span_annotations(info)
@@ -265,13 +238,13 @@ class Span(Node):
             }
         # Merge annotations into the metadata
         metadata = {
-            **attributes,
-            "annotations": annotations,
+            "span_kind": span.span_kind,
+            **({"annotations": annotations} if annotations else {}),
         }
 
         return SpanAsExampleRevision(
-            input=get_dataset_example_input(span_io),
-            output=get_dataset_example_output(span_io),
+            input=get_dataset_example_input(span),
+            output=get_dataset_example_output(span),
             metadata=metadata,
         )
 
@@ -435,19 +408,9 @@ def _convert_metadata_to_string(metadata: Any) -> Optional[str]:
         return str(metadata)
 
 
-@dataclass
-class _SpanIO:
-    """
-    An class 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
+INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
+INPUT_VALUE = SpanAttributes.INPUT_VALUE
+METADATA = SpanAttributes.METADATA
+OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
+RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS

phoenix/server/api/types/TemplateLanguage.py

@@ -5,5 +5,6 @@ import strawberry
 
 @strawberry.enum
 class TemplateLanguage(Enum):
+    NONE = "NONE"
     MUSTACHE = "MUSTACHE"
     F_STRING = "F_STRING"

phoenix/server/static/.vite/manifest.json

@@ -1,32 +1,32 @@
 {
-  "_components-C_HASv83.js": {
-    "file": "assets/components-C_HASv83.js",
+  "_components-Bo-xjXoV.js": {
+    "file": "assets/components-Bo-xjXoV.js",
     "name": "components",
     "imports": [
-      "_vendor-BCxsh5i3.js",
-      "_pages-DYHcAdjT.js",
-      "_vendor-arizeai-C2CDZgMz.js",
-      "_vendor-codemirror-DYbtnCTn.js",
+      "_vendor-UNccFYqq.js",
+      "_pages-BfYVoiUE.js",
+      "_vendor-arizeai-gSIqhzZY.js",
+      "_vendor-codemirror-fSYjW3f-.js",
       "_vendor-three-DwGkEfCM.js"
     ]
   },
-  "_pages-DYHcAdjT.js": {
-    "file": "assets/pages-DYHcAdjT.js",
+  "_pages-BfYVoiUE.js": {
+    "file": "assets/pages-BfYVoiUE.js",
     "name": "pages",
     "imports": [
-      "_vendor-BCxsh5i3.js",
-      "_vendor-arizeai-C2CDZgMz.js",
-      "_components-C_HASv83.js",
-      "_vendor-recharts-P6W8G0Mb.js",
-      "_vendor-codemirror-DYbtnCTn.js"
+      "_vendor-UNccFYqq.js",
+      "_vendor-arizeai-gSIqhzZY.js",
+      "_components-Bo-xjXoV.js",
+      "_vendor-recharts-CVVyA2X0.js",
+      "_vendor-codemirror-fSYjW3f-.js"
     ]
   },
   "_vendor-!~{003}~.js": {
     "file": "assets/vendor-DxkFTwjz.css",
     "src": "_vendor-!~{003}~.js"
   },
-  "_vendor-BCxsh5i3.js": {
-    "file": "assets/vendor-BCxsh5i3.js",
+  "_vendor-UNccFYqq.js": {
+    "file": "assets/vendor-UNccFYqq.js",
     "name": "vendor",
     "imports": [
       "_vendor-three-DwGkEfCM.js"
@@ -35,25 +35,25 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-C2CDZgMz.js": {
-    "file": "assets/vendor-arizeai-C2CDZgMz.js",
+  "_vendor-arizeai-gSIqhzZY.js": {
+    "file": "assets/vendor-arizeai-gSIqhzZY.js",
     "name": "vendor-arizeai",
     "imports": [
-      "_vendor-BCxsh5i3.js"
+      "_vendor-UNccFYqq.js"
     ]
   },
-  "_vendor-codemirror-DYbtnCTn.js": {
-    "file": "assets/vendor-codemirror-DYbtnCTn.js",
+  "_vendor-codemirror-fSYjW3f-.js": {
+    "file": "assets/vendor-codemirror-fSYjW3f-.js",
     "name": "vendor-codemirror",
     "imports": [
-      "_vendor-BCxsh5i3.js"
+      "_vendor-UNccFYqq.js"
     ]
   },
-  "_vendor-recharts-P6W8G0Mb.js": {
-    "file": "assets/vendor-recharts-P6W8G0Mb.js",
+  "_vendor-recharts-CVVyA2X0.js": {
+    "file": "assets/vendor-recharts-CVVyA2X0.js",
     "name": "vendor-recharts",
     "imports": [
-      "_vendor-BCxsh5i3.js"
+      "_vendor-UNccFYqq.js"
     ]
   },
   "_vendor-three-DwGkEfCM.js": {
@@ -61,18 +61,18 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-D7UiCRtr.js",
+    "file": "assets/index-i2KbnOZd.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
-      "_vendor-BCxsh5i3.js",
-      "_vendor-arizeai-C2CDZgMz.js",
-      "_pages-DYHcAdjT.js",
-      "_components-C_HASv83.js",
+      "_vendor-UNccFYqq.js",
+      "_vendor-arizeai-gSIqhzZY.js",
+      "_pages-BfYVoiUE.js",
+      "_components-Bo-xjXoV.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-recharts-P6W8G0Mb.js",
-      "_vendor-codemirror-DYbtnCTn.js"
+      "_vendor-recharts-CVVyA2X0.js",
+      "_vendor-codemirror-fSYjW3f-.js"
     ]
   }
 }