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

phoenix/server/api/subscriptions.py

@@ -1,10 +1,12 @@
 import asyncio
 import logging
+from collections import deque
 from collections.abc import AsyncIterator, Iterator
 from datetime import datetime, timedelta, timezone
 from typing import (
     Any,
     AsyncGenerator,
+    Callable,
     Coroutine,
     Iterable,
     Mapping,
@@ -17,7 +19,7 @@ from typing import (
 import strawberry
 from openinference.instrumentation import safe_json_dumps
 from openinference.semconv.trace import SpanAttributes
-from sqlalchemy import and_, func, insert, select
+from sqlalchemy import and_, insert, select
 from sqlalchemy.orm import load_only
 from strawberry.relay.types import GlobalID
 from strawberry.types import Info
@@ -26,10 +28,15 @@ from typing_extensions import TypeAlias, assert_never
 from phoenix.config import PLAYGROUND_PROJECT_NAME
 from phoenix.datetime_utils import local_now, normalize_datetime
 from phoenix.db import models
-from phoenix.server.api.auth import IsLocked, IsNotReadOnly
+from phoenix.db.helpers import (
+    get_dataset_example_revisions,
+    insert_experiment_with_examples_snapshot,
+)
+from phoenix.server.api.auth import IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, CustomGraphQLError, NotFound
 from phoenix.server.api.helpers.playground_clients import (
+    PlaygroundClientCredential,
     PlaygroundStreamingClient,
     initialize_playground_clients,
 )
@@ -42,6 +49,7 @@ from phoenix.server.api.helpers.playground_spans import (
     get_db_trace,
     streaming_llm_span,
 )
+from phoenix.server.api.helpers.playground_users import get_user
 from phoenix.server.api.helpers.prompts.models import PromptTemplateFormat
 from phoenix.server.api.input_types.ChatCompletionInput import (
     ChatCompletionInput,
@@ -58,10 +66,12 @@ 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.ExperimentRun import ExperimentRun
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Span import Span
+from phoenix.server.daemons.span_cost_calculator import SpanCostCalculator
 from phoenix.server.dml_event import SpanInsertEvent
+from phoenix.server.experiments.utils import generate_experiment_project_name
 from phoenix.server.types import DbSessionFactory
 from phoenix.utilities.template_formatters import (
     FStringTemplateFormatter,
@@ -87,9 +97,109 @@ ChatCompletionResult: TypeAlias = tuple[
 ChatStream: TypeAlias = AsyncGenerator[ChatCompletionSubscriptionPayload, None]
 
 
+async def _stream_single_chat_completion(
+    *,
+    input: ChatCompletionInput,
+    llm_client: PlaygroundStreamingClient,
+    project_id: int,
+    repetition_number: int,
+    results: asyncio.Queue[tuple[Optional[models.Span], int]],
+) -> ChatStream:
+    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
+    ]
+    attributes = None
+    if template_options := input.template:
+        messages = list(
+            _formatted_messages(
+                messages=messages,
+                template_format=template_options.format,
+                template_variables=template_options.variables,
+            )
+        )
+        attributes = {PROMPT_TEMPLATE_VARIABLES: safe_json_dumps(template_options.variables)}
+    invocation_parameters = llm_client.construct_invocation_parameters(input.invocation_parameters)
+    async with streaming_llm_span(
+        input=input,
+        messages=messages,
+        invocation_parameters=invocation_parameters,
+        attributes=attributes,
+    ) as span:
+        try:
+            async for chunk in llm_client.chat_completion_create(
+                messages=messages, tools=input.tools or [], **invocation_parameters
+            ):
+                span.add_response_chunk(chunk)
+                chunk.repetition_number = repetition_number
+                yield chunk
+        finally:
+            span.set_attributes(llm_client.attributes)
+    if span.status_message is not None:
+        yield ChatCompletionSubscriptionError(
+            message=span.status_message,
+            repetition_number=repetition_number,
+        )
+
+    db_trace = get_db_trace(span, project_id)
+    db_span = get_db_span(span, db_trace)
+    await results.put((db_span, repetition_number))
+
+
+async def _chat_completion_span_result_payloads(
+    *,
+    db: DbSessionFactory,
+    results: Sequence[tuple[Optional[models.Span], int]],
+    span_cost_calculator: SpanCostCalculator,
+    on_span_insertion: Callable[[], None],
+) -> ChatStream:
+    if not results:
+        return
+    async with db() as session:
+        for span, repetition_number in results:
+            if span:
+                session.add(span)
+                await session.flush()
+                try:
+                    span_cost = span_cost_calculator.calculate_cost(
+                        start_time=span.start_time,
+                        attributes=span.attributes,
+                    )
+                except Exception as e:
+                    logger.exception(f"Failed to calculate cost for span {span.id}: {e}")
+                    span_cost = None
+                if span_cost:
+                    span_cost.span_rowid = span.id
+                    span_cost.trace_rowid = span.trace_rowid
+                    session.add(span_cost)
+        await session.flush()
+    for span, repetition_number in results:
+        if span:
+            yield ChatCompletionSubscriptionResult(
+                span=Span(id=span.id, db_record=span),
+                repetition_number=repetition_number,
+            )
+            on_span_insertion()
+
+
+def _is_span_result_payloads_stream(
+    stream: ChatStream,
+) -> bool:
+    """
+    Checks if the given generator was instantiated from
+    `_chat_completion_span_result_payloads`
+    """
+    return stream.ag_code == _chat_completion_span_result_payloads.__code__  # type: ignore
+
+
 @strawberry.type
 class Subscription:
-    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def chat_completion(
         self, info: Info[Context, None], input: ChatCompletionInput
     ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
@@ -98,9 +208,17 @@ class Subscription:
         if llm_client_class is None:
             raise BadRequest(f"Unknown LLM provider: '{provider_key.value}'")
         try:
+            # Convert GraphQL credentials to PlaygroundCredential objects
+            playground_credentials = None
+            if input.credentials:
+                playground_credentials = [
+                    PlaygroundClientCredential(env_var_name=cred.env_var_name, value=cred.value)
+                    for cred in input.credentials
+                ]
+
             llm_client = llm_client_class(
                 model=input.model,
-                api_key=input.api_key,
+                credentials=playground_credentials,
             )
         except CustomGraphQLError:
             raise
@@ -110,42 +228,6 @@ class Subscription:
                 f"{str(error)}"
             )
 
-        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
-        ]
-        attributes = None
-        if template_options := input.template:
-            messages = list(
-                _formatted_messages(
-                    messages=messages,
-                    template_format=template_options.format,
-                    template_variables=template_options.variables,
-                )
-            )
-            attributes = {PROMPT_TEMPLATE_VARIABLES: safe_json_dumps(template_options.variables)}
-        invocation_parameters = llm_client.construct_invocation_parameters(
-            input.invocation_parameters
-        )
-        async with streaming_llm_span(
-            input=input,
-            messages=messages,
-            invocation_parameters=invocation_parameters,
-            attributes=attributes,
-        ) as span:
-            async for chunk in llm_client.chat_completion_create(
-                messages=messages, tools=input.tools or [], **invocation_parameters
-            ):
-                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(
@@ -160,14 +242,100 @@ class Subscription:
                         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=Span(span_rowid=db_span.id, db_span=db_span))
-
-    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+
+        results: asyncio.Queue[tuple[Optional[models.Span], int]] = asyncio.Queue()
+        not_started: deque[tuple[int, ChatStream]] = deque(
+            (
+                repetition_number,
+                _stream_single_chat_completion(
+                    input=input,
+                    llm_client=llm_client,
+                    project_id=playground_project_id,
+                    repetition_number=repetition_number,
+                    results=results,
+                ),
+            )
+            for repetition_number in range(1, input.repetitions + 1)
+        )
+        in_progress: list[
+            tuple[
+                Optional[int],
+                ChatStream,
+                asyncio.Task[ChatCompletionSubscriptionPayload],
+            ]
+        ] = []
+        max_in_progress = 3
+        write_batch_size = 10
+        write_interval = timedelta(seconds=10)
+        last_write_time = datetime.now()
+        while not_started or in_progress:
+            while not_started and len(in_progress) < max_in_progress:
+                rep_num, stream = not_started.popleft()
+                task = _create_task_with_timeout(stream)
+                in_progress.append((rep_num, stream, task))
+            async_tasks_to_run = [task for _, _, task in in_progress]
+            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)
+                repetition_number, stream, _ = in_progress[idx]
+                try:
+                    yield completed_task.result()
+                except StopAsyncIteration:
+                    del in_progress[idx]  # removes exhausted stream
+                except asyncio.TimeoutError:
+                    del in_progress[idx]  # removes timed-out stream
+                    if repetition_number is not None:
+                        yield ChatCompletionSubscriptionError(
+                            message="Playground task timed out",
+                            repetition_number=repetition_number,
+                        )
+                except Exception as error:
+                    del in_progress[idx]  # removes failed stream
+                    if repetition_number is not None:
+                        yield ChatCompletionSubscriptionError(
+                            message="An unexpected error occurred",
+                            repetition_number=repetition_number,
+                        )
+                    logger.exception(error)
+                else:
+                    task = _create_task_with_timeout(stream)
+                    in_progress[idx] = (repetition_number, stream, task)
+
+                exceeded_write_batch_size = results.qsize() >= write_batch_size
+                exceeded_write_interval = datetime.now() - last_write_time > write_interval
+                write_already_in_progress = any(
+                    _is_span_result_payloads_stream(stream) for _, stream, _ in in_progress
+                )
+                if (
+                    not results.empty()
+                    and (exceeded_write_batch_size or exceeded_write_interval)
+                    and not write_already_in_progress
+                ):
+                    result_payloads_stream = _chat_completion_span_result_payloads(
+                        db=info.context.db,
+                        results=_drain_no_wait(results),
+                        span_cost_calculator=info.context.span_cost_calculator,
+                        on_span_insertion=lambda: info.context.event_queue.put(
+                            SpanInsertEvent(ids=(playground_project_id,))
+                        ),
+                    )
+                    task = _create_task_with_timeout(result_payloads_stream)
+                    in_progress.append((None, result_payloads_stream, task))
+                    last_write_time = datetime.now()
+        if remaining_results := await _drain(results):
+            async for result_payload in _chat_completion_span_result_payloads(
+                db=info.context.db,
+                results=remaining_results,
+                span_cost_calculator=info.context.span_cost_calculator,
+                on_span_insertion=lambda: info.context.event_queue.put(
+                    SpanInsertEvent(ids=(playground_project_id,))
+                ),
+            ):
+                yield result_payload
+
+    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def chat_completion_over_dataset(
         self, info: Info[Context, None], input: ChatCompletionOverDatasetInput
     ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
@@ -176,9 +344,17 @@ class Subscription:
         if llm_client_class is None:
             raise BadRequest(f"Unknown LLM provider: '{provider_key.value}'")
         try:
+            # Convert GraphQL credentials to PlaygroundCredential objects
+            playground_credentials = None
+            if input.credentials:
+                playground_credentials = [
+                    PlaygroundClientCredential(env_var_name=cred.env_var_name, value=cred.value)
+                    for cred in input.credentials
+                ]
+
             llm_client = llm_client_class(
                 model=input.model,
-                api_key=input.api_key,
+                credentials=playground_credentials,
             )
         except CustomGraphQLError:
             raise
@@ -223,27 +399,22 @@ class Subscription:
                     )
                 ) 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)
-            )
+
+            # Parse split IDs if provided
+            resolved_split_ids: Optional[list[int]] = None
+            if input.split_ids is not None and len(input.split_ids) > 0:
+                resolved_split_ids = [
+                    from_global_id_with_expected_type(split_id, models.DatasetSplit.__name__)
+                    for split_id in input.split_ids
+                ]
+
             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",
-                            )
+                        get_dataset_example_revisions(
+                            resolved_version_id,
+                            split_ids=resolved_split_ids,
                         )
                         .order_by(models.DatasetExampleRevision.dataset_example_id.asc())
                         .options(
@@ -256,31 +427,38 @@ class Subscription:
                 ]
             ):
                 raise NotFound("No examples found for the given dataset and version")
+            project_name = generate_experiment_project_name()
             if (
                 playground_project_id := await session.scalar(
-                    select(models.Project.id).where(models.Project.name == PLAYGROUND_PROJECT_NAME)
+                    select(models.Project.id).where(models.Project.name == project_name)
                 )
             ) is None:
                 playground_project_id = await session.scalar(
                     insert(models.Project)
                     .returning(models.Project.id)
                     .values(
-                        name=PLAYGROUND_PROJECT_NAME,
+                        name=project_name,
                         description="Traces from prompt playground",
                     )
                 )
+            user_id = get_user(info)
             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(input.prompt_name),
                 description=input.experiment_description,
-                repetitions=1,
+                repetitions=input.repetitions,
                 metadata_=input.experiment_metadata or dict(),
-                project_name=PLAYGROUND_PROJECT_NAME,
+                project_name=project_name,
+                user_id=user_id,
             )
-            session.add(experiment)
-            await session.flush()
+            if resolved_split_ids:
+                experiment.experiment_dataset_splits = [
+                    models.ExperimentDatasetSplit(dataset_split_id=split_id)
+                    for split_id in resolved_split_ids
+                ]
+            await insert_experiment_with_examples_snapshot(session, experiment)
         yield ChatCompletionSubscriptionExperiment(
             experiment=to_gql_experiment(experiment)
         )  # eagerly yields experiment so it can be linked by consumers of the subscription
@@ -294,11 +472,15 @@ class Subscription:
                     llm_client=llm_client,
                     revision=revision,
                     results=results,
+                    repetition_number=repetition_number,
                     experiment_id=experiment.id,
                     project_id=playground_project_id,
                 ),
             )
             for revision in revisions
+            for repetition_number in reversed(
+                range(1, input.repetitions + 1)
+            )  # since we pop right, this runs the repetitions in increasing order
         ]
         in_progress: list[
             tuple[
@@ -355,14 +537,18 @@ class Subscription:
                     and not write_already_in_progress
                 ):
                     result_payloads_stream = _chat_completion_result_payloads(
-                        db=info.context.db, results=_drain_no_wait(results)
+                        db=info.context.db,
+                        results=_drain_no_wait(results),
+                        span_cost_calculator=info.context.span_cost_calculator,
                     )
                     task = _create_task_with_timeout(result_payloads_stream)
                     in_progress.append((None, result_payloads_stream, task))
                     last_write_time = datetime.now()
         if remaining_results := await _drain(results):
             async for result_payload in _chat_completion_result_payloads(
-                db=info.context.db, results=remaining_results
+                db=info.context.db,
+                results=remaining_results,
+                span_cost_calculator=info.context.span_cost_calculator,
             ):
                 yield result_payload
 
@@ -372,6 +558,7 @@ async def _stream_chat_completion_over_dataset_example(
     input: ChatCompletionOverDatasetInput,
     llm_client: PlaygroundStreamingClient,
     revision: models.DatasetExampleRevision,
+    repetition_number: int,
     results: asyncio.Queue[ChatCompletionResult],
     experiment_id: int,
     project_id: int,
@@ -398,7 +585,11 @@ async def _stream_chat_completion_over_dataset_example(
         )
     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)
+        yield ChatCompletionSubscriptionError(
+            message=str(error),
+            dataset_example_id=example_id,
+            repetition_number=repetition_number,
+        )
         await results.put(
             (
                 example_id,
@@ -408,7 +599,7 @@ async def _stream_chat_completion_over_dataset_example(
                     dataset_example_id=revision.dataset_example_id,
                     trace_id=None,
                     output={},
-                    repetition_number=1,
+                    repetition_number=repetition_number,
                     start_time=format_start_time,
                     end_time=format_end_time,
                     error=str(error),
@@ -423,22 +614,31 @@ async def _stream_chat_completion_over_dataset_example(
         invocation_parameters=invocation_parameters,
         attributes={PROMPT_TEMPLATE_VARIABLES: safe_json_dumps(revision.input)},
     ) 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)
+        try:
+            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
+                chunk.repetition_number = repetition_number
+                yield chunk
+        finally:
+            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
+        db_span,
+        db_trace,
+        experiment_id=experiment_id,
+        example_id=revision.dataset_example_id,
+        repetition_number=repetition_number,
     )
     await results.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
+            message=span.status_message,
+            dataset_example_id=example_id,
+            repetition_number=repetition_number,
         )
 
 
@@ -446,6 +646,7 @@ async def _chat_completion_result_payloads(
     *,
     db: DbSessionFactory,
     results: Sequence[ChatCompletionResult],
+    span_cost_calculator: SpanCostCalculator,
 ) -> ChatStream:
     if not results:
         return
@@ -453,13 +654,27 @@ async def _chat_completion_result_payloads(
         for _, span, run in results:
             if span:
                 session.add(span)
+                await session.flush()
+                try:
+                    span_cost = span_cost_calculator.calculate_cost(
+                        start_time=span.start_time,
+                        attributes=span.attributes,
+                    )
+                except Exception as e:
+                    logger.exception(f"Failed to calculate cost for span {span.id}: {e}")
+                    span_cost = None
+                if span_cost:
+                    span_cost.span_rowid = span.id
+                    span_cost.trace_rowid = span.trace_rowid
+                    session.add(span_cost)
             session.add(run)
         await session.flush()
     for example_id, span, run in results:
         yield ChatCompletionSubscriptionResult(
-            span=Span(span_rowid=span.id, db_span=span) if span else None,
-            experiment_run=to_gql_experiment_run(run),
+            span=Span(id=span.id, db_record=span) if span else None,
+            experiment_run=ExperimentRun(id=run.id, db_record=run),
             dataset_example_id=example_id,
+            repetition_number=run.repetition_number,
         )
 
 
@@ -577,3 +792,5 @@ 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
 PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
+LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
+LLM_PROVIDER = SpanAttributes.LLM_PROVIDER