refactor(api): tighten OTel decorator typing (#32163)

api/extensions/otel/decorators/base.py

@@ -1,6 +1,6 @@
 import functools
 from collections.abc import Callable
-from typing import Any, TypeVar, cast
+from typing import ParamSpec, TypeVar, cast
 
 from opentelemetry.trace import get_tracer
 
@@ -8,7 +8,8 @@ from configs import dify_config
 from extensions.otel.decorators.handler import SpanHandler
 from extensions.otel.runtime import is_instrument_flag_enabled
 
-T = TypeVar("T", bound=Callable[..., Any])
+P = ParamSpec("P")
+R = TypeVar("R")
 
 _HANDLER_INSTANCES: dict[type[SpanHandler], SpanHandler] = {SpanHandler: SpanHandler()}
 
@@ -20,7 +21,7 @@ def _get_handler_instance(handler_class: type[SpanHandler]) -> SpanHandler:
     return _HANDLER_INSTANCES[handler_class]
 
 
-def trace_span(handler_class: type[SpanHandler] | None = None) -> Callable[[T], T]:
+def trace_span(handler_class: type[SpanHandler] | None = None) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """
     Decorator that traces a function with an OpenTelemetry span.
 
@@ -30,9 +31,9 @@ def trace_span(handler_class: type[SpanHandler] | None = None) -> Callable[[T],
     :param handler_class: Optional handler class to use for this span. If None, uses the default SpanHandler.
     """
 
-    def decorator(func: T) -> T:
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
         @functools.wraps(func)
-        def wrapper(*args: Any, **kwargs: Any) -> Any:
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
             if not (dify_config.ENABLE_OTEL or is_instrument_flag_enabled()):
                 return func(*args, **kwargs)
 
@@ -46,6 +47,6 @@ def trace_span(handler_class: type[SpanHandler] | None = None) -> Callable[[T],
                 kwargs=kwargs,
             )
 
-        return cast(T, wrapper)
+        return cast(Callable[P, R], wrapper)
 
     return decorator