aidial-sdk 0.48.0.dev5__tar.gz → 0.48.0.dev6__tar.gz

{aidial_sdk-0.48.0.dev5 → aidial_sdk-0.48.0.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aidial-sdk
-Version: 0.48.0.dev5
+Version: 0.48.0.dev6
 Summary: Framework to create applications and model adapters for AI DIAL
 License-Expression: Apache-2.0
 License-File: LICENSE

aidial_sdk-0.48.0.dev6/aidial_sdk/_constants.py

@@ -0,0 +1,3 @@
+from aidial_sdk.utils.env import env_float
+
+DIAL_SDK_SSE_HEARTBEAT_INTERVAL = env_float("DIAL_SDK_SSE_HEARTBEAT_INTERVAL")

{aidial_sdk-0.48.0.dev5 → aidial_sdk-0.48.0.dev6}/aidial_sdk/application.py

@@ -9,6 +9,7 @@ from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 from starlette.types import Receive, Scope, Send
 
+from aidial_sdk._constants import DIAL_SDK_SSE_HEARTBEAT_INTERVAL
 from aidial_sdk._errors import (
     Handler,
     dial_exception_handler,
@@ -45,7 +46,7 @@ from aidial_sdk.utils._interpolation import (
 )
 from aidial_sdk.utils._logging_middleware import LoggingMiddleware
 from aidial_sdk.utils._reflection import get_method_implementation
-from aidial_sdk.utils.env import env_float, env_var_list
+from aidial_sdk.utils.env import env_var_list
 from aidial_sdk.utils.log_config import configure_sdk_logger
 from aidial_sdk.utils.logging import (
     reset_log_context,
@@ -53,7 +54,7 @@ from aidial_sdk.utils.logging import (
 )
 from aidial_sdk.utils.pydantic import model_validate_extra_fields
 from aidial_sdk.utils.streaming import (
-    add_heartbeat,
+    start_streaming,
     to_block_response,
     to_streaming_response,
 )
@@ -203,6 +204,10 @@ class DIALApp(FastAPI):
         *,
         heartbeat_interval: float | None = None,
     ) -> "DIALApp":
+        heartbeat_interval = (
+            heartbeat_interval or DIAL_SDK_SSE_HEARTBEAT_INTERVAL
+        )
+
         for endpoint, handler, methods in self._chat_completion_handlers(
             deployment_name, impl, heartbeat_interval
         ):
@@ -394,10 +399,6 @@ class DIALApp(FastAPI):
         *,
         heartbeat_interval: float | None,
     ) -> Handler:
-        heartbeat_interval = heartbeat_interval or env_float(
-            "DIAL_SDK_SSE_HEARTBEAT_INTERVAL"
-        )
-
         async def _handler(original_request: Request) -> Response:
             request = await self._parse_request(
                 ChatCompletionRequest, original_request, deployment_id
@@ -407,15 +408,13 @@ class DIALApp(FastAPI):
             stream = response._generate_stream(impl.chat_completion)
 
             if request.stream:
-                if heartbeat_interval:
-                    stream = add_heartbeat(
-                        stream,
-                        heartbeat_interval=heartbeat_interval,
-                        heartbeat_object=": heartbeat\n\n",
-                    )
-
+                sse_stream = await start_streaming(
+                    stream,
+                    heartbeat_interval=heartbeat_interval,
+                    heartbeat=": heartbeat\n\n",
+                )
                 resp = StreamingResponse(
-                    await to_streaming_response(stream),
+                    await to_streaming_response(sse_stream),
                     media_type="text/event-stream",
                 )
             else:

aidial_sdk-0.48.0.dev6/aidial_sdk/translators/messages_to_chat_completions/backend/client.py

@@ -0,0 +1,5 @@
+class Heartbeat:
+    """Yielded by a chat stream for each heartbeat interval it was idle."""
+
+
+HEARTBEAT = Heartbeat()

{aidial_sdk-0.48.0.dev5 → aidial_sdk-0.48.0.dev6}/aidial_sdk/translators/messages_to_chat_completions/backend/http.py

@@ -10,10 +10,15 @@ from openai import AsyncOpenAI, AsyncStream, Omit, omit
 from openai.types.chat import ChatCompletion, ChatCompletionChunk
 
 from aidial_sdk.exceptions import HTTPException as DIALException
+from aidial_sdk.translators.messages_to_chat_completions.backend.client import (
+    HEARTBEAT,
+    Heartbeat,
+)
 from aidial_sdk.translators.messages_to_chat_completions.settings import (
     get_dial_api_version,
 )
 from aidial_sdk.utils.logging import logger as log
+from aidial_sdk.utils.streaming import with_heartbeat
 
 # The headers of the inbound request that DIAL Core reads when the translator
 # calls it back. Any other header is not forwarded.
@@ -122,12 +127,16 @@ def _core_errors() -> Iterator[None]:
         raise _from_core_error(e) from e
 
 
-async def _chunks(
-    stream: AsyncStream[ChatCompletionChunk],
-) -> AsyncGenerator[ChatCompletionChunk, None]:
+async def _with_heartbeat(
+    stream: AsyncStream[ChatCompletionChunk], heartbeat_interval: float | None
+) -> AsyncGenerator[ChatCompletionChunk | Heartbeat, None]:
     with _core_errors():
         async with stream:
-            async for item in stream:
+            async for item in with_heartbeat(
+                aiter(stream),
+                heartbeat_interval=heartbeat_interval,
+                heartbeat=HEARTBEAT,
+            ):
                 yield item
 
 
@@ -135,10 +144,11 @@ async def _chunks(
 class ChatClientHttp:
     client: AsyncOpenAI
     core_headers: dict[str, str | Omit]
+    heartbeat_interval: float | None
 
     async def create_stream(
         self, body: dict
-    ) -> AsyncGenerator[ChatCompletionChunk, None]:
+    ) -> AsyncGenerator[ChatCompletionChunk | Heartbeat, None]:
         with _core_errors():
             stream = await self.client.post(
                 "/chat/completions",
@@ -148,7 +158,7 @@ class ChatClientHttp:
                 stream=True,
                 stream_cls=AsyncStream[ChatCompletionChunk],
             )
-        return _chunks(stream)
+        return _with_heartbeat(stream, self.heartbeat_interval)
 
     async def create_block(self, body: dict) -> ChatCompletion:
         with _core_errors():

{aidial_sdk-0.48.0.dev5 → aidial_sdk-0.48.0.dev6}/aidial_sdk/translators/messages_to_chat_completions/chat_completions/app.py

@@ -8,6 +8,7 @@ from anthropic.types.beta.message_create_params import MessageCreateParams
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 
+from aidial_sdk._constants import DIAL_SDK_SSE_HEARTBEAT_INTERVAL
 from aidial_sdk._errors import error_handler
 from aidial_sdk.deployment.from_request_mixin import get_client_deployment_id
 from aidial_sdk.exceptions import InvalidRequestError, ResourceNotFoundError
@@ -94,7 +95,9 @@ def create_anthropic_api_app(
         else:
             client = core_v1_client(dial_url, http_client)
 
-        return ChatClientHttp(client, core_headers(request))
+        return ChatClientHttp(
+            client, core_headers(request), DIAL_SDK_SSE_HEARTBEAT_INTERVAL
+        )
 
     @app.post("/anthropic/v1/messages")
     @error_handler("anthropic")

{aidial_sdk-0.48.0.dev5 → aidial_sdk-0.48.0.dev6}/aidial_sdk/translators/messages_to_chat_completions/chat_completions/streaming.py

@@ -16,6 +16,9 @@ from openai.types.chat import ChatCompletionChunk
 from openai.types.chat.chat_completion_chunk import Choice, ChoiceDeltaToolCall
 from openai.types.completion_usage import CompletionUsage
 
+from aidial_sdk.translators.messages_to_chat_completions.backend.client import (
+    Heartbeat,
+)
 from aidial_sdk.translators.messages_to_chat_completions.chat_completions._reasoning_content import (
     ReasoningContentProvider,
 )
@@ -60,10 +63,14 @@ class ChatCompletionsToAnthropicStream(AnthropicStreamState):
         self._seen_citations: set[str] = set()
         self._tools: dict[int, int] = {}
 
-    def handle(self, chunk: ChatCompletionChunk) -> Iterator[bytes]:
+    def handle(self, chunk: ChatCompletionChunk | Heartbeat) -> Iterator[bytes]:
         if self.terminated:
             return
 
+        if isinstance(chunk, Heartbeat):
+            yield format_sse({"type": "ping"})
+            return
+
         if not self.started:
             self.message_id = chunk.id
             yield from self._start()
@@ -213,7 +220,7 @@ def _merge_usage(
 
 
 async def translate_stream(
-    stream: AsyncGenerator[ChatCompletionChunk, None],
+    stream: AsyncGenerator[ChatCompletionChunk | Heartbeat, None],
     requested_model: str,
     translation_state: TranslationState,
 ) -> AsyncIterator[bytes]:

{aidial_sdk-0.48.0.dev5 → aidial_sdk-0.48.0.dev6}/aidial_sdk/utils/streaming.py

@@ -151,3 +151,46 @@ async def add_heartbeat(
 
                 if heartbeat_callback is not None:
                     await _call_heartbeat_callback(heartbeat_callback)
+
+
+def with_heartbeat(
+    stream: AsyncIterator[_T],
+    *,
+    heartbeat_interval: float | None,
+    heartbeat: _T,
+) -> AsyncIterator[_T]:
+    if not heartbeat_interval:
+        return stream
+    return add_heartbeat(
+        stream,
+        heartbeat_interval=heartbeat_interval,
+        heartbeat_object=heartbeat,
+    )
+
+
+async def start_streaming(
+    stream: AsyncIterator[_T],
+    *,
+    heartbeat_interval: float | None,
+    heartbeat: _T,
+) -> AsyncIterator[_T]:
+    """
+    Waits for the first item of the stream or, while it is idle, the first
+    heartbeat. A DIALException as the first item is raised, so that it becomes
+    an HTTP error, while a later one is left in the stream.
+    """
+
+    stream = with_heartbeat(
+        stream, heartbeat_interval=heartbeat_interval, heartbeat=heartbeat
+    )
+
+    first_item = await anext(stream)
+    if isinstance(first_item, DIALException):
+        raise first_item
+
+    async def _stream() -> AsyncIterator[_T]:
+        yield first_item
+        async for item in stream:
+            yield item
+
+    return _stream()

{aidial_sdk-0.48.0.dev5 → aidial_sdk-0.48.0.dev6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aidial-sdk"
-version = "0.48.0.dev5"
+version = "0.48.0.dev6"
 description = "Framework to create applications and model adapters for AI DIAL"
 authors = [{ name = "EPAM RAIL", email = "SpecialEPM-DIALDevTeam@epam.com" }]
 license = "Apache-2.0"