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

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

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

aidial_sdk-0.48.0.dev5/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.dev5}/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.dev5}/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.dev5}/aidial_sdk/translators/messages_to_chat_completions/chat_completions/streaming.py

@@ -16,9 +16,10 @@ 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.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,9 +56,8 @@ 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]:
@@ -84,7 +84,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 +111,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):

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

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