arize-phoenix 5.10.0__tar.gz → 5.11.0__tar.gz

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 5.10.0
+Version: 5.11.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -53,6 +53,7 @@ Requires-Dist: wrapt
 Provides-Extra: container
 Requires-Dist: anthropic; extra == 'container'
 Requires-Dist: fast-hdbscan>=0.2.0; extra == 'container'
+Requires-Dist: google-generativeai; extra == 'container'
 Requires-Dist: numba>=0.60.0; extra == 'container'
 Requires-Dist: openai>=1.0.0; extra == 'container'
 Requires-Dist: opentelemetry-exporter-otlp; extra == 'container'

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/pyproject.toml

@@ -121,6 +121,7 @@ pg = [
 ]
 container = [
   "anthropic",
+  "google-generativeai",
   "prometheus-client",
   "openai>=1.0.0",
   "opentelemetry-sdk",

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/src/phoenix/config.py

@@ -97,6 +97,13 @@ ENV_PHOENIX_SERVER_INSTRUMENTATION_OTLP_TRACE_COLLECTOR_GRPC_ENDPOINT = (
 ENV_PHOENIX_ENABLE_AUTH = "PHOENIX_ENABLE_AUTH"
 ENV_PHOENIX_DISABLE_RATE_LIMIT = "PHOENIX_DISABLE_RATE_LIMIT"
 ENV_PHOENIX_SECRET = "PHOENIX_SECRET"
+ENV_PHOENIX_DEFAULT_ADMIN_INITIAL_PASSWORD = "PHOENIX_DEFAULT_ADMIN_INITIAL_PASSWORD"
+"""
+The initial password for the default admin account, which defaults to ‘admin’ if not
+explicitly set. Note that changing this value will have no effect if the default admin
+record already exists in the database. In such cases, the default admin password must
+be updated manually in the application.
+"""
 ENV_PHOENIX_API_KEY = "PHOENIX_API_KEY"
 ENV_PHOENIX_USE_SECURE_COOKIES = "PHOENIX_USE_SECURE_COOKIES"
 ENV_PHOENIX_ACCESS_TOKEN_EXPIRY_MINUTES = "PHOENIX_ACCESS_TOKEN_EXPIRY_MINUTES"
@@ -274,6 +281,12 @@ def get_env_phoenix_secret() -> Optional[str]:
     return phoenix_secret
 
 
+def get_env_default_admin_initial_password() -> str:
+    from phoenix.auth import DEFAULT_ADMIN_PASSWORD
+
+    return os.environ.get(ENV_PHOENIX_DEFAULT_ADMIN_INITIAL_PASSWORD) or DEFAULT_ADMIN_PASSWORD
+
+
 def get_env_phoenix_use_secure_cookies() -> bool:
     return _bool_val(ENV_PHOENIX_USE_SECURE_COOKIES, False)
 

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/src/phoenix/db/facilitator.py

@@ -13,13 +13,13 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from phoenix.auth import (
     DEFAULT_ADMIN_EMAIL,
-    DEFAULT_ADMIN_PASSWORD,
     DEFAULT_ADMIN_USERNAME,
     DEFAULT_SECRET_LENGTH,
     DEFAULT_SYSTEM_EMAIL,
     DEFAULT_SYSTEM_USERNAME,
     compute_password_hash,
 )
+from phoenix.config import get_env_default_admin_initial_password
 from phoenix.db import models
 from phoenix.db.enums import COLUMN_ENUMS, UserRole
 from phoenix.server.types import DbSessionFactory
@@ -97,7 +97,8 @@ async def _ensure_user_roles(session: AsyncSession) -> None:
         admin_role_id := role_ids.get(admin_role)
     ) is not None:
         salt = secrets.token_bytes(DEFAULT_SECRET_LENGTH)
-        compute = partial(compute_password_hash, password=DEFAULT_ADMIN_PASSWORD, salt=salt)
+        password = get_env_default_admin_initial_password()
+        compute = partial(compute_password_hash, password=password, salt=salt)
         loop = asyncio.get_running_loop()
         hash_ = await loop.run_in_executor(None, compute)
         admin_user = models.User(

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/src/phoenix/server/api/helpers/playground_clients.py

@@ -2,6 +2,7 @@ import asyncio
 import importlib.util
 import inspect
 import json
+import os
 import time
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Callable, Iterator
@@ -25,6 +26,7 @@ from phoenix.evals.models.rate_limiters import (
     RateLimiter,
     RateLimitError,
 )
+from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.helpers.playground_registry import PROVIDER_DEFAULT, register_llm_client
 from phoenix.server.api.input_types.GenerativeModelInput import GenerativeModelInput
 from phoenix.server.api.input_types.InvocationParameters import (
@@ -99,9 +101,9 @@ class PlaygroundRateLimiter(RateLimiter, KeyedSingleton):
         super().__init__(
             rate_limit_error=rate_limit_error,
             max_rate_limit_retries=3,
-            initial_per_second_request_rate=2.0,
-            maximum_per_second_request_rate=10.0,
-            enforcement_window_minutes=1,
+            initial_per_second_request_rate=1.0,
+            maximum_per_second_request_rate=3.0,
+            enforcement_window_minutes=0.05,
             rate_reduction_factor=0.5,
             rate_increase_factor=0.01,
             cooldown_seconds=5,
@@ -128,10 +130,11 @@ class PlaygroundRateLimiter(RateLimiter, KeyedSingleton):
                     self._rate_limit_handling.set()  # Set the event as a failsafe
                 await self._throttler.async_wait_until_ready()
                 request_start_time = time.time()
-                if inspect.iscoroutinefunction(fn):
-                    return await fn(*args, **kwargs)  # type: ignore
+                maybe_coroutine = fn(*args, **kwargs)
+                if inspect.iscoroutine(maybe_coroutine):
+                    return await maybe_coroutine  # type: ignore
                 else:
-                    return fn(*args, **kwargs)
+                    return maybe_coroutine
             except self._rate_limit_error:
                 async with self._rate_limit_handling_lock:
                     self._rate_limit_handling.clear()  # prevent new requests from starting
@@ -258,6 +261,10 @@ class OpenAIStreamingClient(PlaygroundStreamingClient):
         from openai import AsyncOpenAI
         from openai import RateLimitError as OpenAIRateLimitError
 
+        # todo: check if custom base url is set before raising error to allow
+        # for custom endpoints that don't require an API key
+        if not (api_key := api_key or os.environ.get("OPENAI_API_KEY")):
+            raise BadRequest("An API key is required for OpenAI models")
         super().__init__(model=model, api_key=api_key)
         self._attributes[LLM_PROVIDER] = OpenInferenceLLMProviderValues.OPENAI.value
         self._attributes[LLM_SYSTEM] = OpenInferenceLLMSystemValues.OPENAI.value
@@ -341,7 +348,7 @@ class OpenAIStreamingClient(PlaygroundStreamingClient):
         openai_messages = [self.to_openai_chat_completion_param(*message) for message in messages]
         tool_call_ids: dict[int, str] = {}
         token_usage: Optional["CompletionUsage"] = None
-        throttled_create = self.rate_limiter.alimit(self.client.chat.completions.create)
+        throttled_create = self.rate_limiter._alimit(self.client.chat.completions.create)
         async for chunk in await throttled_create(
             messages=openai_messages,
             model=self.model_name,
@@ -510,7 +517,7 @@ class OpenAIO1StreamingClient(OpenAIStreamingClient):
 
         tool_call_ids: dict[int, str] = {}
 
-        throttled_create = self.rate_limiter.alimit(self.client.chat.completions.create)
+        throttled_create = self.rate_limiter._alimit(self.client.chat.completions.create)
         response = await throttled_create(
             messages=openai_messages,
             model=self.model_name,
@@ -618,12 +625,14 @@ class AzureOpenAIStreamingClient(OpenAIStreamingClient):
         super().__init__(model=model, api_key=api_key)
         self._attributes[LLM_PROVIDER] = OpenInferenceLLMProviderValues.AZURE.value
         self._attributes[LLM_SYSTEM] = OpenInferenceLLMSystemValues.OPENAI.value
-        if model.endpoint is None or model.api_version is None:
-            raise ValueError("endpoint and api_version are required for Azure OpenAI models")
+        if not (endpoint := model.endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")):
+            raise BadRequest("An Azure endpoint is required for Azure OpenAI models")
+        if not (api_version := model.api_version or os.environ.get("OPENAI_API_VERSION")):
+            raise BadRequest("An OpenAI API version is required for Azure OpenAI models")
         self.client = AsyncAzureOpenAI(
             api_key=api_key,
-            azure_endpoint=model.endpoint,
-            api_version=model.api_version,
+            azure_endpoint=endpoint,
+            api_version=api_version,
         )
 
 
@@ -648,6 +657,8 @@ class AnthropicStreamingClient(PlaygroundStreamingClient):
         super().__init__(model=model, api_key=api_key)
         self._attributes[LLM_PROVIDER] = OpenInferenceLLMProviderValues.ANTHROPIC.value
         self._attributes[LLM_SYSTEM] = OpenInferenceLLMSystemValues.ANTHROPIC.value
+        if not (api_key := api_key or os.environ.get("ANTHROPIC_API_KEY")):
+            raise BadRequest("An API key is required for Anthropic models")
         self.client = anthropic.AsyncAnthropic(api_key=api_key)
         self.model_name = model.name
         self.rate_limiter = PlaygroundRateLimiter(model.provider_key, anthropic.RateLimitError)
@@ -707,7 +718,6 @@ class AnthropicStreamingClient(PlaygroundStreamingClient):
             "messages": anthropic_messages,
             "model": self.model_name,
             "system": system_prompt,
-            "max_tokens": 1024,
             "tools": tools,
             **invocation_parameters,
         }
@@ -820,6 +830,12 @@ class GeminiStreamingClient(PlaygroundStreamingClient):
         super().__init__(model=model, api_key=api_key)
         self._attributes[LLM_PROVIDER] = OpenInferenceLLMProviderValues.GOOGLE.value
         self._attributes[LLM_SYSTEM] = OpenInferenceLLMSystemValues.VERTEXAI.value
+        if not (
+            api_key := api_key
+            or os.environ.get("GEMINI_API_KEY")
+            or os.environ.get("GOOGLE_API_KEY")
+        ):
+            raise BadRequest("An API key is required for Gemini models")
         google_genai.configure(api_key=api_key)
         self.model_name = model.name
 

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/src/phoenix/server/api/helpers/playground_spans.py

@@ -379,6 +379,11 @@ def _llm_output_messages(
     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_id:
+            yield (
+                f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_ID}",
+                _tool_call_id,
+            )
         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}",
@@ -441,6 +446,7 @@ MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
 MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
 MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS
 
+TOOL_CALL_ID = ToolCallAttributes.TOOL_CALL_ID
 TOOL_CALL_FUNCTION_NAME = ToolCallAttributes.TOOL_CALL_FUNCTION_NAME
 TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
 TOOL_CALL_ID = ToolCallAttributes.TOOL_CALL_ID

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/src/phoenix/server/api/mutations/chat_mutations.py

@@ -1,9 +1,9 @@
 import asyncio
 from dataclasses import asdict, field
 from datetime import datetime, timezone
-from itertools import chain
+from itertools import chain, islice
 from traceback import format_exc
-from typing import Any, Iterable, Iterator, List, Optional, Union
+from typing import Any, Iterable, Iterator, List, Optional, TypeVar, Union
 
 import strawberry
 from openinference.instrumentation import safe_json_dumps
@@ -26,7 +26,7 @@ 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.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,
@@ -127,11 +127,19 @@ class ChatCompletionMutationMixin:
         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)}"
+            )
         dataset_id = from_global_id_with_expected_type(input.dataset_id, Dataset.__name__)
         dataset_version_id = (
             from_global_id_with_expected_type(
@@ -158,7 +166,9 @@ class ChatCompletionMutationMixin:
             revisions = [
                 revision
                 async for revision in await session.stream_scalars(
-                    get_dataset_example_revisions(resolved_version_id)
+                    get_dataset_example_revisions(resolved_version_id).order_by(
+                        models.DatasetExampleRevision.id
+                    )
                 )
             ]
             if not revisions:
@@ -181,28 +191,32 @@ class ChatCompletionMutationMixin:
             session.add(experiment)
             await session.flush()
 
+        results = []
+        batch_size = 3
         start_time = datetime.now(timezone.utc)
-        results = await asyncio.gather(
-            *(
-                cls._chat_completion(
-                    info,
-                    llm_client,
-                    ChatCompletionInput(
-                        model=input.model,
-                        api_key=input.api_key,
-                        messages=input.messages,
-                        tools=input.tools,
-                        invocation_parameters=input.invocation_parameters,
-                        template=TemplateOptions(
-                            language=input.template_language,
-                            variables=revision.input,
+        for batch in _get_batches(revisions, batch_size):
+            batch_results = await asyncio.gather(
+                *(
+                    cls._chat_completion(
+                        info,
+                        llm_client,
+                        ChatCompletionInput(
+                            model=input.model,
+                            api_key=input.api_key,
+                            messages=input.messages,
+                            tools=input.tools,
+                            invocation_parameters=input.invocation_parameters,
+                            template=TemplateOptions(
+                                language=input.template_language,
+                                variables=revision.input,
+                            ),
                         ),
-                    ),
-                )
-                for revision in revisions
-            ),
-            return_exceptions=True,
-        )
+                    )
+                    for revision in batch
+                ),
+                return_exceptions=True,
+            )
+            results.extend(batch_results)
 
         payload = ChatCompletionOverDatasetMutationPayload(
             dataset_id=GlobalID(models.Dataset.__name__, str(dataset.id)),
@@ -266,11 +280,19 @@ class ChatCompletionMutationMixin:
         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)}"
+            )
         return await cls._chat_completion(info, llm_client, input)
 
     @classmethod
@@ -486,6 +508,11 @@ def _llm_output_messages(
     if text_content:
         yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", text_content
     for tool_call_index, tool_call in enumerate(tool_calls.values()):
+        if tool_call_id := tool_call.id:
+            yield (
+                f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_ID}",
+                tool_call_id,
+            )
         yield (
             f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
             tool_call.function.name,
@@ -513,6 +540,19 @@ def _serialize_event(event: SpanException) -> dict[str, Any]:
     return {k: (v.isoformat() if isinstance(v, datetime) else v) for k, v in asdict(event).items()}
 
 
+_AnyT = TypeVar("_AnyT")
+
+
+def _get_batches(
+    iterable: Iterable[_AnyT],
+    batch_size: int,
+) -> Iterator[list[_AnyT]]:
+    """Splits an iterable into batches not exceeding a specified size."""
+    iterator = iter(iterable)
+    while batch := list(islice(iterator, batch_size)):
+        yield batch
+
+
 JSON = OpenInferenceMimeTypeValues.JSON.value
 TEXT = OpenInferenceMimeTypeValues.TEXT.value
 LLM = OpenInferenceSpanKindValues.LLM.value
@@ -534,6 +574,7 @@ MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
 MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
 MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS
 
+TOOL_CALL_ID = ToolCallAttributes.TOOL_CALL_ID
 TOOL_CALL_FUNCTION_NAME = ToolCallAttributes.TOOL_CALL_FUNCTION_NAME
 TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
 

{arize_phoenix-5.10.0 → arize_phoenix-5.11.0}/src/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))