Merge pull request #1117 from roboflow/feat/pass-service-name-received-within-request-path-param

Feat/pass service name received within request path param

Author: grzegorz-roboflow

inference/models/__init__.py

@@ -73,6 +73,7 @@
     pass
 
 from inference.models.resnet import ResNetClassification
+from inference.models.rfdetr import RFDETRObjectDetection
 from inference.models.vit import VitClassification
 from inference.models.yolact import YOLACT
 from inference.models.yolonas import YOLONASObjectDetection
@@ -92,4 +93,3 @@
     YOLOv11ObjectDetection,
 )
 from inference.models.yolov12 import YOLOv12ObjectDetection
-from inference.models.rfdetr import RFDETRObjectDetection

inference/models/rfdetr/__init__.py

@@ -1,3 +1,3 @@
 from inference.models.rfdetr.rfdetr import RFDETRObjectDetection
 
-__all__ = ["RFDETRObjectDetection"]
+__all__ = ["RFDETRObjectDetection"]

inference/usage_tracking/collector.py

@@ -223,7 +223,7 @@ def _calculate_resource_hash(resource_details: Dict[str, Any]) -> str:
         return sha256_hash(json.dumps(resource_details, sort_keys=True))
 
     def _enqueue_payload(self, payload: UsagePayload):
-        logger.debug("Enqueuing usage payload %s", payload)
+        logger.debug("Enqueuing usage payload")
         if not payload:
             return
         with self._queue_lock:
@@ -347,6 +347,8 @@ def _update_usage_payload(
         inference_test_run: bool = False,
         fps: float = 0,
         execution_duration: float = 0,
+        roboflow_service_name: Optional[str] = None,
+        roboflow_internal_secret: Optional[str] = None,
     ):
         source = str(source) if source else ""
         try:
@@ -382,7 +384,8 @@ def _update_usage_payload(
             source_usage["ip_address_hash"] = ip_address_hash
             source_usage["is_gpu_available"] = is_gpu_available
             source_usage["execution_duration"] += execution_duration
-            logger.debug("Updated usage: %s", source_usage)
+            source_usage["roboflow_service_name"] = roboflow_service_name
+            source_usage["roboflow_internal_secret"] = roboflow_internal_secret
 
     def record_usage(
         self,
@@ -395,6 +398,8 @@ def record_usage(
         inference_test_run: bool = False,
         fps: float = 0,
         execution_duration: float = 0,
+        roboflow_service_name: Optional[str] = None,
+        roboflow_internal_secret: Optional[str] = None,
     ):
         if not api_key:
             return
@@ -417,6 +422,8 @@ def record_usage(
             inference_test_run=inference_test_run,
             fps=fps,
             execution_duration=execution_duration,
+            roboflow_service_name=roboflow_service_name,
+            roboflow_internal_secret=roboflow_internal_secret,
         )
 
     async def async_record_usage(
@@ -430,6 +437,8 @@ async def async_record_usage(
         inference_test_run: bool = False,
         fps: float = 0,
         execution_duration: float = 0,
+        roboflow_service_name: Optional[str] = None,
+        roboflow_internal_secret: Optional[str] = None,
     ):
         if self._async_lock:
             async with self._async_lock:
@@ -443,6 +452,8 @@ async def async_record_usage(
                     inference_test_run=inference_test_run,
                     fps=fps,
                     execution_duration=execution_duration,
+                    roboflow_internal_service_name=roboflow_service_name,
+                    roboflow_internal_secret=roboflow_internal_secret,
                 )
         else:
             self.record_usage(
@@ -455,6 +466,8 @@ async def async_record_usage(
                 inference_test_run=inference_test_run,
                 fps=fps,
                 execution_duration=execution_duration,
+                roboflow_internal_service_name=roboflow_service_name,
+                roboflow_internal_secret=roboflow_internal_secret,
             )
 
     def _usage_collector(self):
@@ -515,7 +528,6 @@ def _offload_to_api(self, payloads: List[APIKeyUsage]):
                         self._plan_details._is_enterprise_col_name
                     ]
 
-            logger.debug("Sending usage payload %s", payload)
             api_keys_hashes_failed = send_usage_payload(
                 payload=payload,
                 api_usage_endpoint_url=self._settings.api_usage_endpoint_url,
@@ -644,6 +656,9 @@ def _extract_usage_params_from_func_kwargs(
             else:
                 logger.debug("Could not obtain API key from func kwargs")
 
+        roboflow_service_name = func_kwargs.get("source_info")
+        roboflow_internal_secret = func_kwargs.get("service_secret")
+
         return {
             "source": source,
             "api_key": usage_api_key,
@@ -653,6 +668,8 @@ def _extract_usage_params_from_func_kwargs(
             "inference_test_run": usage_inference_test_run,
             "fps": usage_fps,
             "execution_duration": execution_duration,
+            "roboflow_service_name": roboflow_service_name,
+            "roboflow_internal_secret": roboflow_internal_secret,
         }
 
     def __call__(

inference/usage_tracking/decorator_helpers.py

@@ -125,14 +125,15 @@ def get_request_resource_id_from_kwargs(func_kwargs: Dict[str, Any]) -> Optional
 def get_request_resource_details_from_kwargs(
     func_kwargs: Dict[str, Any],
 ) -> Dict[str, Any]:
+    resource_details = {}
     if "workflow_request" in func_kwargs:
         workflow_request = func_kwargs["workflow_request"]
         if hasattr(workflow_request, "specification") and isinstance(
             workflow_request.specification, dict
         ):
-            return {
-                "steps": get_resource_details_from_workflow_json(
-                    workflow_json=workflow_request.specification,
-                )
-            }
-    return {}
+            resource_details["steps"] = get_resource_details_from_workflow_json(
+                workflow_json=workflow_request.specification,
+            )
+    if "countinference" in func_kwargs:
+        resource_details["billable"] = func_kwargs["countinference"]
+    return resource_details