Implement request type checks for LLM streaming in proxy response handling

Author: fujita-h

endpoints/helpers/endpoint.py

@@ -1,11 +1,12 @@
 import json
-from typing import Any
+from typing import Any, Tuple
 
 import httpx
 import werkzeug
 
 
 def proxy_response(
+    request: werkzeug.Request,
     method: str,
     url: str,
     headers: httpx._types.HeaderTypes | None,
@@ -14,7 +15,54 @@ def proxy_response(
     data: httpx._types.RequestData | None,
     files: httpx._types.RequestFiles | None,
 ) -> werkzeug.Response:
-    with httpx.Client() as client:
+    is_llm, is_stream = check_llm_streaming_request(request)
+
+    if is_llm and is_stream:
+        return proxy_stream_response(
+            method=method,
+            url=url,
+            headers=headers,
+            params=params,
+            json=json,
+            data=data,
+            files=files,
+            timeout=httpx.Timeout(None, read=300, write=10),
+        )
+    elif is_llm:
+        return proxy_blocking_response(
+            method=method,
+            url=url,
+            headers=headers,
+            params=params,
+            json=json,
+            data=data,
+            files=files,
+            timeout=httpx.Timeout(None, read=300, write=10),
+        )
+    else:
+        return proxy_blocking_response(
+            method=method,
+            url=url,
+            headers=headers,
+            params=params,
+            json=json,
+            data=data,
+            files=files,
+            timeout=httpx.Timeout(None, read=10, write=10),
+        )
+
+
+def proxy_blocking_response(
+    method: str,
+    url: str,
+    headers: httpx._types.HeaderTypes | None,
+    params: httpx._types.QueryParamTypes | None,
+    json: Any | None,
+    data: httpx._types.RequestData | None,
+    files: httpx._types.RequestFiles | None,
+    timeout: httpx._types.TimeoutTypes | None,
+) -> werkzeug.Response:
+    with httpx.Client(timeout=timeout) as client:
         response = client.request(
             method=method, url=url, headers=headers, params=params, json=json, data=data, files=files
         )
@@ -34,8 +82,11 @@ def proxy_stream_response(
     json: Any | None,
     data: httpx._types.RequestData | None,
     files: httpx._types.RequestFiles | None,
+    timeout: httpx._types.TimeoutTypes | None,
 ) -> werkzeug.Response:
-    stream_context = httpx.stream(method=method, url=url, headers=headers, json=json, data=data, files=files)
+    stream_context = httpx.stream(
+        timeout=timeout, method=method, url=url, headers=headers, json=json, data=data, files=files
+    )
 
     # Manually enter the context manager to get the response
     stream_response = stream_context.__enter__()
@@ -62,3 +113,18 @@ def OidcApiProxyErrorResponse(message: str, error_code: int = 500) -> werkzeug.R
         status=error_code,
         content_type="application/json",
     )
+
+
+def check_llm_streaming_request(request: werkzeug.Request) -> Tuple[bool, bool]:
+    is_llm = False
+    is_stream = False
+
+    if request.method.lower() in ["post"]:
+        if request.path in ["/chat-messages", "/workflows/run"]:
+            is_llm = True
+            if request.is_json:
+                json = request.get_json()
+                if str(json.get("response_mode", "")).lower() == "streaming":
+                    is_stream = True
+
+    return is_llm, is_stream

endpoints/oidc_api_proxy.py

@@ -4,7 +4,7 @@
 from dify_plugin import Endpoint
 from werkzeug import Request, Response
 
-from endpoints.helpers.endpoint import OidcApiProxyErrorResponse, proxy_response, proxy_stream_response
+from endpoints.helpers.endpoint import OidcApiProxyErrorResponse, proxy_response
 from endpoints.helpers.oidc import OpenIDConnectDiscoveryProvider
 
 
@@ -71,8 +71,8 @@ def _invoke(self, r: Request, values: Mapping, settings: Mapping) -> Response:
 
         # Forward request to Dify API with Syncronous HTTP Client
         try:
-            return proxy_stream_response(
-                method=r.method, url=url, headers=headers, params=r.args, json=json, data=r.form, files=files
+            return proxy_response(
+                request=r, method=r.method, url=url, headers=headers, params=r.args, json=json, data=r.form, files=files
             )
         except Exception as e:
             print(str(e))