feat(http): reuse http session

Author: kevinlin09

dashscope/api_entities/api_request_factory.py

@@ -36,6 +36,7 @@ def _get_protocol_params(kwargs):
     base_address = kwargs.pop("base_address", None)
     flattened_output = kwargs.pop("flattened_output", False)
     extra_url_parameters = kwargs.pop("extra_url_parameters", None)
+    session = kwargs.pop("session", None)
 
     # Extract user-agent from headers if present
     user_agent = ""
@@ -58,6 +59,7 @@ def _get_protocol_params(kwargs):
         flattened_output,
         extra_url_parameters,
         user_agent,
+        session,
     )
 
 
@@ -87,6 +89,7 @@ def _build_api_request(  # pylint: disable=too-many-branches
         flattened_output,
         extra_url_parameters,
         user_agent,
+        session,
     ) = _get_protocol_params(kwargs)
     task_id = kwargs.pop("task_id", None)
     enable_encryption = kwargs.pop("enable_encryption", False)
@@ -130,6 +133,7 @@ def _build_api_request(  # pylint: disable=too-many-branches
             flattened_output=flattened_output,
             encryption=encryption,
             user_agent=user_agent,
+            session=session,
         )
     elif api_protocol == ApiProtocol.WEBSOCKET:
         if base_address is not None:

dashscope/api_entities/http_request.py

@@ -4,7 +4,7 @@
 import json
 import ssl
 from http import HTTPStatus
-from typing import Optional
+from typing import Optional, Dict, Union
 
 import aiohttp
 import certifi
@@ -42,6 +42,9 @@ def __init__(
         flattened_output: bool = False,
         encryption: Optional[Encryption] = None,
         user_agent: str = "",
+        session: Optional[
+            Union[requests.Session, aiohttp.ClientSession]
+        ] = None,
     ) -> None:
         """HttpSSERequest, processing http server sent event stream.
 
@@ -54,14 +57,38 @@ def __init__(
                 Defaults to DEFAULT_REQUEST_TIMEOUT_SECONDS.
             user_agent (str, optional): Additional user agent string to
                 append. Defaults to ''.
+            session (Optional[Union[requests.Session,
+                aiohttp.ClientSession]], optional):
+                Custom Session for connection reuse. Can be either
+                requests.Session for sync calls or aiohttp.ClientSession
+                for async calls. Defaults to None.
         """
 
         super().__init__(user_agent=user_agent)
         self.url = url
         self.flattened_output = flattened_output
         self.async_request = async_request
         self.encryption = encryption
