fix: api_entities module Review question

dashscope/api_entities/aio_session.py

@@ -7,6 +7,7 @@
 is created once and shared across all sessions.
 """
 import asyncio
+import atexit
 import ssl
 import threading
 import weakref
@@ -47,45 +48,55 @@ async def get_shared_aio_session() -> aiohttp.ClientSession:
         connector = aiohttp.TCPConnector(ssl=get_ssl_context())
         session = aiohttp.ClientSession(connector=connector, trust_env=True)
 
-        # Register finalizer to close session when it's garbage collected.
-        # This prevents resource leaks even if
-        # close_shared_aio_session() is not called.
-        weakref.finalize(session, _sync_close_session, id(session))
-
         _aio_sessions[loop] = session
+        # Register GC-safe finalizer to close session when loop is collected
+        weakref.finalize(session, _sync_close_session, id(session))
     return session
 
 
 def _sync_close_session(session_id: int) -> None:
-    """Synchronous callback when session is garbage collected.
-
-    Note: This runs during GC, so we cannot await. We try to close
-    synchronously if possible, otherwise log a warning.
-    """
+    """GC callback to safely close a session outside async context."""
     try:
         loop = asyncio.get_event_loop()
-        if not loop.is_closed() and not loop.is_running():
-            # Loop exists and is not running, we can use it
-            # Find the session in our cache
-            with _lock:
-                for stored_session in _aio_sessions.values():
-                    if (
-                        id(stored_session) == session_id
-                        and not stored_session.closed
-                    ):
-                        try:
-                            loop.run_until_complete(stored_session.close())
-                        except Exception:
-                            pass
-                        break
+        if loop.is_closed():
+            return
+        if loop.is_running():
+            asyncio.ensure_future(_async_close_by_id(session_id))
+        else:
+            loop.run_until_complete(_async_close_by_id(session_id))
     except RuntimeError:
-        # No event loop available, can't close asynchronously
-        pass
-    except Exception:
-        # Ignore errors during cleanup
         pass
 
 
+async def _async_close_by_id(session_id: int) -> None:
+    """Close session by ID, safe against already-collected sessions."""
+    with _lock:
+        for loop, session in list(_aio_sessions.items()):
+            if id(session) == session_id and not session.closed:
+                await session.close()
+                _aio_sessions.pop(loop, None)
+                return
+
+
+def _atexit_cleanup() -> None:
+    """Cleanup all sessions at interpreter exit."""
+    with _lock:
+        sessions = list(_aio_sessions.items())
+        _aio_sessions.clear()
+
+    for loop, session in sessions:
+        if not session.closed:
+            try:
+                if not loop.is_closed() and not loop.is_running():
+                    loop.run_until_complete(session.close())
+            except Exception:
+                pass
+
+
+# Register atexit handler
+atexit.register(_atexit_cleanup)
+
+
 async def close_shared_aio_session() -> None:
     """Close the shared session for the current event loop."""
     loop = asyncio.get_running_loop()

dashscope/api_entities/aiohttp_request.py

@@ -78,7 +78,7 @@ def __init__(
             self.headers["X-Accel-Buffering"] = "no"
             self.headers["X-DashScope-SSE"] = "enable"
         if self.query:
-            self.url = self.url.replace("api", "api-task")
+            self.url = self.url.replace("/api/", "/api-task/", 1)
             self.url += f"{task_id}"
         if timeout is None:
             self.timeout = DEFAULT_REQUEST_TIMEOUT_SECONDS
@@ -89,7 +89,7 @@ def add_header(self, key, value):
         self.headers[key] = value
 
     def add_headers(self, headers):
-        self.headers = {**self.headers, **headers}
+        self.headers.update(headers)
 
     def call(self):
         response = async_to_sync(self._handle_request())
@@ -142,6 +142,7 @@ async def _handle_response(  # pylint: disable=too-many-branches
         response: aiohttp.ClientResponse,
     ):
         request_id = ""
+        headers = dict(response.headers)
         if (
             response.status == HTTPStatus.OK
             and self.stream
@@ -167,6 +168,7 @@ async def _handle_response(  # pylint: disable=too-many-branches
                         status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                         code="Unknown",
                         message=data,
+                        headers=headers,
                     )
                     continue
                 if is_error:
@@ -175,13 +177,15 @@ async def _handle_response(  # pylint: disable=too-many-branches
                         status_code=status_code,
                         code=msg["code"],
                         message=msg["message"],
+                        headers=headers,
                     )
                 else:
                     yield DashScopeAPIResponse(
                         request_id=request_id,
                         status_code=HTTPStatus.OK,
                         output=output,
                         usage=usage,
+                        headers=headers,
                     )
         elif (
             response.status == HTTPStatus.OK
@@ -201,6 +205,7 @@ async def _handle_response(  # pylint: disable=too-many-branches
                 request_id=request_id,
                 status_code=HTTPStatus.OK,
                 output=output,
+                headers=headers,
             )
         elif response.status == HTTPStatus.OK:
             json_content = await response.json()
@@ -217,6 +222,7 @@ async def _handle_response(  # pylint: disable=too-many-branches
                 status_code=HTTPStatus.OK,
                 output=output,
                 usage=usage,
+                headers=headers,
             )
         else:
             if "application/json" in response.content_type:
@@ -243,6 +249,7 @@ async def _handle_response(  # pylint: disable=too-many-branches
                     status_code=response.status,
                     code=error["code"],
                     message=message,
+                    headers=headers,
                 )
             else:
                 msg = await response.read()
@@ -251,17 +258,16 @@ async def _handle_response(  # pylint: disable=too-many-branches
                     status_code=response.status,
                     code="Unknown",
                     message=msg.decode("utf-8"),
+                    headers=headers,
                 )
 
     # pylint: disable=too-many-branches
     async def _handle_request(self):
         try:
             if self._external_aio_session is not None:
                 session = self._external_aio_session
-                should_close = False
             else:
                 session = await get_shared_aio_session()
-                should_close = False
 
             if self.stream:
                 request_timeout = aiohttp.ClientTimeout(
@@ -271,52 +277,48 @@ async def _handle_request(self):
             else:
                 request_timeout = aiohttp.ClientTimeout(total=self.timeout)
 
-            try:
-                logger.debug("Starting request: %s", self.url)
-                if self.method == HTTPMethod.POST:
-                    is_form, obj = False, {}
-                    if hasattr(self, "data") and self.data is not None:
-                        is_form, obj = self.data.get_aiohttp_payload()
-                    if is_form:
-                        headers = {**self.headers, **obj.headers}
-                        response = await session.post(
-                            url=self.url,
-                            data=obj,
-                            headers=headers,
-                            timeout=request_timeout,
-                        )
-                    else:
-                        body = json.dumps(obj, ensure_ascii=False).encode(
-                            "utf-8",
-                        )
-                        response = await session.request(
-                            "POST",
-                            url=self.url,
-                            data=body,
-                            headers=self.headers,
-                            timeout=request_timeout,
-                        )
-                elif self.method == HTTPMethod.GET:
-                    params = {}
-                    if hasattr(self, "data") and self.data is not None:
-                        params = getattr(self.data, "parameters", {})
-                    response = await session.get(
+            logger.debug("Starting request: %s", self.url)
+            if self.method == HTTPMethod.POST:
+                is_form, obj = False, {}
+                if hasattr(self, "data") and self.data is not None:
+                    is_form, obj = self.data.get_aiohttp_payload()
+                if is_form:
+                    headers = {**self.headers, **obj.headers}
+                    response = await session.post(
                         url=self.url,
-                        params=params,
-                        headers=self.headers,
+                        data=obj,
+                        headers=headers,
                         timeout=request_timeout,
                     )
                 else:
-                    raise UnsupportedHTTPMethod(
-                        f"Unsupported http method: {self.method}",
+                    body = json.dumps(obj, ensure_ascii=False).encode(
+                        "utf-8",
                     )
-                logger.debug("Response returned: %s", self.url)
-                async with response:
-                    async for rsp in self._handle_response(response):
-                        yield rsp
-            finally:
-                if should_close:
-                    await session.close()
-        except Exception as e:
-            logger.debug(e)
-            raise e
+                    response = await session.request(
+                        "POST",
+                        url=self.url,
+                        data=body,
+                        headers=self.headers,
+                        timeout=request_timeout,
+                    )
+            elif self.method == HTTPMethod.GET:
+                params = {}
+                if hasattr(self, "data") and self.data is not None:
+                    params = getattr(self.data, "parameters", {})
+                response = await session.get(
+                    url=self.url,
+                    params=params,
+                    headers=self.headers,
+                    timeout=request_timeout,
+                )
+            else:
+                raise UnsupportedHTTPMethod(
+                    f"Unsupported http method: {self.method}",
+                )
+            logger.debug("Response returned: %s", self.url)
+            async with response:
+                async for rsp in self._handle_response(response):
+                    yield rsp
+        except Exception:
+            logger.exception("Request failed")
+            raise

dashscope/api_entities/api_request_data.py

@@ -2,6 +2,7 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
 import json
+from typing import Optional
 from urllib.parse import urlencode
 
 import aiohttp
@@ -50,19 +51,18 @@ def add_parameters(self, **params):
     def add_resources(self, resources):
         self.resources = resources
 
-    def to_request_object(self) -> str:
-        """Convert data to json, called from http request.
+    def to_request_object(self) -> dict:
+        """Convert data to request dict, called from http request.
         Returns:
-            str: Json string.
+            dict: Request payload.
         """
         self.input = next(self._input_resolver)
         o = {
             k: v
             for k, v in self.__dict__.items()
             if not (
                 k.startswith("_")
-                or k.startswith("task")
-                or k.startswith("function")
+                or k in ("task_group", "task", "function")
                 or v is None
             )
         }
