Add user parameter replacement functionality in OIDC API Proxy

Author: fujita-h

endpoints/helpers/endpoint.py

@@ -3,6 +3,7 @@
 
 import httpx
 import werkzeug
+import werkzeug.datastructures
 
 
 def proxy_response(
@@ -128,3 +129,26 @@ def check_llm_streaming_request(request: werkzeug.Request) -> Tuple[bool, bool]:
                     is_stream = True
 
     return is_llm, is_stream
+
+
+def replace_user_params(
+    user: str,
+    args: werkzeug.datastructures.MultiDict[str, str],
+    json: Any | None,
+    data: werkzeug.datastructures.ImmutableMultiDict[str, str],
+) -> Tuple[
+    werkzeug.datastructures.MultiDict[str, str],
+    Any | None,
+    werkzeug.datastructures.ImmutableMultiDict[str, str],
+]:
+    if not user:
+        return args, json, data
+
+    if args:
+        args["user"] = user
+    if json is not None:
+        json["user"] = user
+    if data:
+        data["user"] = user
+
+    return args, json, data

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
+from endpoints.helpers.endpoint import OidcApiProxyErrorResponse, proxy_response, replace_user_params
 from endpoints.helpers.oidc import OpenIDConnectDiscoveryProvider
 
 
@@ -16,6 +16,7 @@ def _invoke(self, r: Request, values: Mapping, settings: Mapping) -> Response:
         oidc_scope = str(settings.get("oidc_scope", ""))
         dify_api_url = str(settings.get("dify_api_url", ""))
         dify_api_key = str(settings.get("dify_api_key", ""))
+        dify_replace_user_param_claim = str(settings.get("dify_replace_user_param_claim", ""))
 
         # prepare dify api url by removing trailing slash
         if dify_api_url.endswith("/"):
@@ -43,7 +44,7 @@ def _invoke(self, r: Request, values: Mapping, settings: Mapping) -> Response:
         # Verify access token
         try:
             oidc_provider = OpenIDConnectDiscoveryProvider(self.session, oidc_issuer, oidc_audience, oidc_scope)
-            _ = oidc_provider.verify_access_token(access_token)
+            oidc_claims = oidc_provider.verify_access_token(access_token)
         except Exception as e:
             return OidcApiProxyErrorResponse(str(e), 401)
 
@@ -61,8 +62,14 @@ def _invoke(self, r: Request, values: Mapping, settings: Mapping) -> Response:
             **({"Content-Type": r.headers["Content-Type"]} if r.headers.get("Content-Type") else {}),
         }
 
-        # prepare json if request is json
+        # prepare params, json and data
+        params = r.args
         json = r.get_json() if r.is_json else None
+        data = r.form
+
+        # replace user params
+        user = str(oidc_claims.get(dify_replace_user_param_claim, ""))
+        params, json, data = replace_user_params(user, params, json, data)
 
         # prepare files if request has files
         files = [
@@ -72,7 +79,7 @@ def _invoke(self, r: Request, values: Mapping, settings: Mapping) -> Response:
         # Forward request to Dify API with Syncronous HTTP Client
         try:
             return proxy_response(
-                request=r, method=r.method, url=url, headers=headers, params=r.args, json=json, data=r.form, files=files
+                request=r, method=r.method, url=url, headers=headers, params=params, json=json, data=data, files=files
             )
         except Exception as e:
             print(str(e))

group/oidc_api_proxy.yaml

@@ -34,6 +34,13 @@ settings:
       en_US: Dify - API Key
     placeholder:
       en_US: Please input Dify API Key
+  - name: dify_replace_user_param_claim
+    type: text-input
+    required: false
+    label:
+      en_US: Dify - Replace user parameters using OIDC claims
+    placeholder:
+      en_US: Please input the claim name to use for replacement (if required).
 endpoints:
   - endpoints/definitions/get/conversations.yaml
   - endpoints/definitions/get/info.yaml