Skip to content

Commit c84d268

Browse files
zhansheng.lzsclaude
andcommitted
feat: add external session support for AioHttpRequest and export close_shared_aio_session
- Add optional `session` parameter to AioHttpRequest.__init__ for injecting external aiohttp.ClientSession (aligned with HttpRequest) - Add should_close lifecycle pattern to async request handlers in both AioHttpRequest and HttpRequest (consistent with sync _handle_request) - Export close_shared_aio_session from dashscope package for explicit connection pool cleanup Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 13aee1a commit c84d268

3 files changed

Lines changed: 107 additions & 82 deletions

File tree

dashscope/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
HttpSpeechSynthesizer,
2222
)
2323
from dashscope.audio.tts.speech_synthesizer import SpeechSynthesizer
24+
from dashscope.api_entities.aio_session import close_shared_aio_session
2425
from dashscope.common.api_key import save_api_key
2526
from dashscope.common.env import (
2627
api_key,
@@ -75,6 +76,7 @@
7576
"api_key",
7677
"api_key_file_path",
7778
"save_api_key",
79+
"close_shared_aio_session",
7880
"AioGeneration",
7981
"Conversation",
8082
"Generation",

dashscope/api_entities/aiohttp_request.py

Lines changed: 54 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import json
55
from http import HTTPStatus
6+
from typing import Optional
67

78
import aiohttp
89

@@ -31,6 +32,7 @@ def __init__(
3132
timeout: int = DEFAULT_REQUEST_TIMEOUT_SECONDS,
3233
task_id: str = None,
3334
user_agent: str = "",
35+
session: Optional[aiohttp.ClientSession] = None,
3436
) -> None:
3537
"""HttpSSERequest, processing http server sent event stream.
3638
@@ -45,16 +47,20 @@ def __init__(
4547
Defaults to DEFAULT_REQUEST_TIMEOUT_SECONDS.
4648
user_agent (str, optional): Additional user agent string to
4749
append. Defaults to ''.
50+
session (aiohttp.ClientSession, optional): External aiohttp
51+
session to use instead of the shared session. The caller is
52+
responsible for closing it. Defaults to None.
4853
"""
4954

5055
super().__init__(user_agent=user_agent)
5156
self.url = url
5257
self.async_request = async_request
58+
self._external_aio_session = session
5359
self.headers = {
5460
"Accept": "application/json",
5561
"Authorization": f"Bearer {api_key}",
5662
"Cache-Control": "no-cache",
57-
**self.headers,
63+
**self.headers, # type: ignore[has-type]
5864
}
5965
self.query = query
6066
if self.async_request and self.query is False:
@@ -247,9 +253,16 @@ async def _handle_response( # pylint: disable=too-many-branches
247253
message=msg.decode("utf-8"),
248254
)
249255

256+
# pylint: disable=too-many-branches
250257
async def _handle_request(self):
251258
try:
252-
session = await get_shared_aio_session()
259+
if self._external_aio_session is not None:
260+
session = self._external_aio_session
261+
should_close = False
262+
else:
263+
session = await get_shared_aio_session()
264+
should_close = False
265+
253266
if self.stream:
254267
request_timeout = aiohttp.ClientTimeout(
255268
total=None,
@@ -258,45 +271,49 @@ async def _handle_request(self):
258271
else:
259272
request_timeout = aiohttp.ClientTimeout(total=self.timeout)
260273

261-
logger.debug("Starting request: %s", self.url)
262-
if self.method == HTTPMethod.POST:
263-
is_form, obj = False, {}
264-
if hasattr(self, "data") and self.data is not None:
265-
is_form, obj = self.data.get_aiohttp_payload()
266-
if is_form:
267-
headers = {**self.headers, **obj.headers}
268-
response = await session.post(
274+
try:
275+
logger.debug("Starting request: %s", self.url)
276+
if self.method == HTTPMethod.POST:
277+
is_form, obj = False, {}
278+
if hasattr(self, "data") and self.data is not None:
279+
is_form, obj = self.data.get_aiohttp_payload()
280+
if is_form:
281+
headers = {**self.headers, **obj.headers}
282+
response = await session.post(
283+
url=self.url,
284+
data=obj,
285+
headers=headers,
286+
timeout=request_timeout,
287+
)
288+
else:
289+
response = await session.request(
290+
"POST",
291+
url=self.url,
292+
json=obj,
293+
headers=self.headers,
294+
timeout=request_timeout,
295+
)
296+
elif self.method == HTTPMethod.GET:
297+
params = {}
298+
if hasattr(self, "data") and self.data is not None:
299+
params = getattr(self.data, "parameters", {})
300+
response = await session.get(
269301
url=self.url,
270-
data=obj,
271-
headers=headers,
302+
params=params,
303+
headers=self.headers,
272304
timeout=request_timeout,
273305
)
274306
else:
275-
response = await session.request(
276-
"POST",
277-
url=self.url,
278-
json=obj,
279-
headers=self.headers,
280-
timeout=request_timeout,
307+
raise UnsupportedHTTPMethod(
308+
f"Unsupported http method: {self.method}",
281309
)
282-
elif self.method == HTTPMethod.GET:
283-
params = {}
284-
if hasattr(self, "data") and self.data is not None:
285-
params = getattr(self.data, "parameters", {})
286-
response = await session.get(
287-
url=self.url,
288-
params=params,
289-
headers=self.headers,
290-
timeout=request_timeout,
291-
)
292-
else:
293-
raise UnsupportedHTTPMethod(
294-
f"Unsupported http method: {self.method}",
295-
)
296-
logger.debug("Response returned: %s", self.url)
297-
async with response:
298-
async for rsp in self._handle_response(response):
299-
yield rsp
310+
logger.debug("Response returned: %s", self.url)
311+
async with response:
312+
async for rsp in self._handle_response(response):
313+
yield rsp
314+
finally:
315+
if should_close:
316+
await session.close()
300317
except Exception as e:
301-
logger.error(e)
318+
logger.debug(e)
302319
raise e

dashscope/api_entities/http_request.py

Lines changed: 51 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -166,60 +166,66 @@ async def _handle_aio_request(self): # pylint: disable=too-many-branches
166166
# otherwise use shared session with connection pooling
167167
if self._external_aio_session is not None:
168168
session = self._external_aio_session
169+
should_close = False
169170
else:
170171
session = await get_shared_aio_session()
172+
should_close = False
171173

172-
if self.stream:
173-
request_timeout = aiohttp.ClientTimeout(
174-
total=None,
175-
sock_read=self.timeout,
176-
)
177-
else:
178-
request_timeout = aiohttp.ClientTimeout(total=self.timeout)
179-
180-
logger.debug("Starting request: %s", self.url)
181-
if self.method == HTTPMethod.POST:
182-
is_form, obj = False, {}
183-
if hasattr(self, "data") and self.data is not None:
184-
is_form, obj = self.data.get_aiohttp_payload()
185-
if is_form:
186-
headers = {**self.headers, **obj.headers}
187-
response = await session.post(
188-
url=self.url,
189-
data=obj,
190-
headers=headers,
191-
timeout=request_timeout,
174+
try:
175+
if self.stream:
176+
request_timeout = aiohttp.ClientTimeout(
177+
total=None,
178+
sock_read=self.timeout,
192179
)
193180
else:
194-
response = await session.request(
195-
"POST",
181+
request_timeout = aiohttp.ClientTimeout(total=self.timeout)
182+
183+
logger.debug("Starting request: %s", self.url)
184+
if self.method == HTTPMethod.POST:
185+
is_form, obj = False, {}
186+
if hasattr(self, "data") and self.data is not None:
187+
is_form, obj = self.data.get_aiohttp_payload()
188+
if is_form:
189+
headers = {**self.headers, **obj.headers}
190+
response = await session.post(
191+
url=self.url,
192+
data=obj,
193+
headers=headers,
194+
timeout=request_timeout,
195+
)
196+
else:
197+
response = await session.request(
198+
"POST",
199+
url=self.url,
200+
json=obj,
201+
headers=self.headers,
202+
timeout=request_timeout,
203+
)
204+
elif self.method == HTTPMethod.GET:
205+
params = {}
206+
if hasattr(self, "data") and self.data is not None:
207+
params = getattr(self.data, "parameters", {})
208+
if params:
209+
params = self.__handle_parameters(params)
210+
response = await session.get(
196211
url=self.url,
197-
json=obj,
212+
params=params,
198213
headers=self.headers,
199214
timeout=request_timeout,
200215
)
201-
elif self.method == HTTPMethod.GET:
202-
params = {}
203-
if hasattr(self, "data") and self.data is not None:
204-
params = getattr(self.data, "parameters", {})
205-
if params:
206-
params = self.__handle_parameters(params)
207-
response = await session.get(
208-
url=self.url,
209-
params=params,
210-
headers=self.headers,
211-
timeout=request_timeout,
212-
)
213-
else:
214-
raise UnsupportedHTTPMethod(
215-
f"Unsupported http method: {self.method}",
216-
)
217-
logger.debug("Response returned: %s", self.url)
218-
async with response:
219-
async for rsp in self._handle_aio_response(response):
220-
yield rsp
216+
else:
217+
raise UnsupportedHTTPMethod(
218+
f"Unsupported http method: {self.method}",
219+
)
220+
logger.debug("Response returned: %s", self.url)
221+
async with response:
222+
async for rsp in self._handle_aio_response(response):
223+
yield rsp
224+
finally:
225+
if should_close:
226+
await session.close()
221227
except Exception as e:
222-
logger.error(e)
228+
logger.debug(e)
223229
raise e
224230

225231
@staticmethod
@@ -505,5 +511,5 @@ def _handle_request(self): # pylint: disable=too-many-branches
505511
if should_close:
506512
session.close()
507513
except Exception as e:
508-
logger.error(e)
514+
logger.debug(e)
509515
raise e

0 commit comments

Comments
 (0)