Agentic RL support cancel async call

README.md

@@ -1,12 +1,3 @@
-<h4 align="center">
-    <p>
-        <b>English</b>
-    <p>
-</h4>
-
-
-</div>
-
 # DashScope Python SDK
 
 The DashScope Python SDK provides a comprehensive interface to [Alibaba Cloud Model Studio (Bailian)](https://www.alibabacloud.com/help/en/model-studio/) APIs, covering text generation, multi-modal understanding, embeddings, reranking, image/video generation, speech synthesis & recognition, and more.

dashscope/__init__.py

@@ -21,6 +21,7 @@
     HttpSpeechSynthesizer,
 )
 from dashscope.audio.tts.speech_synthesizer import SpeechSynthesizer
+from dashscope.api_entities.aio_session import close_shared_aio_session
 from dashscope.common.api_key import save_api_key
 from dashscope.common.env import (
     api_key,
@@ -46,7 +47,7 @@
 from dashscope.files import Files
 from dashscope.models import Models
 from dashscope.nlp.understanding import Understanding
-from dashscope.rerank.text_rerank import TextReRank
+from dashscope.rerank import AioTextReRank, TextReRank
 from dashscope.threads import (
     MessageFile,
     Messages,
@@ -75,6 +76,7 @@
     "api_key",
     "api_key_file_path",
     "save_api_key",
+    "close_shared_aio_session",
     "AioGeneration",
     "Conversation",
     "Generation",
@@ -106,6 +108,7 @@
     "list_tokenizers",
     "Application",
     "TextReRank",
+    "AioTextReRank",
     "Assistants",
     "Threads",
     "Messages",

dashscope/aigc/generation.py

@@ -179,11 +179,13 @@ def call(  # pylint: disable=arguments-renamed,too-many-branches,too-many-statem
             to_merge_incremental_output = True
             parameters["incremental_output"] = True
 
-        # Pass incremental_to_full flag via headers user-agent
-        if "headers" not in parameters:
-            parameters["headers"] = {}
+        # Pass incremental_to_full flag via user_agent parameter
         flag = "1" if to_merge_incremental_output else "0"
-        parameters["headers"]["user-agent"] = f"incremental_to_full/{flag}"
+        existing_ua = parameters.get("user_agent", "")
+        new_ua = f"incremental_to_full/{flag}"
+        parameters["user_agent"] = (
+            f"{existing_ua}; {new_ua}".strip() if existing_ua else new_ua
+        )
 
         response = super().call(
             model=model,
@@ -434,11 +436,13 @@ async def call(  # type: ignore[override] # pylint: disable=arguments-renamed,to
             to_merge_incremental_output = True
             parameters["incremental_output"] = True
 
-        # Pass incremental_to_full flag via headers user-agent
-        if "headers" not in parameters:
-            parameters["headers"] = {}
+        # Pass incremental_to_full flag via user_agent parameter
         flag = "1" if to_merge_incremental_output else "0"
-        parameters["headers"]["user-agent"] = f"incremental_to_full/{flag}"
+        existing_ua = parameters.get("user_agent", "")
+        new_ua = f"incremental_to_full/{flag}"
+        parameters["user_agent"] = (
+            f"{existing_ua}; {new_ua}".strip() if existing_ua else new_ua
+        )
 
         response = await super().call(
             model=model,

dashscope/aigc/image_generation.py

@@ -99,14 +99,15 @@ def call(  # type: ignore[override]
             to_merge_incremental_output = True
             kwargs["incremental_output"] = True
 
-        # Pass incremental_to_full flag via headers user-agent
-        if "headers" not in kwargs:
-            kwargs["headers"] = {}
-
+        # Pass incremental_to_full flag via user_agent parameter
         flag = "1" if to_merge_incremental_output else "0"
-        kwargs["headers"]["user-agent"] = f"incremental_to_full/{flag}"
+        existing_ua = kwargs.get("user_agent", "")
+        new_ua = f"incremental_to_full/{flag}"
+        kwargs["user_agent"] = (
+            f"{existing_ua}; {new_ua}".strip() if existing_ua else new_ua
+        )
         if kwargs.get("is_async", False):
-            kwargs["headers"]["X-DashScope-Async"] = "enable"
+            kwargs.setdefault("headers", {})["X-DashScope-Async"] = "enable"
             task = cls.async_task
         else:
             task = cls.sync_task
@@ -177,6 +178,7 @@ def wait(
         task: Union[str, ImageGenerationResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> DashScopeAPIResponse:
         """Wait for image(s) synthesis task to complete, and return the result.
@@ -186,11 +188,18 @@ def wait(
                 ImageGenerationResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             DashScopeAPIResponse: The task result.
         """
-        response = super().wait(task, api_key, workspace=workspace)
+        response = super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return ImageGenerationResponse.from_api_response(response)
 
     @classmethod
@@ -408,14 +417,15 @@ async def call(  # type: ignore[override]
             to_merge_incremental_output = True
             kwargs["incremental_output"] = True
 
-        # Pass incremental_to_full flag via headers user-agent
-        if "headers" not in kwargs:
-            kwargs["headers"] = {}
-
+        # Pass incremental_to_full flag via user_agent parameter
         flag = "1" if to_merge_incremental_output else "0"
-        kwargs["headers"]["user-agent"] = f"incremental_to_full/{flag}"
+        existing_ua = kwargs.get("user_agent", "")
+        new_ua = f"incremental_to_full/{flag}"
+        kwargs["user_agent"] = (
+            f"{existing_ua}; {new_ua}".strip() if existing_ua else new_ua
+        )
         if kwargs.get("is_async", False):
-            kwargs["headers"]["X-DashScope-Async"] = "enable"
+            kwargs.setdefault("headers", {})["X-DashScope-Async"] = "enable"
             task = cls.async_task
         else:
             task = cls.sync_task
@@ -487,6 +497,7 @@ async def wait(
         task: Union[str, ImageGenerationResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> DashScopeAPIResponse:
         """Wait for image(s) synthesis task to complete, and return the result.
@@ -496,11 +507,18 @@ async def wait(
                 ImageGenerationResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             DashScopeAPIResponse: The task result.
         """
-        response = await super().wait(task, api_key, workspace=workspace)
+        response = await super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return ImageGenerationResponse.from_api_response(response)
 
     @classmethod

dashscope/aigc/image_synthesis.py

@@ -261,16 +261,16 @@ def _get_input(  # pylint: disable=too-many-branches
             kwargs["headers"] = headers
 
         def __get_i2i_task(task, model) -> str:
-            # 处理task参数：优先使用有效的task值
+            # Handle task parameter: prefer valid task value
             if task is not None and task != "":
                 return task
 
-            # 根据model确定任务类型
+            # Determine task type based on model
             if model is not None and model != "":
                 if "imageedit" in model or "wan2.5-i2i" in model:
                     return "image2image"
 
-            # 默认返回文本到图像任务
+            # Default to text-to-image task
             return ImageSynthesis.task
 
         task = __get_i2i_task(task, model)
@@ -387,6 +387,7 @@ def wait(  # type: ignore[override]
         task: Union[str, ImageSynthesisResponse],
         api_key: str = None,
         workspace: str = None,
+        **kwargs,
     ) -> ImageSynthesisResponse:
         """Wait for image(s) synthesis task to complete, and return the result.
 
@@ -399,7 +400,12 @@ def wait(  # type: ignore[override]
         Returns:
             ImageSynthesisResponse: The task result.
         """
-        response = super().wait(task, api_key, workspace=workspace)
+        response = super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            **kwargs,
+        )
         return ImageSynthesisResponse.from_api_response(response)
 
     @classmethod
@@ -733,6 +739,7 @@ async def wait(
         task: Union[str, ImageSynthesisResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> ImageSynthesisResponse:
         """Wait for image(s) synthesis task to complete, and return the result.
@@ -742,11 +749,18 @@ async def wait(
                 ImageSynthesisResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             ImageSynthesisResponse: The task result.
         """
-        response = await super().wait(task, api_key, workspace=workspace)
+        response = await super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return ImageSynthesisResponse.from_api_response(response)
 
     @classmethod

dashscope/aigc/multimodal_conversation.py

@@ -163,11 +163,13 @@ def call(  # pylint: disable=arguments-renamed,too-many-branches,too-many-statem
             to_merge_incremental_output = True
             kwargs["incremental_output"] = True
 
-        # Pass incremental_to_full flag via headers user-agent
-        if "headers" not in kwargs:
-            kwargs["headers"] = {}
+        # Pass incremental_to_full flag via user_agent parameter
         flag = "1" if to_merge_incremental_output else "0"
-        kwargs["headers"]["user-agent"] = f"incremental_to_full/{flag}"
+        existing_ua = kwargs.get("user_agent", "")
+        new_ua = f"incremental_to_full/{flag}"
+        kwargs["user_agent"] = (
+            f"{existing_ua}; {new_ua}".strip() if existing_ua else new_ua
+        )
 
         response = super().call(
             model=model,

dashscope/aigc/video_synthesis.py

@@ -509,6 +509,7 @@ def wait(  # type: ignore[override]
         task: Union[str, VideoSynthesisResponse],
         api_key: str = None,
         workspace: str = None,
+        **kwargs,
     ) -> VideoSynthesisResponse:
         """Wait for video synthesis task to complete, and return the result.
 
@@ -521,7 +522,12 @@ def wait(  # type: ignore[override]
         Returns:
             VideoSynthesisResponse: The task result.
         """
-        response = super().wait(task, api_key, workspace=workspace)
+        response = super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            **kwargs,
+        )
         return VideoSynthesisResponse.from_api_response(response)
 
     @classmethod
@@ -888,6 +894,7 @@ async def wait(
         task: Union[str, VideoSynthesisResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> VideoSynthesisResponse:
         """Wait for video synthesis task to complete, and return the result.
@@ -897,11 +904,18 @@ async def wait(
                 VideoSynthesisResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             VideoSynthesisResponse: The task result.
         """
-        response = await super().wait(task, api_key, workspace=workspace)
+        response = await super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return VideoSynthesisResponse.from_api_response(response)
 
     @classmethod

dashscope/api_entities/aio_session.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""Shared aiohttp session pool with cached SSL context.
+
+Provides connection reuse across async API calls. Each event loop gets
+its own ClientSession (aiohttp sessions are loop-bound). The SSL context
+is created once and shared across all sessions.
+"""
+import asyncio
+import ssl
+import threading
+import weakref
+from typing import Optional
+
+import aiohttp
+import certifi
+
+_shared_ssl_context: Optional[ssl.SSLContext] = None
+_aio_sessions: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
+_lock = threading.RLock()
+
+
+def get_ssl_context() -> ssl.SSLContext:
+    global _shared_ssl_context
+    with _lock:
+        if _shared_ssl_context is None:
+            _shared_ssl_context = ssl.create_default_context(
+                cafile=certifi.where(),
+            )
+    return _shared_ssl_context
+
+
+async def get_shared_aio_session() -> aiohttp.ClientSession:
+    """Return a shared aiohttp.ClientSession bound to the running event loop.
+
+    The session is lazily created on first use and reused for all
+    subsequent calls on the same event loop. Connection pooling (keep-alive)
+    is handled by the underlying TCPConnector.
+    """
+    loop = asyncio.get_running_loop()
+
+    with _lock:
+        session = _aio_sessions.get(loop)
+        if session is not None and not session.closed:
+            return session
+
+        connector = aiohttp.TCPConnector(ssl=get_ssl_context())
+        session = aiohttp.ClientSession(connector=connector, trust_env=True)
+        _aio_sessions[loop] = session
+    return session
+
+
+async def close_shared_aio_session() -> None:
+    """Close the shared session for the current event loop."""
+    loop = asyncio.get_running_loop()
+    with _lock:
+        session = _aio_sessions.pop(loop, None)
+    if session is not None and not session.closed:
+        await session.close()