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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aidial-sdk
-Version: 0.48.0.dev4
+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.dev4 → 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.dev4 → 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.dev6/aidial_sdk/translators/messages_to_chat_completions/chat_completions/_reasoning_content.py

@@ -0,0 +1,95 @@
+import re
+from dataclasses import dataclass
+from typing import Literal
+
+from openai.types.chat import ChatCompletionMessage
+from openai.types.chat.chat_completion_chunk import ChoiceDelta
+
+from aidial_sdk.translators.messages_to_chat_completions.chat_completions.dial_extensions import (
+    CustomContent,
+    parse_extras,
+)
+
+_Source = Literal["azure", "claude", "dial"]
+
+_REASONING_STAGE = re.compile(r"think|thought|reason", re.IGNORECASE)
+
+
+@dataclass
+class ReasoningContent:
+    content: str | None
+    signature: str | None
+
+
+class ReasoningContentProvider:
+    """
+    Reads the reasoning of a message, or of a stream of message chunks, from
+    `reasoning_content`, the Claude state and the DIAL reasoning stages. They
+    carry the same text, chunked differently: only the text none of them has
+    delivered before is returned.
+    """
+
+    def __init__(self) -> None:
+        # Every source is a prefix of the same text, so offsets suffice.
+        self._returned = 0
+        self._received: dict[_Source, int] = {
+            "azure": 0,
+            "claude": 0,
+            "dial": 0,
+        }
+        self._reasoning_stages: set[int | None] = set()
+
+    def get_reasoning_content(
+        self, message: ChatCompletionMessage | ChoiceDelta
+    ) -> ReasoningContent:
+        custom_content = parse_extras(message.model_extra).custom_content
+        thinking, signature = _get_claude_reasoning_content(custom_content)
+        content = (
+            self._unseen("azure", _get_azure_reasoning_content(message))
+            + self._unseen("claude", thinking)
+            + self._unseen("dial", self._dial_reasoning_content(custom_content))
+        )
+        return ReasoningContent(
+            content=content or None,
+            signature=signature or None,
+        )
+
+    def _unseen(self, source: _Source, text: str | None) -> str:
+        text = text or ""
+        start = self._received[source]
+        self._received[source] += len(text)
+        unseen = text[self._returned - start :]
+        self._returned = max(self._returned, self._received[source])
+        return unseen
+
+    def _dial_reasoning_content(
+        self, custom_content: CustomContent | None
+    ) -> str:
+        text = ""
+        for stage in (custom_content.stages if custom_content else None) or []:
+            # In a stream, only the first chunk of a stage carries its name.
+            if _is_reasoning_stage(stage.name):
+                self._reasoning_stages.add(stage.index)
+            if stage.index in self._reasoning_stages:
+                text += stage.content or ""
+        return text
+
+
+def _get_azure_reasoning_content(
+    message: ChatCompletionMessage | ChoiceDelta,
+) -> str | None:
+    return (message.model_extra or {}).get("reasoning_content")
+
+
+def _get_claude_reasoning_content(
+    custom_content: CustomContent | None,
+) -> tuple[str | None, str | None]:
+    state = custom_content.state if custom_content else None
+    for block in (state.claude_message_content if state else None) or []:
+        if block.type == "thinking":
+            return block.thinking, block.signature
+    return None, None
+
+
+def _is_reasoning_stage(name: str | None) -> bool:
+    return name is not None and bool(_REASONING_STAGE.search(name))

{aidial_sdk-0.48.0.dev4 → 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.dev4 → aidial_sdk-0.48.0.dev6}/aidial_sdk/translators/messages_to_chat_completions/chat_completions/dial_extensions.py

@@ -1,4 +1,3 @@
-import re
 from collections.abc import Mapping
 
 from openai.types.chat.chat_completion_message import Annotation
@@ -7,8 +6,6 @@ from pydantic import ValidationError
 from aidial_sdk.utils.logging import logger as log
 from aidial_sdk.utils.pydantic import ExtraAllowModel
 
-_REASONING_STAGE = re.compile(r"think|thought|reason", re.IGNORECASE)
-
 
 class ClaudeMessageBlock(ExtraAllowModel):
     type: str | None = None
@@ -42,25 +39,3 @@ def parse_extras(model_extra: Mapping[str, object] | None) -> DialExtras:
     except ValidationError as e:
         log.error(f"Failed to parse DIAL extension fields: {e}")
         return DialExtras()
-
-
-def is_reasoning_stage(name: str | None) -> bool:
-    return name is not None and bool(_REASONING_STAGE.search(name))
-
-
-def get_claude_thinking_block(
-    custom_content: CustomContent | None,
-) -> ClaudeMessageBlock | None:
-    state = custom_content.state if custom_content else None
-    for block in (state.claude_message_content if state else None) or []:
-        if block.type == "thinking" and block.thinking:
-            return block
-    return None
-
-
-def get_dial_reasoning_content(custom_content: CustomContent | None) -> str:
-    return "".join(
-        stage.content or ""
-        for stage in (custom_content.stages if custom_content else None) or []
-        if is_reasoning_stage(stage.name)
-    )

