diff --git a/vision_agent/tools/tool_utils.py b/vision_agent/tools/tool_utils.py
index 2a260c41..185563a4 100644
--- a/vision_agent/tools/tool_utils.py
+++ b/vision_agent/tools/tool_utils.py
@@ -1,7 +1,7 @@
 import inspect
 import logging
 import os
-from typing import Any, Callable, Dict, List, MutableMapping, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, MutableMapping, Optional, Tuple
 
 import pandas as pd
 from IPython.display import display
@@ -34,7 +34,7 @@ def send_inference_request(
     files: Optional[List[Tuple[Any, ...]]] = None,
     v2: bool = False,
     metadata_payload: Optional[Dict[str, Any]] = None,
-) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+) -> Dict[str, Any]:
     # TODO: runtime_tag and function_name should be metadata_payload and now included
     # in the service payload
     try:
diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
index a4eee6ac..08bf0370 100644
--- a/vision_agent/tools/tools.py
+++ b/vision_agent/tools/tools.py
@@ -540,10 +540,10 @@ def countgd_counting(
         "box_threshold": box_threshold,
     }
     metadata_payload = {"function_name": "countgd_counting"}
-    data: List[Dict[str, Any]] = send_inference_request(
+    resp: List[Dict[str, Any]] = send_inference_request(
         payload, "countgd", v2=True, metadata_payload=metadata_payload
-    )
-    return data
+    )  # type: ignore
+    return resp["data"]
 
 
 def countgd_example_based_counting(
@@ -589,10 +589,10 @@ def countgd_example_based_counting(
         "box_threshold": box_threshold,
     }
     metadata_payload = {"function_name": "countgd_example_based_counting"}
-    data: List[Dict[str, Any]] = send_inference_request(
+    resp: List[Dict[str, Any]] = send_inference_request(
         payload, "countgd", v2=True, metadata_payload=metadata_payload
-    )
-    return data
+    )  # type: ignore
+    return resp["data"]
 
 
 def florence2_roberta_vqa(prompt: str, image: np.ndarray) -> str: