arize-phoenix 5.6.0__py3-none-any.whl → 5.8.0__py3-none-any.whl

phoenix/server/api/subscriptions.py

@@ -1,419 +1,85 @@
-import json
-from abc import ABC, abstractmethod
-from collections import defaultdict
-from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping
-from dataclasses import asdict
+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, timezone
-from enum import Enum
-from itertools import chain
-from traceback import format_exc
-from typing import TYPE_CHECKING, Annotated, Any, Optional, Union, cast
+from typing import (
+    Any,
+    Iterable,
+    Mapping,
+    Optional,
+    Sequence,
+    TypeVar,
+    cast,
+)
 
 import strawberry
-from openinference.instrumentation import safe_json_dumps
-from openinference.semconv.trace import (
-    MessageAttributes,
-    OpenInferenceMimeTypeValues,
-    OpenInferenceSpanKindValues,
-    SpanAttributes,
-    ToolAttributes,
-    ToolCallAttributes,
-)
-from opentelemetry.sdk.trace.id_generator import RandomIdGenerator as DefaultOTelIDGenerator
-from opentelemetry.trace import StatusCode
-from sqlalchemy import insert, select
-from strawberry import UNSET
-from strawberry.scalars import JSON as JSONScalarType
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import and_, func, insert, select
+from sqlalchemy.orm import load_only
+from strawberry.relay.types import GlobalID
 from strawberry.types import Info
 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.context import Context
-from phoenix.server.api.exceptions import BadRequest
-from phoenix.server.api.input_types.ChatCompletionMessageInput import ChatCompletionMessageInput
-from phoenix.server.api.input_types.InvocationParameters import InvocationParameters
+from phoenix.server.api.exceptions import BadRequest, NotFound
+from phoenix.server.api.helpers.playground_clients import (
+    PlaygroundStreamingClient,
+    initialize_playground_clients,
+)
+from phoenix.server.api.helpers.playground_registry import (
+    PLAYGROUND_CLIENT_REGISTRY,
+)
+from phoenix.server.api.helpers.playground_spans import (
+    get_db_experiment_run,
+    get_db_span,
+    get_db_trace,
+    streaming_llm_span,
+)
+from phoenix.server.api.input_types.ChatCompletionInput import (
+    ChatCompletionInput,
+    ChatCompletionOverDatasetInput,
+)
 from phoenix.server.api.types.ChatCompletionMessageRole import ChatCompletionMessageRole
-from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
-from phoenix.server.api.types.Span import Span, to_gql_span
-from phoenix.server.dml_event import SpanInsertEvent
-from phoenix.trace.attributes import unflatten
-from phoenix.trace.schemas import (
-    SpanEvent,
-    SpanException,
+from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
+    ChatCompletionSubscriptionError,
+    ChatCompletionSubscriptionExperiment,
+    ChatCompletionSubscriptionPayload,
+    ChatCompletionSubscriptionResult,
 )
-from phoenix.utilities.json import jsonify
+from phoenix.server.api.types.Dataset import Dataset
+from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetVersion import DatasetVersion
+from phoenix.server.api.types.Experiment import to_gql_experiment
+from phoenix.server.api.types.ExperimentRun import to_gql_experiment_run
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.Span import to_gql_span
+from phoenix.server.api.types.TemplateLanguage import TemplateLanguage
+from phoenix.server.dml_event import SpanInsertEvent
+from phoenix.server.types import DbSessionFactory
 from phoenix.utilities.template_formatters import (
     FStringTemplateFormatter,
     MustacheTemplateFormatter,
     TemplateFormatter,
+    TemplateFormatterError,
 )
 
-if TYPE_CHECKING:
-    from anthropic.types import MessageParam
-    from openai.types import CompletionUsage
-    from openai.types.chat import (
-        ChatCompletionMessageParam,
-        ChatCompletionMessageToolCallParam,
-    )
-
-PLAYGROUND_PROJECT_NAME = "playground"
-
-ToolCallID: TypeAlias = str
-SetSpanAttributesFn: TypeAlias = Callable[[Mapping[str, Any]], None]
-
-
-@strawberry.enum
-class TemplateLanguage(Enum):
-    MUSTACHE = "MUSTACHE"
-    F_STRING = "F_STRING"
-
-
-@strawberry.input
-class TemplateOptions:
-    variables: JSONScalarType
-    language: TemplateLanguage
-
-
-@strawberry.type
-class TextChunk:
-    content: str
-
-
-@strawberry.type
-class FunctionCallChunk:
-    name: str
-    arguments: str
-
-
-@strawberry.type
-class ToolCallChunk:
-    id: str
-    function: FunctionCallChunk
-
+GenericType = TypeVar("GenericType")
 
-@strawberry.type
-class ChatCompletionSubscriptionError:
-    message: str
+logger = logging.getLogger(__name__)
 
+initialize_playground_clients()
 
