dify-oapi2 0.4.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

dify_oapi/core/http/transport/async_transport.py

@@ -16,6 +16,7 @@ from dify_oapi.core.model.request_option import RequestOption
 from dify_oapi.core.type import T
 
 from ._misc import _build_header, _build_url, _get_sleep_time, _merge_dicts, _unmarshaller
+from .connection_pool import connection_pool
 
 
 def _format_log_details(
@@ -33,6 +34,12 @@ def _format_log_details(
     return ", ".join(details)
 
 
+def _update_response_mode(body_data: dict | None, stream: bool) -> None:
+    """Update response_mode field in request body based on stream parameter"""
+    if body_data and "response_mode" in body_data:
+        body_data["response_mode"] = "streaming" if stream else "blocking"
+
+
 async def _handle_async_stream_error(response: httpx.Response) -> bytes:
     """Handle async streaming response errors"""
     try:
@@ -58,6 +65,16 @@ async def _async_stream_generator(
     method_name = http_method.name
     body_data = _merge_dicts(json_, files, data)
 
+    # Use connection pool for async streaming requests
+    client = connection_pool.get_async_client(
+        conf.domain or "",
+        conf.timeout,
+        getattr(conf, "max_keepalive_connections", 20),
+        getattr(conf, "max_connections", 100),
+        getattr(conf, "keepalive_expiry", 30.0),
+        getattr(conf, "verify_ssl", True),
+    )
+
     for retry in range(conf.max_retry_count + 1):
         if retry > 0:
             sleep_time = _get_sleep_time(retry)
@@ -65,19 +82,16 @@ async def _async_stream_generator(
             await asyncio.sleep(sleep_time)
 
         try:
-            async with (
-                httpx.AsyncClient() as client,
-                client.stream(
-                    method_name,
-                    url,
-                    headers=headers,
-                    params=tuple(req.queries),
-                    json=json_,
-                    data=data,
-                    files=files,
-                    timeout=conf.timeout,
-                ) as response,
-            ):
+            async with client.stream(
+                method_name,
+                url,
+                headers=headers,
+                params=tuple(req.queries),
+                json=json_,
+                data=data,
+                files=files,
+                timeout=conf.timeout,
+            ) as response:
                 logger.debug(
                     f"{_format_log_details(method_name, url, headers, req.queries, body_data)}, stream response"
                 )
@@ -162,8 +176,10 @@ class ATransport:
             files = req.files
             if req.body is not None:
                 data = json.loads(JSON.marshal(req.body))
+                _update_response_mode(data, stream)
         elif req.body is not None:
             json_ = json.loads(JSON.marshal(req.body))
+            _update_response_mode(json_, stream)
 
         if stream:
             return _async_stream_generator(
@@ -173,45 +189,52 @@ class ATransport:
         method_name = req.http_method.name
         body_data = _merge_dicts(json_, files, data)
 
-        async with httpx.AsyncClient() as client:
-            for retry in range(conf.max_retry_count + 1):
-                if retry > 0:
-                    sleep_time = _get_sleep_time(retry)
-                    logger.info(f"in-request: sleep {sleep_time}s")
-                    await asyncio.sleep(sleep_time)
+        # Use connection pool for async regular requests
+        client = connection_pool.get_async_client(
+            conf.domain or "",
+            conf.timeout,
+            getattr(conf, "max_keepalive_connections", 20),
+            getattr(conf, "max_connections", 100),
+            getattr(conf, "keepalive_expiry", 30.0),
+            getattr(conf, "verify_ssl", True),
+        )
+
+        for retry in range(conf.max_retry_count + 1):
+            if retry > 0:
+                sleep_time = _get_sleep_time(retry)
+                logger.info(f"in-request: sleep {sleep_time}s")
+                await asyncio.sleep(sleep_time)
+
+            try:
+                response = await client.request(
+                    method_name,
+                    url,
+                    headers=headers,
+                    params=tuple(req.queries),
+                    json=json_,
+                    data=data,
+                    files=files,
+                    timeout=conf.timeout,
+                )
+                break
+            except httpx.RequestError as e:
+                err_msg = f"{e.__class__.__name__}: {e!r}"
+                log_details = _format_log_details(method_name, url, headers, req.queries, body_data)
 
-                try:
-                    response = await client.request(
-                        method_name,
-                        url,
-                        headers=headers,
-                        params=tuple(req.queries),
-                        json=json_,
-                        data=data,
-                        files=files,
-                        timeout=conf.timeout,
-                    )
-                    break
-                except httpx.RequestError as e:
-                    err_msg = f"{e.__class__.__name__}: {e!r}"
-                    log_details = _format_log_details(method_name, url, headers, req.queries, body_data)
-
-                    if retry < conf.max_retry_count:
-                        logger.info(
-                            f"in-request: retrying ({retry + 1}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
-                        )
-                        continue
+                if retry < conf.max_retry_count:
                     logger.info(
-                        f"in-request: request failed, retried ({retry}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
+                        f"in-request: retrying ({retry + 1}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
                     )
-                    raise
+                    continue
+                logger.info(
+                    f"in-request: request failed, retried ({retry}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
+                )
+                raise
 
-            logger.debug(
-                f"{_format_log_details(method_name, url, headers, req.queries, body_data)} {response.status_code}"
-            )
+        logger.debug(f"{_format_log_details(method_name, url, headers, req.queries, body_data)} {response.status_code}")
 
-            raw_resp = RawResponse()
-            raw_resp.status_code = response.status_code
-            raw_resp.headers = dict(response.headers)
-            raw_resp.content = response.content
-            return _unmarshaller(raw_resp, unmarshal_as)
+        raw_resp = RawResponse()
+        raw_resp.status_code = response.status_code
+        raw_resp.headers = dict(response.headers)
+        raw_resp.content = response.content
+        return _unmarshaller(raw_resp, unmarshal_as)

dify_oapi/core/http/transport/connection_pool.py

@@ -0,0 +1,131 @@
+"""HTTP connection pool management for efficient TCP connection reuse."""
+
+import threading
+from typing import Optional
+
+import httpx
+
+
+class ConnectionPoolManager:
+    """Manages HTTP connection pools to reduce TCP connection overhead."""
+
+    _instance: Optional["ConnectionPoolManager"] = None
+    _lock = threading.Lock()
+
+    def __new__(cls) -> "ConnectionPoolManager":
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if not hasattr(self, "_initialized"):
+            self._sync_clients: dict[str, httpx.Client] = {}
+            self._async_clients: dict[str, httpx.AsyncClient] = {}
+            self._client_lock = threading.Lock()
+            self._initialized = True
+
+    def get_sync_client(
+        self,
+        domain: str,
+        timeout: Optional[float] = None,
+        max_keepalive: int = 20,
+        max_connections: int = 100,
+        keepalive_expiry: float = 30.0,
+        verify_ssl: bool = True,
+    ) -> httpx.Client:
+        """Get or create a sync HTTP client for the given domain."""
+        client_key = f"{domain}:{timeout}:{max_keepalive}:{max_connections}:{keepalive_expiry}:{verify_ssl}"
+
+        with self._client_lock:
+            if client_key not in self._sync_clients:
+                # Configure connection limits to prevent excessive connections
+                limits = httpx.Limits(
+                    max_keepalive_connections=max_keepalive,
+                    max_connections=max_connections,
+                    keepalive_expiry=keepalive_expiry,
+                )
+
+                self._sync_clients[client_key] = httpx.Client(
+                    timeout=timeout,
+                    limits=limits,
+                    verify=verify_ssl,
+                    # Note: HTTP/2 disabled to avoid h2 dependency requirement
+                )
+
+            return self._sync_clients[client_key]
+
+    def get_async_client(
+        self,
+        domain: str,
+        timeout: Optional[float] = None,
+        max_keepalive: int = 20,
+        max_connections: int = 100,
+        keepalive_expiry: float = 30.0,
+        verify_ssl: bool = True,
+    ) -> httpx.AsyncClient:
+        """Get or create an async HTTP client for the given domain."""
+        client_key = f"{domain}:{timeout}:{max_keepalive}:{max_connections}:{keepalive_expiry}:{verify_ssl}"
+
+        with self._client_lock:
+            if client_key not in self._async_clients:
+                # Configure connection limits to prevent excessive connections
+                limits = httpx.Limits(
+                    max_keepalive_connections=max_keepalive,
+                    max_connections=max_connections,
+                    keepalive_expiry=keepalive_expiry,
+                )
+
+                self._async_clients[client_key] = httpx.AsyncClient(
+                    timeout=timeout,
+                    limits=limits,
+                    verify=verify_ssl,
+                    # Note: HTTP/2 disabled to avoid h2 dependency requirement
+                )
+
+            return self._async_clients[client_key]
+
+    def close_all(self):
+        """Close all HTTP clients and clean up connections."""
+        with self._client_lock:
+            # Close sync clients
+            for client in self._sync_clients.values():
+                try:
+                    client.close()
+                except Exception:
+                    pass
+            self._sync_clients.clear()
+
+            # Close async clients
+            for _client in self._async_clients.values():
+                try:
+                    # Note: This should be called from an async context in practice
+                    # For cleanup purposes, we'll just clear the dict
+                    pass
+                except Exception:
+                    pass
+            self._async_clients.clear()
+
+    async def aclose_all(self):
+        """Async version of close_all for proper async client cleanup."""
+        with self._client_lock:
+            # Close sync clients
+            for sync_client in self._sync_clients.values():
+                try:
+                    sync_client.close()
+                except Exception:
+                    pass
+            self._sync_clients.clear()
+
+            # Close async clients properly
+            for async_client in self._async_clients.values():
+                try:
+                    await async_client.aclose()
+                except Exception:
+                    pass
+            self._async_clients.clear()
+
+
+# Global connection pool manager instance
+connection_pool = ConnectionPoolManager()

dify_oapi/core/http/transport/sync_transport.py

@@ -17,6 +17,7 @@ from dify_oapi.core.model.request_option import RequestOption
 from dify_oapi.core.type import T
 
 from ._misc import _build_header, _build_url, _get_sleep_time, _merge_dicts, _unmarshaller
+from .connection_pool import connection_pool
 
 
 def _format_log_details(
@@ -45,6 +46,12 @@ def _handle_stream_error(response: httpx.Response) -> bytes:
     return f"data: [ERROR] {error_message}\n\n".encode()
 
 
+def _update_response_mode(body_data: dict | None, stream: bool) -> None:
+    """Update response_mode field in request body based on stream parameter"""
+    if body_data and "response_mode" in body_data:
+        body_data["response_mode"] = "streaming" if stream else "blocking"
+
+
 def _stream_generator(
     conf: Config,
     req: BaseRequest,
@@ -59,6 +66,16 @@ def _stream_generator(
     method_name = http_method.name
     body_data = _merge_dicts(json_, files, data)
 
+    # Use connection pool for streaming requests
+    client = connection_pool.get_sync_client(
+        conf.domain or "",
+        conf.timeout,
+        getattr(conf, "max_keepalive_connections", 20),
+        getattr(conf, "max_connections", 100),
+        getattr(conf, "keepalive_expiry", 30.0),
+        getattr(conf, "verify_ssl", True),
+    )
+
     for retry in range(conf.max_retry_count + 1):
         if retry > 0:
             sleep_time = _get_sleep_time(retry)
@@ -66,19 +83,16 @@ def _stream_generator(
             time.sleep(sleep_time)
 
         try:
-            with (
-                httpx.Client() as client,
-                client.stream(
-                    method_name,
-                    url,
-                    headers=headers,
-                    params=tuple(req.queries),
-                    json=json_,
-                    data=data,
-                    files=files,
-                    timeout=conf.timeout,
-                ) as response,
-            ):
+            with client.stream(
+                method_name,
+                url,
+                headers=headers,
+                params=tuple(req.queries),
+                json=json_,
+                data=data,
+                files=files,
+                timeout=conf.timeout,
+            ) as response:
                 logger.debug(
                     f"{_format_log_details(method_name, url, headers, req.queries, body_data)}, stream response"
                 )
@@ -160,8 +174,10 @@ class Transport:
             files = req.files
             if req.body is not None:
                 data = json.loads(JSON.marshal(req.body))
+                _update_response_mode(data, stream)
         elif req.body is not None:
             json_ = json.loads(JSON.marshal(req.body))
+            _update_response_mode(json_, stream)
 
         if stream:
             return _stream_generator(
@@ -176,45 +192,52 @@ class Transport:
         method_name = req.http_method.name
         body_data = _merge_dicts(json_, files, data)
 
-        with httpx.Client() as client:
-            for retry in range(conf.max_retry_count + 1):
-                if retry > 0:
-                    sleep_time = _get_sleep_time(retry)
-                    logger.info(f"in-request: sleep {sleep_time}s")
-                    time.sleep(sleep_time)
+        # Use connection pool for regular requests
+        client = connection_pool.get_sync_client(
+            conf.domain or "",
+            conf.timeout,
+            getattr(conf, "max_keepalive_connections", 20),
+            getattr(conf, "max_connections", 100),
+            getattr(conf, "keepalive_expiry", 30.0),
+            getattr(conf, "verify_ssl", True),
+        )
+
+        for retry in range(conf.max_retry_count + 1):
+            if retry > 0:
+                sleep_time = _get_sleep_time(retry)
+                logger.info(f"in-request: sleep {sleep_time}s")
+                time.sleep(sleep_time)
+
+            try:
+                response = client.request(
+                    method_name,
+                    url,
+                    headers=headers,
+                    params=tuple(req.queries),
+                    json=json_,
+                    data=data,
+                    files=files,
+                    timeout=conf.timeout,
+                )
+                break
+            except httpx.RequestError as e:
+                err_msg = f"{e.__class__.__name__}: {e!r}"
+                log_details = _format_log_details(method_name, url, headers, req.queries, body_data)
 
-                try:
-                    response = client.request(
-                        method_name,
-                        url,
-                        headers=headers,
-                        params=tuple(req.queries),
-                        json=json_,
-                        data=data,
-                        files=files,
-                        timeout=conf.timeout,
-                    )
-                    break
-                except httpx.RequestError as e:
-                    err_msg = f"{e.__class__.__name__}: {e!r}"
-                    log_details = _format_log_details(method_name, url, headers, req.queries, body_data)
-
-                    if retry < conf.max_retry_count:
-                        logger.info(
-                            f"in-request: retrying ({retry + 1}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
-                        )
-                        continue
+                if retry < conf.max_retry_count:
                     logger.info(
-                        f"in-request: request failed, retried ({retry}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
+                        f"in-request: retrying ({retry + 1}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
                     )
-                    raise
+                    continue
+                logger.info(
+                    f"in-request: request failed, retried ({retry}/{conf.max_retry_count}) {log_details}, exp: {err_msg}"
+                )
+                raise
 
-            logger.debug(
-                f"{_format_log_details(method_name, url, headers, req.queries, body_data)} {response.status_code}"
-            )
+        logger.debug(f"{_format_log_details(method_name, url, headers, req.queries, body_data)} {response.status_code}")
 
-            raw_resp = RawResponse()
-            raw_resp.status_code = response.status_code
-            raw_resp.headers = dict(response.headers)
-            raw_resp.content = response.content
-            return _unmarshaller(raw_resp, unmarshal_as)
+        raw_resp = RawResponse()
+        raw_resp.status_code = response.status_code
+        raw_resp.headers = dict(response.headers)
+        raw_resp.content = response.content
+        return _unmarshaller(raw_resp, unmarshal_as)

dify_oapi/core/model/config.py

@@ -1,3 +1,5 @@
+import ssl
+
 from dify_oapi.core.enum import LogLevel
 
 
@@ -7,3 +9,11 @@ class Config:
         self.timeout: float | None = None  # Client timeout in seconds, default is no timeout
         self.log_level: LogLevel = LogLevel.WARNING  # Log level, default is WARNING
         self.max_retry_count: int = 3  # Maximum retry count after request failure. Default is 3
+
+        # Connection pool settings
+        self.max_keepalive_connections: int = 20  # Max keepalive connections per pool
+        self.max_connections: int = 100  # Max total connections per pool
+        self.keepalive_expiry: float = 30.0  # Keepalive connection expiry time in seconds
+
+        # SSL settings
+        self.verify_ssl: ssl.SSLContext | str | bool = True  # SSL certificate verification