File size: 21,960 Bytes
9792ea7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
# -*- coding: utf-8 -*-
"""The DashScope chat model class (OpenAI-compatible implementation)."""
import base64
import warnings
from collections import OrderedDict
from datetime import datetime
from typing import Any, AsyncGenerator, List, Literal, Type, TYPE_CHECKING

from pydantic import BaseModel, Field

from ..._utils._audio import _build_streaming_wav_header
from ..._utils._common import _generate_id
from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES
from .._model_response import ChatResponse, StructuredResponse
from .._model_usage import ChatUsage
from ...credential import DashScopeCredential
from ...formatter import FormatterBase, DashScopeChatFormatter
from ...message import (
    Msg,
    TextBlock,
    ThinkingBlock,
    ToolCallBlock,
)
from ...tool import ToolChoice

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletion
    from openai import AsyncStream
else:
    ChatCompletion = Any
    AsyncStream = Any


class DashScopeChatModel(ChatModelBase):
    """The DashScope chat model (OpenAI-compatible implementation).

    This implementation uses the OpenAI Python SDK to call DashScope's
    OpenAI-compatible endpoint (``compatible-mode/v1``), which supports
    both text-only and multimodal (image/video) inputs through the same
    unified API.
    """

    class Parameters(BaseModel):
        """The parameters for DashScope LLM API."""

        max_tokens: int | None = Field(
            default=None,
            title="Max Tokens",
            description="The maximum number of tokens for the LLM output.",
            gt=0,
        )

        thinking_enable: bool = Field(
            default=False,
            title="Thinking",
            description="The thinking enable for the LLM output.",
        )

        thinking_budget: int | None = Field(
            default=None,
            title="Thinking Budget",
            description="The thinking budget for the LLM output.",
            gt=0,
        )

        temperature: float | None = Field(
            default=None,
            title="Temperature",
            description="The temperature for the LLM output.",
            ge=0,
            lt=2,
        )

        top_p: float | None = Field(
            default=None,
            title="Top P",
            description="The top P value for the LLM output.",
            gt=0,
            le=1,
        )

        top_k: int | None = Field(
            default=None,
            title="Top K",
            description="The top K value for the LLM output.",
            gt=0,
            le=100,
        )

        parallel_tool_calls: bool = Field(
            default=True,
            title="Parallel Tool Calls",
            description="If enable parallel tool calls for the LLM output.",
        )

        voice: str | None = Field(
            default=None,
            title="Voice",
            description=(
                "Voice for audio output on omni-style models (e.g. "
                "``qwen3.5-omni-plus``). Setting this implicitly asks the "
                "model to speak its response — ``modalities`` is filled in "
                "automatically. Supported voices vary by model — see the "
                "model card's ``voice.suggestions``. Any value the API "
                "accepts works — the suggestions are convenience-only. "
                "Leave unset for text-only "
                "responses."
            ),
        )

    type: Literal["dashscope_chat"] = "dashscope_chat"
    """The type of the chat model."""

    def __init__(
        self,
        credential: DashScopeCredential,
        model: str,
        parameters: "DashScopeChatModel.Parameters | None" = None,
        stream: bool = True,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        context_size: int = 131072,
        formatter: FormatterBase | None = None,
        client_kwargs: dict[str, Any] | None = None,
    ) -> None:
        """Initialize the DashScope chat model.

        Args:
            credential (`DashScopeCredential`):
                The DashScope credential used to authenticate API calls.
            model (`str`):
                The DashScope model name, e.g. ``qwen-plus``.
            parameters (`DashScopeChatModel.Parameters | None`, defaults to \
            `None`):
                The DashScope API parameters. When ``None``, the default
                parameters will be used.
            stream (`bool`, defaults to `True`):
                Whether to enable streaming output.
            max_retries (`int`, defaults to `3`):
                The maximum number of retries for the DashScope API.
            retry_delay (`float`, defaults to `1.0`):
                Seconds to sleep between retry attempts.
            context_size (`int`, defaults to `131072`):
                The model context size used for context compression.
            formatter (`FormatterBase | None`, defaults to `None`):
                The formatter that converts ``Msg`` objects to the format
                required by the DashScope API. When ``None``, a
                ``DashScopeChatFormatter`` instance will be used.
            client_kwargs (`dict[str, Any] | None`, defaults to `None`):
                Extra keyword arguments forwarded to ``openai.AsyncClient``
                (e.g. ``timeout``, ``default_headers``, ``http_client``).
        """
        super().__init__(
            credential=credential,
            model=model,
            parameters=parameters or self.Parameters(),
            stream=stream,
            max_retries=max_retries,
            retry_delay=retry_delay,
            context_size=context_size,
        )
        self.formatter = formatter or DashScopeChatFormatter()
        self.client_kwargs = client_kwargs or {}

    @classmethod
    def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]:
        import openai

        return (
            openai.APIConnectionError,
            openai.APITimeoutError,
            openai.RateLimitError,
            openai.InternalServerError,
        )

    async def _call_api(
        self,
        model_name: str,
        messages: list[Msg],
        tools: list[dict] | None = None,
        tool_choice: ToolChoice | None = None,
        **kwargs: Any,
    ) -> ChatResponse | AsyncGenerator[ChatResponse, None]:
        """Call the DashScope chat completions API via OpenAI-compatible
        endpoint.

        Args:
            model_name (`str`):
                The model name to use for this call.
            messages (`list`):
                The Msg objects that will be formatted and sent to the API.
            tools (`list[dict] | None`, default `None`):
                The tools JSON schemas that the model can use.
            tool_choice (`ToolChoice | None`, default `None`):
                Controls which (if any) tool is called by the model.
            **kwargs (`Any`):
                The keyword arguments for DashScope chat completions API,
                e.g. ``temperature``, ``max_tokens``, ``top_p``, etc.
        """
        import openai

        client = openai.AsyncClient(
            **{
                "api_key": self.credential.api_key.get_secret_value(),
                "base_url": self.credential.base_url,
                **self.client_kwargs,
            },
        )

        formatted_messages = await self.formatter.format(messages)

        request_kwargs: dict[str, Any] = {
            "model": model_name,
            "messages": formatted_messages,
            "stream": self.stream,
        }

        if self.parameters.max_tokens is not None:
            request_kwargs["max_tokens"] = self.parameters.max_tokens

        if self.parameters.temperature is not None:
            request_kwargs["temperature"] = self.parameters.temperature

        if self.parameters.top_p is not None:
            request_kwargs["top_p"] = self.parameters.top_p

        if self.parameters.voice is not None:
            # Requesting audio output implies ``modalities`` must include
            # ``"audio"``; set it automatically so callers don't have to.
            # ``format`` is forced to ``pcm16``: omni streaming delivers raw
            # PCM upstream regardless of the requested format, and we wrap
            # it as WAV in ``_parse_stream_response`` before yielding.
            request_kwargs["audio"] = {
                "voice": self.parameters.voice,
                "format": "pcm16",
            }
            request_kwargs["modalities"] = ["text", "audio"]

        request_kwargs.update(kwargs)

        fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice)
        if fmt_tools is not None:
            request_kwargs["tools"] = fmt_tools
            if not self.parameters.parallel_tool_calls:
                request_kwargs["parallel_tool_calls"] = False
        if fmt_tool_choice is not None:
            request_kwargs["tool_choice"] = fmt_tool_choice

        extra_body: dict[str, Any] = {}
        if self.parameters.thinking_enable is not None:
            extra_body["enable_thinking"] = self.parameters.thinking_enable
        if self.parameters.thinking_budget is not None:
            extra_body["thinking_budget"] = self.parameters.thinking_budget
        if self.parameters.top_k is not None:
            extra_body["top_k"] = self.parameters.top_k

        if extra_body:
            request_kwargs.setdefault("extra_body", {})
            request_kwargs["extra_body"].update(extra_body)

        if self.stream:
            request_kwargs["stream_options"] = {"include_usage": True}

        start_datetime = datetime.now()
        response = await client.chat.completions.create(**request_kwargs)

        if self.stream:
            return self._parse_stream_response(start_datetime, response)

        return self._parse_completion_response(start_datetime, response)

    async def _parse_stream_response(
        self,
        start_datetime: datetime,
        response: AsyncStream,
    ) -> AsyncGenerator[ChatResponse, None]:
        """Parse the DashScope streaming response (OpenAI-compatible format).

        Args:
            start_datetime (`datetime`):
                The start datetime of the response generation.
            response (`AsyncStream`):
                The OpenAI-compatible async stream object.

        Yields:
            `ChatResponse`:
                Incremental ``ChatResponse`` objects with ``is_last=False``
                followed by a final one with ``is_last=True``.
        """
        # ``True`` once the first audio chunk has been prefixed with a
        # streaming WAV header and yielded.
        audio_header_sent: bool = False

        usage = None
        response_id: str = _generate_id()
        text_id: str = _generate_id()
        thinking_id: str = _generate_id()
        audio_id = _generate_id()
        # The mapping from index to tool call id
        tool_call_mapping: dict = OrderedDict()

        async with response as stream:
            async for chunk in stream:
                delta_res = ChatResponse(
                    content=[],
                    is_last=False,
                    id=response_id,
                )

                # Update the response ID if exists
                response_id = getattr(chunk, "id", None) or response_id
                delta_res.id = response_id

                if chunk.usage:
                    u = chunk.usage
                    ptd = getattr(u, "prompt_tokens_details", None)
                    if ptd and hasattr(ptd, "cached_tokens"):
                        cache_read = ptd.cached_tokens or 0
                    else:
                        cache_read = 0
                    usage = ChatUsage(
                        input_tokens=u.prompt_tokens or 0,
                        output_tokens=u.completion_tokens or 0,
                        time=(datetime.now() - start_datetime).total_seconds(),
                        cache_input_tokens=cache_read,
                    )

                if not chunk.choices:
                    if usage is not None:
                        delta_res.usage = usage
                        yield delta_res
                    continue

                choice = chunk.choices[0]
                delta = choice.delta

                # Thinking
                if getattr(delta, "reasoning_content", None):
                    delta_res.append_thinking(
                        block_id=thinking_id,
                        thinking=delta.reasoning_content,
                    )

                # Text
                if getattr(delta, "content", None):
                    delta_res.append_text(
                        block_id=text_id,
                        text=delta.content,
                    )

                # Tool call
                for tool_call in getattr(delta, "tool_calls", None) or []:
                    index = tool_call.index
                    fn = getattr(tool_call, "function", None)
                    delta_name = getattr(fn, "name", None) if fn else None
                    delta_args = getattr(fn, "arguments", None) if fn else None

                    # Record the id and name in case following deltas
                    # don't provide them
                    if index not in tool_call_mapping:
                        tool_call_mapping[index] = (
                            tool_call.id,
                            delta_name or "unknown",
                        )

                    stored_id, stored_name = tool_call_mapping[index]

                    delta_res.append_tool_call(
                        block_id=tool_call.id or stored_id,
                        name=delta_name or stored_name,
                        input=delta_args or "",
                    )

                # Data block
                if getattr(delta, "audio", None):
                    delta_audio = getattr(delta, "audio", None)
                    if isinstance(delta_audio, dict):
                        audio_chunk = delta_audio.get("data", "")
                    else:
                        audio_chunk = getattr(delta_audio, "data", "")

                    if audio_chunk:
                        pcm_bytes = base64.b64decode(audio_chunk)
                        if not audio_header_sent:
                            payload = _build_streaming_wav_header() + pcm_bytes
                            audio_header_sent = True
                        else:
                            payload = pcm_bytes
                        # ``append_data_block`` expects the raw incremental
                        # media bytes and handles base64 encoding internally
                        # (see ``ChatResponse.append_data_block``); passing an
                        # already base64-encoded string here would result in
                        # double-encoding.
                        delta_res.append_data_block(
                            block_id=audio_id,
                            data=payload,
                            media_type="audio/wav",
                        )

                if delta_res.content or usage:
                    delta_res.usage = usage
                    yield delta_res

    def _parse_completion_response(
        self,
        start_datetime: datetime,
        response: ChatCompletion,
    ) -> ChatResponse:
        """Parse the DashScope non-streaming response (OpenAI-compatible
        format).

        Args:
            start_datetime (`datetime`):
                The start datetime of the response generation.
            response (`ChatCompletion`):
                The OpenAI-compatible chat completion object.

        Returns:
            `ChatResponse`:
                A single ``ChatResponse`` with ``is_last=True``.
        """
        content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = []

        if response.choices:
            choice = response.choices[0]
            reasoning = getattr(choice.message, "reasoning_content", None)
            if isinstance(reasoning, str) and reasoning:
                content_blocks.append(ThinkingBlock(thinking=reasoning))

            if choice.message.content:
                content_blocks.append(TextBlock(text=choice.message.content))

            for tool_call in choice.message.tool_calls or []:
                content_blocks.append(
                    ToolCallBlock(
                        id=tool_call.id,
                        name=tool_call.function.name,
                        input=tool_call.function.arguments,
                    ),
                )

        usage = None
        if response.usage:
            u = response.usage
            ptd = getattr(u, "prompt_tokens_details", None)
            if ptd and hasattr(ptd, "cached_tokens"):
                cache_read = ptd.cached_tokens or 0
            else:
                cache_read = 0
            usage = ChatUsage(
                input_tokens=u.prompt_tokens,
                output_tokens=u.completion_tokens,
                time=(datetime.now() - start_datetime).total_seconds(),
                cache_input_tokens=cache_read,
            )

        resp_kwargs: dict[str, Any] = {
            "content": content_blocks,
            "is_last": True,
            "usage": usage,
        }
        response_id = getattr(response, "id", None)
        if response_id:
            resp_kwargs["id"] = response_id

        return ChatResponse(**resp_kwargs)

    def _format_tools(
        self,
        tools: list[dict] | None,
        tool_choice: ToolChoice | None,
    ) -> tuple[list[dict] | None, str | dict | None]:
        """Validate and format tools and tool_choice for DashScope.

        DashScope supports "auto", "none", and "required" modes in
        OpenAI-compatible format. When ``tool_choice.tools`` is specified
        the schemas list is filtered to only those tools. When
        ``tool_choice.mode`` is a specific tool name (str) the model is
        forced to call exactly that tool.

        Args:
            tools (`list[dict] | None`, optional):
                The raw tool schemas.
            tool_choice (`ToolChoice | None`, optional):
                The tool choice configuration.

        Returns:
            `tuple[list[dict] | None, str | dict | None]`:
                A tuple of (formatted_tools, formatted_tool_choice).
        """
        if tool_choice and tools:
            self._validate_tool_choice(tool_choice, tools)
            if tool_choice.tools:
                allowed = set(tool_choice.tools)
                tools = [t for t in tools if t["function"]["name"] in allowed]

        fmt_tools = None
        if tools:
            for value in tools:
                if (
                    not isinstance(value, dict)
                    or "type" not in value
                    or value["type"] != "function"
                    or "function" not in value
                ):
                    raise ValueError(
                        f"Each schema must be a dict with 'type' as "
                        f"'function' and 'function' key, got {value}",
                    )
            fmt_tools = tools

        if not tool_choice:
            return fmt_tools, None

        mode = tool_choice.mode

        if mode not in _TOOL_CHOICE_LITERAL_MODES:
            return fmt_tools, {
                "type": "function",
                "function": {"name": mode},
            }

        if mode == "required":
            warnings.warn(
                f"'{mode}' is not fully supported by DashScope API. "
                "It will be converted to 'auto'.",
                DeprecationWarning,
                stacklevel=2,
            )
            return fmt_tools, "auto"

        return fmt_tools, mode

    async def _call_api_with_structured_output(
        self,
        model_name: str,
        messages: list[Msg],
        structured_model: Type[BaseModel] | dict,
        tool_choice: ToolChoice | None = None,
        **kwargs: Any,
    ) -> StructuredResponse:
        """DashScope-specific override for structured output.

        DashScope rejects ``tool_choice="required"`` or an object-form
        ``tool_choice`` when thinking mode is enabled. In that case we
        default ``tool_choice`` to ``"auto"`` and rely on the base class's
        injected system-reminder prompt to guide the model. When thinking
        is disabled, this falls through to the base implementation.

        See: https://help.aliyun.com/en/model-studio/qwen-function-calling

        Args:
            model_name (`str`):
                The model name to use for this call.
            messages (`list[Msg]`):
                The context for the LLM to generate the structured output.
            structured_model (`Type[BaseModel] | dict`):
                A Pydantic model class or a JSON schema dict describing the
                required output structure.
            tool_choice (`ToolChoice | None`, defaults to `None`):
                The tool_choice forwarded to ``_call_api``. When ``None``
                and thinking mode is enabled, it is downgraded to
                ``ToolChoice(mode="auto")``; otherwise the base default
                (force the structured-output tool) is used.
            **kwargs (`Any`):
                Additional keyword arguments forwarded to ``_call_api``.

        Returns:
            `StructuredResponse`:
                The structured response whose ``content`` is the validated
                output dict matching ``structured_model``.
        """
        if tool_choice is None and self.parameters.thinking_enable:
            tool_choice = ToolChoice(mode="auto")
        return await super()._call_api_with_structured_output(
            model_name=model_name,
            messages=messages,
            structured_model=structured_model,
            tool_choice=tool_choice,
            **kwargs,
        )