-@strawberry.type
-class FinishedChatCompletion:
-    span: Span
-
-
-ChatCompletionChunk: TypeAlias = Union[TextChunk, ToolCallChunk]
-
-ChatCompletionSubscriptionPayload: TypeAlias = Annotated[
-    Union[TextChunk, ToolCallChunk, FinishedChatCompletion, ChatCompletionSubscriptionError],
-    strawberry.union("ChatCompletionSubscriptionPayload"),
+ChatCompletionMessage: TypeAlias = tuple[
+    ChatCompletionMessageRole, str, Optional[str], Optional[list[str]]
 ]
-
-
-@strawberry.input
-class GenerativeModelInput:
-    provider_key: GenerativeProviderKey
-    name: str
-    """ The name of the model. Or the Deployment Name for Azure OpenAI models. """
-    endpoint: Optional[str] = UNSET
-    """ The endpoint to use for the model. Only required for Azure OpenAI models. """
-    api_version: Optional[str] = UNSET
-    """ The API version to use for the model. """
-
-
-@strawberry.input
-class ChatCompletionInput:
-    messages: list[ChatCompletionMessageInput]
-    model: GenerativeModelInput
-    invocation_parameters: InvocationParameters = strawberry.field(default_factory=dict)
-    tools: Optional[list[JSONScalarType]] = UNSET
-    template: Optional[TemplateOptions] = UNSET
-    api_key: Optional[str] = strawberry.field(default=None)
-
-
-PLAYGROUND_STREAMING_CLIENT_REGISTRY: dict[
-    GenerativeProviderKey, type["PlaygroundStreamingClient"]
-] = {}
-
-
-def register_llm_client(
-    provider_key: GenerativeProviderKey,
-) -> Callable[[type["PlaygroundStreamingClient"]], type["PlaygroundStreamingClient"]]:
-    def decorator(cls: type["PlaygroundStreamingClient"]) -> type["PlaygroundStreamingClient"]:
-        PLAYGROUND_STREAMING_CLIENT_REGISTRY[provider_key] = cls
-        return cls
-
-    return decorator
-
-
-class PlaygroundStreamingClient(ABC):
-    def __init__(
-        self,
-        model: GenerativeModelInput,
-        api_key: Optional[str] = None,
-        set_span_attributes: Optional[SetSpanAttributesFn] = None,
-    ) -> None:
-        self._set_span_attributes = set_span_attributes
-
-    @abstractmethod
-    async def chat_completion_create(
-        self,
-        messages: list[
-            tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[JSONScalarType]]]
-        ],
-        tools: list[JSONScalarType],
-        **invocation_parameters: Any,
-    ) -> AsyncIterator[ChatCompletionChunk]:
-        # a yield statement is needed to satisfy the type-checker
-        # https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
-        yield TextChunk(content="")
-
-
-@register_llm_client(GenerativeProviderKey.OPENAI)
-class OpenAIStreamingClient(PlaygroundStreamingClient):
-    def __init__(
-        self,
-        model: GenerativeModelInput,
-        api_key: Optional[str] = None,
-        set_span_attributes: Optional[SetSpanAttributesFn] = None,
-    ) -> None:
-        from openai import AsyncOpenAI
-
-        super().__init__(model=model, api_key=api_key, set_span_attributes=set_span_attributes)
-        self.client = AsyncOpenAI(api_key=api_key)
-        self.model_name = model.name
-
-    async def chat_completion_create(
-        self,
-        messages: list[
-            tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[JSONScalarType]]]
-        ],
-        tools: list[JSONScalarType],
-        **invocation_parameters: Any,
-    ) -> AsyncIterator[ChatCompletionChunk]:
-        from openai import NOT_GIVEN
-        from openai.types.chat import ChatCompletionStreamOptionsParam
-
-        # Convert standard messages to OpenAI messages
-        openai_messages = [self.to_openai_chat_completion_param(*message) for message in messages]
-        tool_call_ids: dict[int, str] = {}
-        token_usage: Optional["CompletionUsage"] = None
-        async for chunk in await self.client.chat.completions.create(
-            messages=openai_messages,
-            model=self.model_name,
-            stream=True,
-            stream_options=ChatCompletionStreamOptionsParam(include_usage=True),
-            tools=tools or NOT_GIVEN,
-            **invocation_parameters,
-        ):
-            if (usage := chunk.usage) is not None:
-                token_usage = usage
-                continue
-            choice = chunk.choices[0]
-            delta = choice.delta
-            if choice.finish_reason is None:
-                if isinstance(chunk_content := delta.content, str):
-                    text_chunk = TextChunk(content=chunk_content)
-                    yield text_chunk
-                if (tool_calls := delta.tool_calls) is not None:
-                    for tool_call_index, tool_call in enumerate(tool_calls):
-                        tool_call_id = (
-                            tool_call.id
-                            if tool_call.id is not None
-                            else tool_call_ids[tool_call_index]
-                        )
-                        tool_call_ids[tool_call_index] = tool_call_id
-                        if (function := tool_call.function) is not None:
-                            tool_call_chunk = ToolCallChunk(
-                                id=tool_call_id,
-                                function=FunctionCallChunk(
-                                    name=function.name or "",
-                                    arguments=function.arguments or "",
-                                ),
-                            )
-                            yield tool_call_chunk
-        if token_usage is not None and self._set_span_attributes:
-            self._set_span_attributes(dict(self._llm_token_counts(token_usage)))
-
-    def to_openai_chat_completion_param(
-        self,
-        role: ChatCompletionMessageRole,
-        content: JSONScalarType,
-        tool_call_id: Optional[str] = None,
-        tool_calls: Optional[list[JSONScalarType]] = None,
-    ) -> "ChatCompletionMessageParam":
-        from openai.types.chat import (
-            ChatCompletionAssistantMessageParam,
-            ChatCompletionSystemMessageParam,
-            ChatCompletionToolMessageParam,
-            ChatCompletionUserMessageParam,
-        )
-
-        if role is ChatCompletionMessageRole.USER:
-            return ChatCompletionUserMessageParam(
-                {
-                    "content": content,
-                    "role": "user",
-                }
-            )
-        if role is ChatCompletionMessageRole.SYSTEM:
-            return ChatCompletionSystemMessageParam(
-                {
-                    "content": content,
-                    "role": "system",
-                }
-            )
-        if role is ChatCompletionMessageRole.AI:
-            if tool_calls is None:
-                return ChatCompletionAssistantMessageParam(
-                    {
-                        "content": content,
-                        "role": "assistant",
-                    }
-                )
-            else:
-                return ChatCompletionAssistantMessageParam(
-                    {
-                        "content": content,
-                        "role": "assistant",
-                        "tool_calls": [
-                            self.to_openai_tool_call_param(tool_call) for tool_call in tool_calls
-                        ],
-                    }
-                )
-        if role is ChatCompletionMessageRole.TOOL:
-            if tool_call_id is None:
-                raise ValueError("tool_call_id is required for tool messages")
-        return ChatCompletionToolMessageParam(
-            {"content": content, "role": "tool", "tool_call_id": tool_call_id}
-        )
-        assert_never(role)
-
-    def to_openai_tool_call_param(
-        self,
-        tool_call: JSONScalarType,
-    ) -> "ChatCompletionMessageToolCallParam":
-        from openai.types.chat import ChatCompletionMessageToolCallParam
-
-        return ChatCompletionMessageToolCallParam(
-            id=tool_call.get("id", ""),
-            function={
-                "name": tool_call.get("function", {}).get("name", ""),
-                "arguments": safe_json_dumps(tool_call.get("function", {}).get("arguments", "")),
-            },
-            type="function",
-        )
-
-    @staticmethod
-    def _llm_token_counts(usage: "CompletionUsage") -> Iterator[tuple[str, Any]]:
-        yield LLM_TOKEN_COUNT_PROMPT, usage.prompt_tokens
-        yield LLM_TOKEN_COUNT_COMPLETION, usage.completion_tokens
-        yield LLM_TOKEN_COUNT_TOTAL, usage.total_tokens
-
-
-@register_llm_client(GenerativeProviderKey.AZURE_OPENAI)
-class AzureOpenAIStreamingClient(OpenAIStreamingClient):
-    def __init__(
-        self,
-        model: GenerativeModelInput,
-        api_key: Optional[str] = None,
-        set_span_attributes: Optional[SetSpanAttributesFn] = None,
-    ):
-        from openai import AsyncAzureOpenAI
-
-        super().__init__(model=model, api_key=api_key, set_span_attributes=set_span_attributes)
-        if model.endpoint is None or model.api_version is None:
-            raise ValueError("endpoint and api_version are required for Azure OpenAI models")
-        self.client = AsyncAzureOpenAI(
-            api_key=api_key,
-            azure_endpoint=model.endpoint,
-            api_version=model.api_version,
-        )
-
-
-@register_llm_client(GenerativeProviderKey.ANTHROPIC)
-class AnthropicStreamingClient(PlaygroundStreamingClient):
-    def __init__(
-        self,
-        model: GenerativeModelInput,
-        api_key: Optional[str] = None,
-        set_span_attributes: Optional[SetSpanAttributesFn] = None,
-    ) -> None:
-        import anthropic
-
-        super().__init__(model=model, api_key=api_key, set_span_attributes=set_span_attributes)
-        self.client = anthropic.AsyncAnthropic(api_key=api_key)
-        self.model_name = model.name
-
-    async def chat_completion_create(
-        self,
-        messages: list[
-            tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[JSONScalarType]]]
-        ],
-        tools: list[JSONScalarType],
-        **invocation_parameters: Any,
-    ) -> AsyncIterator[ChatCompletionChunk]:
-        import anthropic.lib.streaming as anthropic_streaming
-        import anthropic.types as anthropic_types
-
-        anthropic_messages, system_prompt = self._build_anthropic_messages(messages)
-
-        anthropic_params = {
-            "messages": anthropic_messages,
-            "model": self.model_name,
-            "system": system_prompt,
-            "max_tokens": 1024,
-            **invocation_parameters,
-        }
-        async with self.client.messages.stream(**anthropic_params) as stream:
-            async for event in stream:
-                if isinstance(event, anthropic_types.RawMessageStartEvent):
-                    if self._set_span_attributes:
-                        self._set_span_attributes(
-                            {LLM_TOKEN_COUNT_PROMPT: event.message.usage.input_tokens}
-                        )
-                elif isinstance(event, anthropic_streaming.TextEvent):
-                    yield TextChunk(content=event.text)
-                elif isinstance(event, anthropic_streaming.MessageStopEvent):
-                    if self._set_span_attributes:
-                        self._set_span_attributes(
-                            {LLM_TOKEN_COUNT_COMPLETION: event.message.usage.output_tokens}
-                        )
-                elif isinstance(
-                    event,
-                    (
-                        anthropic_types.RawContentBlockStartEvent,
-                        anthropic_types.RawContentBlockDeltaEvent,
-                        anthropic_types.RawMessageDeltaEvent,
-                        anthropic_streaming.ContentBlockStopEvent,
-                    ),
-                ):
-                    # event types emitted by the stream that don't contain useful information
-                    pass
-                elif isinstance(event, anthropic_streaming.InputJsonEvent):
-                    raise NotImplementedError
-                else:
-                    assert_never(event)
-
-    def _build_anthropic_messages(
-        self,
-        messages: list[tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[str]]]],
-    ) -> tuple[list["MessageParam"], str]:
-        anthropic_messages: list["MessageParam"] = []
-        system_prompt = ""
-        for role, content, _tool_call_id, _tool_calls in messages:
-            if role == ChatCompletionMessageRole.USER:
-                anthropic_messages.append({"role": "user", "content": content})
-            elif role == ChatCompletionMessageRole.AI:
-                anthropic_messages.append({"role": "assistant", "content": content})
-            elif role == ChatCompletionMessageRole.SYSTEM:
-                system_prompt += content + "\n"
-            elif role == ChatCompletionMessageRole.TOOL:
-                raise NotImplementedError
-            else:
-                assert_never(role)
-
-        return anthropic_messages, system_prompt
+DatasetExampleID: TypeAlias = GlobalID
+ChatCompletionResult: TypeAlias = tuple[
+    DatasetExampleID, Optional[models.Span], models.ExperimentRun
+]
+PLAYGROUND_PROJECT_NAME = "playground"
 
 
 @strawberry.type
