diff --git a/inference/usage_tracking/collector.py b/inference/usage_tracking/collector.py index 608a556f2..5dd48d146 100644 --- a/inference/usage_tracking/collector.py +++ b/inference/usage_tracking/collector.py @@ -20,7 +20,6 @@ List, Optional, ParamSpec, - Tuple, TypeVar, ) @@ -593,6 +592,23 @@ def _extract_usage_params_from_func_kwargs( source = image.get("value") elif hasattr(image, "_image_reference"): source = image._image_reference + + if not usage_api_key: + _self = func_kwargs.get("self") + if "api_key" in func_kwargs and func_kwargs["api_key"]: + usage_api_key = func_kwargs["api_key"] + elif _self and hasattr(_self, "api_key") and _self.api_key: + usage_api_key = _self.api_key + elif ( + "kwargs" in func_kwargs + and isinstance(func_kwargs["kwargs"], dict) + and "api_key" in func_kwargs["kwargs"] + and func_kwargs["kwargs"]["api_key"] + ): + usage_api_key = func_kwargs["kwargs"]["api_key"] + else: + logger.debug("Could not obtain API key from func kwargs") + return { "source": source, "api_key": usage_api_key,