{aidial_sdk-0.48.0.dev4 → aidial_sdk-0.48.0.dev6}/aidial_sdk/translators/messages_to_chat_completions/chat_completions/from_chat_completions.py

@@ -25,11 +25,8 @@ from openai.types.chat.chat_completion_message_tool_call import (
 )
 from openai.types.completion_usage import CompletionUsage, PromptTokensDetails
 
-from aidial_sdk.translators.messages_to_chat_completions.chat_completions.dial_extensions import (
-    CustomContent,
-    get_claude_thinking_block,
-    get_dial_reasoning_content,
-    parse_extras,
+from aidial_sdk.translators.messages_to_chat_completions.chat_completions._reasoning_content import (
+    ReasoningContentProvider,
 )
 from aidial_sdk.translators.messages_to_chat_completions.translation_state import (
     TranslationState,
@@ -76,8 +73,7 @@ def _to_anthropic_content_blocks(
 
     content: list[AnthropicContentBlock] = []
 
-    extras = parse_extras(message.model_extra)
-    if thinking := to_anthropic_thinking_block(extras.custom_content):
+    if thinking := to_anthropic_thinking_block(message):
         content.append(thinking)
 
     content.extend(
@@ -100,17 +96,16 @@ def _text_block(text: str) -> TextBlock:
 
 
 def to_anthropic_thinking_block(
-    custom_content: CustomContent | None,
+    message: ChatCompletionMessage,
 ) -> ThinkingBlock | None:
-    if block := get_claude_thinking_block(custom_content):
-        return ThinkingBlock(
-            type="thinking",
-            thinking=block.thinking or "",
-            signature=block.signature or "",
-        )
-    if text := get_dial_reasoning_content(custom_content):
-        return ThinkingBlock(type="thinking", thinking=text, signature="")
-    return None
+    reasoning = ReasoningContentProvider().get_reasoning_content(message)
+    if reasoning.content is None:
+        return None
+    return ThinkingBlock(
+        type="thinking",
+        thinking=reasoning.content,
+        signature=reasoning.signature or "",
+    )
 
 
 def _citations(annotations: list[Annotation] | None) -> list[tuple[str, str]]:

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

@@ -16,9 +16,13 @@ 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,
+)
 from aidial_sdk.translators.messages_to_chat_completions.chat_completions.dial_extensions import (
-    CustomContent,
-    is_reasoning_stage,
     parse_extras,
 )
 from aidial_sdk.translators.messages_to_chat_completions.chat_completions.from_chat_completions import (
@@ -55,15 +59,18 @@ class ChatCompletionsToAnthropicStream(AnthropicStreamState):
         self.is_refusal = False
         self.has_tool_call = False
         self._translation_state = translation_state
-        self._reasoning_stages: set[int | None] = set()
+        self._reasoning = ReasoningContentProvider()
         self._seen_citations: set[str] = set()
-        self._signed = False
         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()
@@ -84,7 +91,12 @@ class ChatCompletionsToAnthropicStream(AnthropicStreamState):
 
         if choice.finish_reason:
             self.finish_reason = choice.finish_reason
-        yield from self._on_custom_content(extras.custom_content)
+
+        reasoning = self._reasoning.get_reasoning_content(delta)
+        if reasoning.content:
+            yield from self._thinking_delta(reasoning.content)
+        if reasoning.signature:
+            yield from self._signature_delta(reasoning.signature)
 
         if text := delta.content:
             yield from self._emit_text(text)
@@ -106,26 +118,6 @@ class ChatCompletionsToAnthropicStream(AnthropicStreamState):
         for call in delta.tool_calls or []:
             yield from self._on_tool_call_delta(call)
 
-    def _on_custom_content(
-        self, custom_content: CustomContent | None
-    ) -> Iterator[bytes]:
-        if custom_content is None:
-            return
-
-        for stage in custom_content.stages or []:
-            if is_reasoning_stage(stage.name):
-                self._reasoning_stages.add(stage.index)
-            if stage.index in self._reasoning_stages and stage.content:
-                yield from self._thinking_delta(stage.content)
-
-        if self._signed or custom_content.state is None:
-            return
-        for block in custom_content.state.claude_message_content or []:
-            if block.type == "thinking" and block.signature:
-                self._signed = True
-                yield from self._signature_delta(block.signature)
-                break
-
     def _thinking_delta(self, text: str) -> Iterator[bytes]:
         yield from self._close_tools()
         if not self.is_open(_THINKING_KEY):
@@ -228,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.dev4 → 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.dev4 → aidial_sdk-0.48.0.dev6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aidial-sdk"
-version = "0.48.0.dev4"
+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"