From 6cf6af3ad193a9e8e806ddabc3530208d69dc9fd Mon Sep 17 00:00:00 2001
From: Grzegorz Klimaszewski
 <166530809+grzegorz-roboflow@users.noreply.github.com>
Date: Thu, 22 Aug 2024 20:23:49 +0200
Subject: [PATCH 1/2] Fallback API key extraction for UsageCollector

---
 inference/usage_tracking/collector.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/inference/usage_tracking/collector.py b/inference/usage_tracking/collector.py
index 608a556f2..c7fdc0082 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,18 @@ 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,

From bcf8bc3aa776c875799e653d8066016b0c0d0b41 Mon Sep 17 00:00:00 2001
From: Grzegorz Klimaszewski
 <166530809+grzegorz-roboflow@users.noreply.github.com>
Date: Thu, 22 Aug 2024 20:25:34 +0200
Subject: [PATCH 2/2] formatting

---
 inference/usage_tracking/collector.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/inference/usage_tracking/collector.py b/inference/usage_tracking/collector.py
index c7fdc0082..5dd48d146 100644
--- a/inference/usage_tracking/collector.py
+++ b/inference/usage_tracking/collector.py
@@ -599,7 +599,12 @@ def _extract_usage_params_from_func_kwargs(
                 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"]:
+            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")