@@ -422,15 +88,15 @@ class Subscription:
     async def chat_completion(
         self, info: Info[Context, None], input: ChatCompletionInput
     ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
-        # Determine which LLM client to use based on provider_key
         provider_key = input.model.provider_key
-        if (llm_client_class := PLAYGROUND_STREAMING_CLIENT_REGISTRY.get(provider_key)) is None:
+        llm_client_class = PLAYGROUND_CLIENT_REGISTRY.get_client(provider_key, input.model.name)
+        if llm_client_class is None:
             raise BadRequest(f"No LLM client registered for provider '{provider_key}'")
         llm_client = llm_client_class(
             model=input.model,
             api_key=input.api_key,
-            set_span_attributes=lambda attrs: attributes.update(attrs),
         )
+
         messages = [
             (
                 message.role,
@@ -441,70 +107,129 @@ class Subscription:
             for message in input.messages
         ]
         if template_options := input.template:
-            messages = list(_formatted_messages(messages, template_options))
-        invocation_parameters = jsonify(input.invocation_parameters)
-        attributes = dict(
-            chain(
-                _llm_span_kind(),
-                _llm_model_name(input.model.name),
-                _llm_tools(input.tools or []),
-                _llm_input_messages(messages),
-                _llm_invocation_parameters(invocation_parameters),
-                _input_value_and_mime_type(input),
+            messages = list(
+                _formatted_messages(
+                    messages=messages,
+                    template_language=template_options.language,
+                    template_variables=template_options.variables,
+                )
             )
+        invocation_parameters = llm_client.construct_invocation_parameters(
+            input.invocation_parameters
         )
-        status_code: StatusCode
-        status_message = ""
-        events: list[SpanEvent] = []
-        start_time: datetime
-        end_time: datetime
-        response_chunks = []
-        text_chunks: list[TextChunk] = []
-        tool_call_chunks: defaultdict[ToolCallID, list[ToolCallChunk]] = defaultdict(list)
-        try:
-            start_time = cast(datetime, normalize_datetime(dt=local_now(), tz=timezone.utc))
+        async with streaming_llm_span(
+            input=input,
+            messages=messages,
+            invocation_parameters=invocation_parameters,
+        ) as span:
             async for chunk in llm_client.chat_completion_create(
-                messages=messages,
-                tools=input.tools or [],
-                **invocation_parameters,
+                messages=messages, tools=input.tools or [], **invocation_parameters
             ):
-                response_chunks.append(chunk)
-                if isinstance(chunk, TextChunk):
-                    yield chunk
-                    text_chunks.append(chunk)
-                elif isinstance(chunk, ToolCallChunk):
-                    yield chunk
-                    tool_call_chunks[chunk.id].append(chunk)
-                else:
-                    assert_never(chunk)
-            status_code = StatusCode.OK
-        except Exception as error:
-            end_time = cast(datetime, normalize_datetime(dt=local_now(), tz=timezone.utc))
-            status_code = StatusCode.ERROR
-            status_message = str(error)
-            events.append(
-                SpanException(
-                    timestamp=end_time,
-                    message=status_message,
-                    exception_type=type(error).__name__,
-                    exception_escaped=False,
-                    exception_stacktrace=format_exc(),
+                span.add_response_chunk(chunk)
+                yield chunk
+            span.set_attributes(llm_client.attributes)
+        if span.status_message is not None:
+            yield ChatCompletionSubscriptionError(message=span.status_message)
+        async with info.context.db() as session:
+            if (
+                playground_project_id := await session.scalar(
+                    select(models.Project.id).where(models.Project.name == PLAYGROUND_PROJECT_NAME)
                 )
-            )
-            yield ChatCompletionSubscriptionError(message=status_message)
-        else:
-            end_time = cast(datetime, normalize_datetime(dt=local_now(), tz=timezone.utc))
-            attributes.update(
-                chain(
-                    _output_value_and_mime_type(response_chunks),
-                    _llm_output_messages(text_chunks, tool_call_chunks),
+            ) is None:
+                playground_project_id = await session.scalar(
+                    insert(models.Project)
+                    .returning(models.Project.id)
+                    .values(
+                        name=PLAYGROUND_PROJECT_NAME,
+                        description="Traces from prompt playground",
+                    )
                 )
+            db_trace = get_db_trace(span, playground_project_id)
+            db_span = get_db_span(span, db_trace)
+            session.add(db_span)
+            await session.flush()
+        info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
+        yield ChatCompletionSubscriptionResult(span=to_gql_span(db_span))
+
+    @strawberry.subscription
+    async def chat_completion_over_dataset(
+        self, info: Info[Context, None], input: ChatCompletionOverDatasetInput
+    ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
+        provider_key = input.model.provider_key
+        llm_client_class = PLAYGROUND_CLIENT_REGISTRY.get_client(provider_key, input.model.name)
+        if llm_client_class is None:
+            raise BadRequest(f"No LLM client registered for provider '{provider_key}'")
+
+        dataset_id = from_global_id_with_expected_type(input.dataset_id, Dataset.__name__)
+        version_id = (
+            from_global_id_with_expected_type(
+                global_id=input.dataset_version_id, expected_type_name=DatasetVersion.__name__
             )
-        prompt_tokens = attributes.get(LLM_TOKEN_COUNT_PROMPT, 0)
-        completion_tokens = attributes.get(LLM_TOKEN_COUNT_COMPLETION, 0)
-        trace_id = _generate_trace_id()
-        span_id = _generate_span_id()
+            if input.dataset_version_id
+            else None
+        )
         async with info.context.db() as session:
+            if (
+                dataset := await session.scalar(
+                    select(models.Dataset).where(models.Dataset.id == dataset_id)
+                )
+            ) is None:
+                raise NotFound(f"Could not find dataset with ID {dataset_id}")
+            if version_id is None:
+                if (
+                    resolved_version_id := await session.scalar(
+                        select(models.DatasetVersion.id)
+                        .where(models.DatasetVersion.dataset_id == dataset_id)
+                        .order_by(models.DatasetVersion.id.desc())
+                        .limit(1)
+                    )
+                ) is None:
+                    raise NotFound(f"No versions found for dataset with ID {dataset_id}")
+            else:
+                if (
+                    resolved_version_id := await session.scalar(
+                        select(models.DatasetVersion.id).where(
+                            and_(
+                                models.DatasetVersion.dataset_id == dataset_id,
+                                models.DatasetVersion.id == version_id,
+                            )
+                        )
+                    )
+                ) is None:
+                    raise NotFound(f"Could not find dataset version with ID {version_id}")
+            revision_ids = (
+                select(func.max(models.DatasetExampleRevision.id))
+                .join(models.DatasetExample)
+                .where(
+                    and_(
+                        models.DatasetExample.dataset_id == dataset_id,
+                        models.DatasetExampleRevision.dataset_version_id <= resolved_version_id,
+                    )
+                )
+                .group_by(models.DatasetExampleRevision.dataset_example_id)
+            )
+            if not (
+                revisions := [
+                    rev
+                    async for rev in await session.stream_scalars(
+                        select(models.DatasetExampleRevision)
+                        .where(
+                            and_(
+                                models.DatasetExampleRevision.id.in_(revision_ids),
+                                models.DatasetExampleRevision.revision_kind != "DELETE",
+                            )
+                        )
+                        .order_by(models.DatasetExampleRevision.dataset_example_id.asc())
+                        .options(
+                            load_only(
+                                models.DatasetExampleRevision.dataset_example_id,
+                                models.DatasetExampleRevision.input,
+                            )
+                        )
+                    )
+                ]
+            ):
+                raise NotFound("No examples found for the given dataset and version")
             if (
                 playground_project_id := await session.scalar(
                     select(models.Project.id).where(models.Project.name == PLAYGROUND_PROJECT_NAME)
@@ -518,138 +243,208 @@ class Subscription:
                         description="Traces from prompt playground",
                     )
                 )
-            playground_trace = models.Trace(
-                project_rowid=playground_project_id,
-                trace_id=trace_id,
-                start_time=start_time,
-                end_time=end_time,
+            experiment = models.Experiment(
+                dataset_id=from_global_id_with_expected_type(input.dataset_id, Dataset.__name__),
+                dataset_version_id=resolved_version_id,
+                name=input.experiment_name or _default_playground_experiment_name(),
+                description=input.experiment_description
+                or _default_playground_experiment_description(dataset_name=dataset.name),
+                repetitions=1,
+                metadata_=input.experiment_metadata
+                or _default_playground_experiment_metadata(
+                    dataset_name=dataset.name,
+                    dataset_id=input.dataset_id,
+                    version_id=GlobalID(DatasetVersion.__name__, str(resolved_version_id)),
+                ),
+                project_name=PLAYGROUND_PROJECT_NAME,
             )
-            playground_span = models.Span(
-                trace_rowid=playground_trace.id,
-                span_id=span_id,
-                parent_id=None,
-                name="ChatCompletion",
-                span_kind=LLM,
-                start_time=start_time,
-                end_time=end_time,
-                attributes=unflatten(attributes.items()),
-                events=[_serialize_event(event) for event in events],
-                status_code=status_code.name,
-                status_message=status_message,
-                cumulative_error_count=int(status_code is StatusCode.ERROR),
-                cumulative_llm_token_count_prompt=prompt_tokens,
-                cumulative_llm_token_count_completion=completion_tokens,
-                llm_token_count_prompt=prompt_tokens,
-                llm_token_count_completion=completion_tokens,
-                trace=playground_trace,
-            )
-            session.add(playground_trace)
-            session.add(playground_span)
+            session.add(experiment)
             await session.flush()
-            yield FinishedChatCompletion(span=to_gql_span(playground_span))
-        info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
-
-
-def _llm_span_kind() -> Iterator[tuple[str, Any]]:
-    yield OPENINFERENCE_SPAN_KIND, LLM
-
-
-def _llm_model_name(model_name: str) -> Iterator[tuple[str, Any]]:
-    yield LLM_MODEL_NAME, model_name
-
-
-def _llm_invocation_parameters(invocation_parameters: dict[str, Any]) -> Iterator[tuple[str, Any]]:
-    yield LLM_INVOCATION_PARAMETERS, safe_json_dumps(invocation_parameters)
-
-
-def _llm_tools(tools: list[JSONScalarType]) -> Iterator[tuple[str, Any]]:
-    for tool_index, tool in enumerate(tools):
-        yield f"{LLM_TOOLS}.{tool_index}.{TOOL_JSON_SCHEMA}", json.dumps(tool)
-
-
-def _input_value_and_mime_type(input: ChatCompletionInput) -> Iterator[tuple[str, Any]]:
-    assert (api_key := "api_key") in (input_data := jsonify(input))
-    input_data = {k: v for k, v in input_data.items() if k != api_key}
-    assert api_key not in input_data
-    yield INPUT_MIME_TYPE, JSON
-    yield INPUT_VALUE, safe_json_dumps(input_data)
-
-
-def _output_value_and_mime_type(output: Any) -> Iterator[tuple[str, Any]]:
-    yield OUTPUT_MIME_TYPE, JSON
-    yield OUTPUT_VALUE, safe_json_dumps(jsonify(output))
+        yield ChatCompletionSubscriptionExperiment(
+            experiment=to_gql_experiment(experiment)
+        )  # eagerly yields experiment so it can be linked by consumers of the subscription
+
+        results_queue: Queue[ChatCompletionResult] = Queue()
+        chat_completion_streams = [
+            _stream_chat_completion_over_dataset_example(
+                input=input,
+                llm_client_class=llm_client_class,
+                revision=revision,
+                results_queue=results_queue,
+                experiment_id=experiment.id,
+                project_id=playground_project_id,
+            )
+            for revision in revisions
+        ]
+        stream_to_async_tasks: dict[
+            AsyncIterator[ChatCompletionSubscriptionPayload],
+            Task[ChatCompletionSubscriptionPayload],
+        ] = {iterator: _create_task_with_timeout(iterator) for iterator in chat_completion_streams}
+        batch_size = 10
+        while stream_to_async_tasks:
+            async_tasks_to_run = [task for task in stream_to_async_tasks.values()]
+            completed_tasks, _ = await wait(async_tasks_to_run, return_when=FIRST_COMPLETED)
+            for task in completed_tasks:
+                iterator = next(it for it, t in stream_to_async_tasks.items() if t == task)
+                try:
+                    yield task.result()
+                except (StopAsyncIteration, asyncio.TimeoutError):
+                    del stream_to_async_tasks[iterator]  # removes exhausted iterator
+                except Exception as error:
+                    del stream_to_async_tasks[iterator]  # removes failed iterator
+                    logger.exception(error)
+                else:
+                    stream_to_async_tasks[iterator] = _create_task_with_timeout(iterator)
+                if results_queue.qsize() >= batch_size:
+                    result_iterator = _chat_completion_result_payloads(
+                        db=info.context.db, results=_drain_no_wait(results_queue)
+                    )
+                    stream_to_async_tasks[result_iterator] = _create_task_with_timeout(
+                        result_iterator
+                    )
+        if remaining_results := await _drain(results_queue):
+            async for result_payload in _chat_completion_result_payloads(
+                db=info.context.db, results=remaining_results
+            ):
+                yield result_payload
+
+
+async def _stream_chat_completion_over_dataset_example(
+    *,
+    input: ChatCompletionOverDatasetInput,
+    llm_client_class: type["PlaygroundStreamingClient"],
+    revision: models.DatasetExampleRevision,
+    results_queue: Queue[ChatCompletionResult],
+    experiment_id: int,
+    project_id: int,
+) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
+    example_id = GlobalID(DatasetExample.__name__, str(revision.dataset_example_id))
+    llm_client = llm_client_class(
+        model=input.model,
+        api_key=input.api_key,
+    )
+    invocation_parameters = llm_client.construct_invocation_parameters(input.invocation_parameters)
+    messages = [
+        (
+            message.role,
+            message.content,
+            message.tool_call_id if isinstance(message.tool_call_id, str) else None,
+            message.tool_calls if isinstance(message.tool_calls, list) else None,
+        )
+        for message in input.messages
+    ]
+    try:
+        format_start_time = cast(datetime, normalize_datetime(dt=local_now(), tz=timezone.utc))
+        messages = list(
+            _formatted_messages(
+                messages=messages,
+                template_language=input.template_language,
+                template_variables=revision.input,
+            )
+        )
+    except TemplateFormatterError as error:
+        format_end_time = cast(datetime, normalize_datetime(dt=local_now(), tz=timezone.utc))
+        yield ChatCompletionSubscriptionError(message=str(error), dataset_example_id=example_id)
+        await results_queue.put(
+            (
+                example_id,
+                None,
+                models.ExperimentRun(
+                    experiment_id=experiment_id,
+                    dataset_example_id=revision.dataset_example_id,
+                    trace_id=None,
+                    output={},
+                    repetition_number=1,
+                    start_time=format_start_time,
+                    end_time=format_end_time,
+                    error=str(error),
+                    trace=None,
+                ),
+            )
+        )
+        return
+    async with streaming_llm_span(
+        input=input,
+        messages=messages,
+        invocation_parameters=invocation_parameters,
+    ) as span:
+        async for chunk in llm_client.chat_completion_create(
+            messages=messages, tools=input.tools or [], **invocation_parameters
+        ):
+            span.add_response_chunk(chunk)
+            chunk.dataset_example_id = example_id
+            yield chunk
+        span.set_attributes(llm_client.attributes)
+    db_trace = get_db_trace(span, project_id)
+    db_span = get_db_span(span, db_trace)
+    db_run = get_db_experiment_run(
+        db_span, db_trace, experiment_id=experiment_id, example_id=revision.dataset_example_id
+    )
+    await results_queue.put((example_id, db_span, db_run))
+    if span.status_message is not None:
+        yield ChatCompletionSubscriptionError(
+            message=span.status_message, dataset_example_id=example_id
+        )
 
 
-def _llm_input_messages(
-    messages: Iterable[
-        tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[JSONScalarType]]]
-    ],
-) -> Iterator[tuple[str, Any]]:
-    for i, (role, content, _tool_call_id, tool_calls) in enumerate(messages):
-        yield f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_ROLE}", role.value.lower()
-        yield f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_CONTENT}", content
-        if tool_calls is not None:
-            for tool_call_index, tool_call in enumerate(tool_calls):
-                yield (
-                    f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
-                    tool_call["function"]["name"],
-                )
-                if arguments := tool_call["function"]["arguments"]:
-                    yield (
-                        f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
-                        safe_json_dumps(jsonify(arguments)),
-                    )
+async def _chat_completion_result_payloads(
+    *,
+    db: DbSessionFactory,
+    results: Sequence[ChatCompletionResult],
+) -> AsyncIterator[ChatCompletionSubscriptionResult]:
+    if not results:
+        return
+    async with db() as session:
+        for _, span, run in results:
+            if span:
+                session.add(span)
+            session.add(run)
+        await session.flush()
+    for example_id, span, run in results:
+        yield ChatCompletionSubscriptionResult(
+            span=to_gql_span(span) if span else None,
+            experiment_run=to_gql_experiment_run(run),
+            dataset_example_id=example_id,
+        )
 
 
-def _llm_output_messages(
-    text_chunks: list[TextChunk],
-    tool_call_chunks: defaultdict[ToolCallID, list[ToolCallChunk]],
-) -> Iterator[tuple[str, Any]]:
-    yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"
-    if content := "".join(chunk.content for chunk in text_chunks):
-        yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", content
-    for tool_call_index, (_tool_call_id, tool_call_chunks_) in enumerate(tool_call_chunks.items()):
-        if tool_call_chunks_ and (name := tool_call_chunks_[0].function.name):
-            yield (
-                f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
-                name,
-            )
-        if arguments := "".join(chunk.function.arguments for chunk in tool_call_chunks_):
-            yield (
-                f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
-                arguments,
-            )
+def _create_task_with_timeout(
+    iterable: AsyncIterator[GenericType], timeout_in_seconds: int = 60
+) -> Task[GenericType]:
+    return create_task(wait_for(_as_coroutine(iterable), timeout=timeout_in_seconds))
 
 
-def _generate_trace_id() -> str:
-    """
-    Generates a random trace ID in hexadecimal format.
-    """
-    return _hex(DefaultOTelIDGenerator().generate_trace_id())
+async def _drain(queue: Queue[GenericType]) -> list[GenericType]:
+    values: list[GenericType] = []
+    while not queue.empty():
+        values.append(await queue.get())
+    return values
 
 
-def _generate_span_id() -> str:
-    """
-    Generates a random span ID in hexadecimal format.
-    """
-    return _hex(DefaultOTelIDGenerator().generate_span_id())
+def _drain_no_wait(queue: Queue[GenericType]) -> list[GenericType]:
+    values: list[GenericType] = []
+    while True:
+        try:
+            values.append(queue.get_nowait())
+        except QueueEmpty:
+            break
+    return values
 
 
-def _hex(number: int) -> str:
-    """
-    Converts an integer to a hexadecimal string.
-    """
-    return hex(number)[2:]
+async def _as_coroutine(iterable: AsyncIterator[GenericType]) -> GenericType:
+    return await iterable.__anext__()
 
 
 def _formatted_messages(
-    messages: Iterable[tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[str]]]],
-    template_options: TemplateOptions,
+    *,
+    messages: Iterable[ChatCompletionMessage],
+    template_language: TemplateLanguage,
+    template_variables: Mapping[str, Any],
 ) -> Iterator[tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[str]]]]:
     """
     Formats the messages using the given template options.
     """
-    template_formatter = _template_formatter(template_language=template_options.language)
+    template_formatter = _template_formatter(template_language=template_language)
     (
         roles,
         templates,
@@ -657,7 +452,7 @@ def _formatted_messages(
         tool_calls,
     ) = zip(*messages)
     formatted_templates = map(
-        lambda template: template_formatter.format(template, **template_options.variables),
+        lambda template: template_formatter.format(template, **template_variables),
         templates,
     )
     formatted_messages = zip(roles, formatted_templates, tool_call_id, tool_calls)
@@ -675,36 +470,24 @@ def _template_formatter(template_language: TemplateLanguage) -> TemplateFormatte
     assert_never(template_language)
 
 
-def _serialize_event(event: SpanEvent) -> dict[str, Any]:
-    """
-    Serializes a SpanEvent to a dictionary.
-    """
-    return {k: (v.isoformat() if isinstance(v, datetime) else v) for k, v in asdict(event).items()}
+def _default_playground_experiment_name() -> str:
+    return "playground-experiment"
 
 
-JSON = OpenInferenceMimeTypeValues.JSON.value
+def _default_playground_experiment_description(dataset_name: str) -> str:
+    return f'Playground experiment for dataset "{dataset_name}"'
 
-LLM = OpenInferenceSpanKindValues.LLM.value
 
-OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
-INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
-INPUT_VALUE = SpanAttributes.INPUT_VALUE
-OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
-OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
-LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
-LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
-LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
-LLM_INVOCATION_PARAMETERS = SpanAttributes.LLM_INVOCATION_PARAMETERS
-LLM_TOOLS = SpanAttributes.LLM_TOOLS
-LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
-LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
-LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
+def _default_playground_experiment_metadata(
+    dataset_name: str, dataset_id: GlobalID, version_id: GlobalID
+) -> dict[str, Any]:
+    return {
+        "dataset_name": dataset_name,
+        "dataset_id": str(dataset_id),
+        "dataset_version_id": str(version_id),
+    }
 
-MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
-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
-
-TOOL_JSON_SCHEMA = ToolAttributes.TOOL_JSON_SCHEMA
+LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
+LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
+LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT