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

phoenix/server/api/subscriptions.py

@@ -2,9 +2,10 @@ 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 datetime import datetime, timedelta, timezone
 from typing import (
     Any,
+    AsyncGenerator,
     Iterable,
     Mapping,
     Optional,
@@ -24,7 +25,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.context import Context
-from phoenix.server.api.exceptions import BadRequest, NotFound
+from phoenix.server.api.exceptions import BadRequest, CustomGraphQLError, NotFound
 from phoenix.server.api.helpers.playground_clients import (
     PlaygroundStreamingClient,
     initialize_playground_clients,
@@ -79,6 +80,7 @@ DatasetExampleID: TypeAlias = GlobalID
 ChatCompletionResult: TypeAlias = tuple[
     DatasetExampleID, Optional[models.Span], models.ExperimentRun
 ]
+ChatStream: TypeAlias = AsyncGenerator[ChatCompletionSubscriptionPayload, None]
 PLAYGROUND_PROJECT_NAME = "playground"
 
 
@@ -91,11 +93,19 @@ class Subscription:
         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}'")
-        llm_client = llm_client_class(
-            model=input.model,
-            api_key=input.api_key,
-        )
+            raise BadRequest(f"Unknown LLM provider: '{provider_key.value}'")
+        try:
+            llm_client = llm_client_class(
+                model=input.model,
+                api_key=input.api_key,
+            )
+        except CustomGraphQLError:
+            raise
+        except Exception as error:
+            raise BadRequest(
+                f"Failed to connect to LLM API for {provider_key.value} {input.model.name}: "
+                f"{str(error)}"
+            )
 
         messages = [
             (
@@ -158,7 +168,19 @@ class Subscription:
         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}'")
+            raise BadRequest(f"Unknown LLM provider: '{provider_key.value}'")
+        try:
+            llm_client = llm_client_class(
+                model=input.model,
+                api_key=input.api_key,
+            )
+        except CustomGraphQLError:
+            raise
+        except Exception as error:
+            raise BadRequest(
+                f"Failed to connect to LLM API for {provider_key.value} {input.model.name}: "
+                f"{str(error)}"
+            )
 
         dataset_id = from_global_id_with_expected_type(input.dataset_id, Dataset.__name__)
         version_id = (
@@ -264,45 +286,76 @@ class Subscription:
             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,
+        results: Queue[ChatCompletionResult] = Queue()
+        not_started: list[tuple[DatasetExampleID, ChatStream]] = [
+            (
+                GlobalID(DatasetExample.__name__, str(revision.dataset_example_id)),
+                _stream_chat_completion_over_dataset_example(
+                    input=input,
+                    llm_client=llm_client,
+                    revision=revision,
+                    results=results,
+                    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()]
+        in_progress: list[
+            tuple[Optional[DatasetExampleID], ChatStream, 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:
+                ex_id, stream = not_started.pop()
+                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)
-            for task in completed_tasks:
-                iterator = next(it for it, t in stream_to_async_tasks.items() if t == task)
+            for completed_task in completed_tasks:
+                idx = [task for _, _, task in in_progress].index(completed_task)
+                example_id, stream, _ = in_progress[idx]
                 try:
-                    yield task.result()
-                except (StopAsyncIteration, asyncio.TimeoutError):
-                    del stream_to_async_tasks[iterator]  # removes exhausted iterator
+                    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 example_id is not None:
+                        yield ChatCompletionSubscriptionError(
+                            message="Timed out", dataset_example_id=example_id
+                        )
                 except Exception as error:
-                    del stream_to_async_tasks[iterator]  # removes failed iterator
+                    del in_progress[idx]  # removes failed stream
+                    if example_id is not None:
+                        yield ChatCompletionSubscriptionError(
+                            message="An unexpected error occurred", dataset_example_id=example_id
+                        )
                     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
+                    task = _create_task_with_timeout(stream)
+                    in_progress[idx] = (example_id, 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_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_result_payloads(
+                        db=info.context.db, results=_drain_no_wait(results)
                     )
-        if remaining_results := await _drain(results_queue):
+                    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
             ):
@@ -312,17 +365,13 @@ class Subscription:
 async def _stream_chat_completion_over_dataset_example(
     *,
     input: ChatCompletionOverDatasetInput,
-    llm_client_class: type["PlaygroundStreamingClient"],
+    llm_client: PlaygroundStreamingClient,
     revision: models.DatasetExampleRevision,
-    results_queue: Queue[ChatCompletionResult],
+    results: Queue[ChatCompletionResult],
     experiment_id: int,
     project_id: int,
-) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
+) -> ChatStream:
     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 = [
         (
@@ -345,7 +394,7 @@ 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)
-        await results_queue.put(
+        await results.put(
             (
                 example_id,
                 None,
@@ -380,7 +429,7 @@ async def _stream_chat_completion_over_dataset_example(
     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))
+    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
@@ -391,7 +440,7 @@ async def _chat_completion_result_payloads(
     *,
     db: DbSessionFactory,
     results: Sequence[ChatCompletionResult],
-) -> AsyncIterator[ChatCompletionSubscriptionResult]:
+) -> ChatStream:
     if not results:
         return
     async with db() as session:
@@ -408,8 +457,18 @@ async def _chat_completion_result_payloads(
         )
 
 
+def _is_result_payloads_stream(
+    stream: ChatStream,
+) -> bool:
+    """
+    Checks if the given generator was instantiated from
+    `_chat_completion_result_payloads`
+    """
+    return stream.ag_code == _chat_completion_result_payloads.__code__
+
+
 def _create_task_with_timeout(
-    iterable: AsyncIterator[GenericType], timeout_in_seconds: int = 60
+    iterable: AsyncIterator[GenericType], timeout_in_seconds: int = 90
 ) -> Task[GenericType]:
     return create_task(wait_for(_as_coroutine(iterable), timeout=timeout_in_seconds))
 

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

@@ -1,32 +1,32 @@
 {
-  "_components-BXIz9ZO8.js": {
-    "file": "assets/components-BXIz9ZO8.js",
+  "_components-C_HASv83.js": {
+    "file": "assets/components-C_HASv83.js",
     "name": "components",
     "imports": [
-      "_vendor-BX8_Znqy.js",
-      "_pages-B8FpJuXu.js",
-      "_vendor-arizeai-CtHir-Ua.js",
-      "_vendor-codemirror-DLlGiguX.js",
+      "_vendor-BCxsh5i3.js",
+      "_pages-DYHcAdjT.js",
+      "_vendor-arizeai-C2CDZgMz.js",
+      "_vendor-codemirror-DYbtnCTn.js",
       "_vendor-three-DwGkEfCM.js"
     ]
   },
-  "_pages-B8FpJuXu.js": {
-    "file": "assets/pages-B8FpJuXu.js",
+  "_pages-DYHcAdjT.js": {
+    "file": "assets/pages-DYHcAdjT.js",
     "name": "pages",
     "imports": [
-      "_vendor-BX8_Znqy.js",
-      "_vendor-arizeai-CtHir-Ua.js",
-      "_components-BXIz9ZO8.js",
-      "_vendor-recharts-CJRple0d.js",
-      "_vendor-codemirror-DLlGiguX.js"
+      "_vendor-BCxsh5i3.js",
+      "_vendor-arizeai-C2CDZgMz.js",
+      "_components-C_HASv83.js",
+      "_vendor-recharts-P6W8G0Mb.js",
+      "_vendor-codemirror-DYbtnCTn.js"
     ]
   },
   "_vendor-!~{003}~.js": {
     "file": "assets/vendor-DxkFTwjz.css",
     "src": "_vendor-!~{003}~.js"
   },
-  "_vendor-BX8_Znqy.js": {
-    "file": "assets/vendor-BX8_Znqy.js",
+  "_vendor-BCxsh5i3.js": {
+    "file": "assets/vendor-BCxsh5i3.js",
     "name": "vendor",
     "imports": [
       "_vendor-three-DwGkEfCM.js"
@@ -35,25 +35,25 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-CtHir-Ua.js": {
-    "file": "assets/vendor-arizeai-CtHir-Ua.js",
+  "_vendor-arizeai-C2CDZgMz.js": {
+    "file": "assets/vendor-arizeai-C2CDZgMz.js",
     "name": "vendor-arizeai",
     "imports": [
-      "_vendor-BX8_Znqy.js"
+      "_vendor-BCxsh5i3.js"
     ]
   },
-  "_vendor-codemirror-DLlGiguX.js": {
-    "file": "assets/vendor-codemirror-DLlGiguX.js",
+  "_vendor-codemirror-DYbtnCTn.js": {
+    "file": "assets/vendor-codemirror-DYbtnCTn.js",
     "name": "vendor-codemirror",
     "imports": [
-      "_vendor-BX8_Znqy.js"
+      "_vendor-BCxsh5i3.js"
     ]
   },
-  "_vendor-recharts-CJRple0d.js": {
-    "file": "assets/vendor-recharts-CJRple0d.js",
+  "_vendor-recharts-P6W8G0Mb.js": {
+    "file": "assets/vendor-recharts-P6W8G0Mb.js",
     "name": "vendor-recharts",
     "imports": [
-      "_vendor-BX8_Znqy.js"
+      "_vendor-BCxsh5i3.js"
     ]
   },
   "_vendor-three-DwGkEfCM.js": {
@@ -61,18 +61,18 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-DTut7g1y.js",
+    "file": "assets/index-D7UiCRtr.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
-      "_vendor-BX8_Znqy.js",
-      "_vendor-arizeai-CtHir-Ua.js",
-      "_pages-B8FpJuXu.js",
-      "_components-BXIz9ZO8.js",
+      "_vendor-BCxsh5i3.js",
+      "_vendor-arizeai-C2CDZgMz.js",
+      "_pages-DYHcAdjT.js",
+      "_components-C_HASv83.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-recharts-CJRple0d.js",
-      "_vendor-codemirror-DLlGiguX.js"
+      "_vendor-recharts-P6W8G0Mb.js",
+      "_vendor-codemirror-DYbtnCTn.js"
     ]
   }
 }