-        self.headers = {
+
+        # Auto-detect session type and store accordingly
+        if session is not None:
+            session_type = type(session).__name__
+            session_module = type(session).__module__
+
+            # Check if it's an aiohttp ClientSession
+            if (
+                session_type == "ClientSession" and "aiohttp" in session_module
+            ) or isinstance(session, aiohttp.ClientSession):
+                self._external_session = None
+                self._external_aio_session = session
+            else:
+                # Treat as requests Session
+                self._external_session = session
+                self._external_aio_session = None
+        else:
+            self._external_session = None
+            self._external_aio_session = None
+        self.headers: Dict = {
             "Accept": "application/json",
             "Authorization": f"Bearer {api_key}",
             **self.headers,
@@ -132,18 +159,27 @@ async def aio_call(self):
                 pass
             return result
 
-    async def _handle_aio_request(self):
+    async def _handle_aio_request(self):  # pylint: disable=too-many-branches
         try:
-            connector = aiohttp.TCPConnector(
-                ssl=ssl.create_default_context(
-                    cafile=certifi.where(),
-                ),
-            )
-            async with aiohttp.ClientSession(
-                connector=connector,
-                timeout=aiohttp.ClientTimeout(total=self.timeout),
-                headers=self.headers,
-            ) as session:
+            # Use external aio_session if provided,
+            # otherwise create temporary session
+            if self._external_aio_session is not None:
+                session = self._external_aio_session
+                should_close = False
+            else:
+                connector = aiohttp.TCPConnector(
+                    ssl=ssl.create_default_context(
+                        cafile=certifi.where(),
+                    ),
+                )
+                session = aiohttp.ClientSession(
+                    connector=connector,
+                    timeout=aiohttp.ClientTimeout(total=self.timeout),
+                    headers=self.headers,
+                )
+                should_close = True
+
+            try:
                 logger.debug("Starting request: %s", self.url)
                 if self.method == HTTPMethod.POST:
                     is_form, obj = False, {}
@@ -183,6 +219,10 @@ async def _handle_aio_request(self):
                 async with response:
                     async for rsp in self._handle_aio_response(response):
                         yield rsp
+            finally:
+                # Only close if we created the session
+                if should_close:
+                    await session.close()
         except aiohttp.ClientConnectorError as e:
             logger.error(e)
             raise e
@@ -408,7 +448,16 @@ def _handle_response(  # pylint: disable=too-many-branches
 
     def _handle_request(self):
         try:
-            with requests.Session() as session:
+            # Use external session if provided,
+            # otherwise create temporary session
+            if self._external_session is not None:
+                session = self._external_session
+                should_close = False
+            else:
+                session = requests.Session()
+                should_close = True
+
+            try:
                 if self.method == HTTPMethod.POST:
                     is_form, form, obj = self.data.get_http_payload()
                     if is_form:
@@ -443,6 +492,10 @@ def _handle_request(self):
                     )
                 for rsp in self._handle_response(response):
                     yield rsp
+            finally:
+                # Only close if we created the session
+                if should_close:
+                    session.close()
         except BaseException as e:
             logger.error(e)
             raise e

samples/test_aio_generation.py

@@ -3,6 +3,9 @@
 
 import asyncio
 import os
+import aiohttp
+import ssl
+import certifi
 from dashscope.aigc.generation import AioGeneration
 
 
@@ -439,6 +442,118 @@ async def process_responses(responses, step_name):
         await call_deep_research_model(messages, "第二步：深入研究")
         print("\n 研究完成！")
 
+    @staticmethod
+    async def test_with_custom_session():
+        """示例：使用自定义 ClientSession 进行连接复用"""
+        print("\n=== 使用自定义 ClientSession 示例 ===")
+
+        # 配置 TCPConnector
+        connector = aiohttp.TCPConnector(
+            limit=100,
+            limit_per_host=30,
+            ssl=ssl.create_default_context(cafile=certifi.where()),
+        )
+
+        # 创建自定义 ClientSession
+        async with aiohttp.ClientSession(connector=connector) as session:
+            # 使用同一个 session 进行多次请求
+            for i in range(3):
+                print(f"\n--- 请求 {i+1} ---")
+
+                messages = [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": f"你好"},
+                ]
+
+                response = await AioGeneration.call(
+                    api_key=os.getenv("DASHSCOPE_API_KEY"),
+                    model="qwen-turbo",
+                    messages=messages,
+                    result_format="message",
+                    session=session,  # ← 传入自定义 session
+                )
+
+                print(f"响应: {response.output.choices[0].message.content}")
+
+        print("\n✅ ClientSession 已自动关闭")
+
+    @staticmethod
+    async def test_with_custom_session_streaming():
+        """示例：使用自定义 ClientSession 进行流式输出"""
+        print("\n=== 使用自定义 ClientSession 流式输出示例 ===")
+
+        # 配置连接池
+        connector = aiohttp.TCPConnector(
+            limit=100,
+            ssl=ssl.create_default_context(cafile=certifi.where()),
+        )
+
+        async with aiohttp.ClientSession(connector=connector) as session:
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "请写一首关于秋天的诗"},
+            ]
+
+            print("\n流式输出:")
+            response = await AioGeneration.call(
+                api_key=os.getenv("DASHSCOPE_API_KEY"),
+                model="qwen-turbo",
+                messages=messages,
+                result_format="message",
+                stream=True,
+                incremental_output=True,
+                session=session,  # ← 传入自定义 session
+            )
+
+            async for chunk in response:
+                if chunk.status_code == 200:
+                    print(chunk.output.choices[0].message.content, end='', flush=True)
+
+            print("\n")
+
+        print("✅ ClientSession 已自动关闭")
+
+    @staticmethod
+    async def test_with_custom_session_concurrent():
+        """示例：使用自定义 ClientSession 进行并发请求"""
+        print("\n=== 使用自定义 ClientSession 并发请求示例 ===")
+
+        # 配置连接池
+        connector = aiohttp.TCPConnector(
+            limit=100,
+            ssl=ssl.create_default_context(cafile=certifi.where()),
+        )
+
+        async with aiohttp.ClientSession(connector=connector) as session:
+            # 创建多个并发任务
+            tasks = []
+            topics = ["Python", "JavaScript", "Go"]
+
+            for topic in topics:
+                messages = [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": f"请用一句话介绍：{topic}"},
+                ]
+
+                task = AioGeneration.call(
+                    api_key=os.getenv("DASHSCOPE_API_KEY"),
+                    model="qwen-turbo",
+                    messages=messages,
+                    result_format="message",
+                    session=session,  # ← 所有请求共享同一个 session
+                )
+                tasks.append(task)
+
+            # 并发执行所有任务
+            print("\n开始并发请求...")
+            responses = await asyncio.gather(*tasks)
+
+            # 处理响应
+            for i, response in enumerate(responses):
+                print(f"\n请求 {i+1} 响应: {response.output.choices[0].message.content}")
+
+        print("\n✅ ClientSession 已自动关闭")
+
 
 async def main():
     """Main function to run all async tests."""
@@ -450,6 +565,11 @@ async def main():
     # await TestAioGeneration.test_response_with_search_info()
     # await TestAioGeneration.test_response_with_reasoning_content()
 
+    # 自定义 Session 示例
+    # await TestAioGeneration.test_with_custom_session()
+    # await TestAioGeneration.test_with_custom_session_streaming()
+    # await TestAioGeneration.test_with_custom_session_concurrent()
+
     print("\n所有异步测试用例执行完成！")