@@ -92,15 +92,6 @@ def get_aiohttp_payload(self):
                 json.dumps(data["parameters"], ensure_ascii=False),
             )
             return True, form()
-            # pylint: disable=unreachable,pointless-string-statement
-            """
-            mp_writer = aiohttp.MultipartWriter('mixed')
-            mp_writer.append('model=%s'%self.model)
-            mp_writer.append('input=%s' % json.dumps(self._input))
-            mp_writer.append('parameters=%s'%json.dumps(self.parameters))
-            mp_writer.append(form())
-            return True, mp_writer
-            """
         else:
             return False, data
 
@@ -143,15 +134,7 @@ def get_websocket_continue_data(self):
         for content in self._input_resolver:
             yield content
 
-    def _to_json_only_data(self) -> str:
-        o = {
-            k: v
-            for k, v in self.__dict__.items()
-            if not (k.startswith("_") or k.startswith("param"))
-        }
-        return json.dumps(o, default=lambda o: o.__dict__)
-
-    def get_batch_binary_data(self) -> bytes:  # type: ignore[return]
+    def get_batch_binary_data(self) -> Optional[bytes]:
         """Get binary data. used in streaming mode none and
            out (input is not streaming), we send data in one package.
            In this case only has one field input.
@@ -161,12 +144,14 @@ def get_batch_binary_data(self) -> bytes:  # type: ignore[return]
         """
         for content in self._input_resolver:
             return content
+        return None
 
-    def _only_parameters(self) -> str:
+    def _only_parameters(self) -> dict:
         temp_input = None
-        if "raw_input" in self.parameters:
-            temp_input = self.parameters.pop("raw_input")
-        obj = {"model": self.model, "parameters": self.parameters, "input": {}}
+        params = dict(self.parameters)
+        if "raw_input" in params:
+            temp_input = params.pop("raw_input")
+        obj = {"model": self.model, "parameters": params, "input": {}}
         if temp_input is not None:
             obj["input"] = temp_input
         if self.task is not None: