fix import name issue

Author: Hao Ni (CSI Interfusion Co Ltd)

agentlightning/instrumentation/agentops.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Callable, cast
+from typing import Any, Callable
 
 import requests
 from agentops.client.api import V3Client, V4Client
@@ -38,7 +38,7 @@ def enable_agentops_service(enabled: bool = True) -> None:
     """
     global _agentops_service_enabled
     _agentops_service_enabled = enabled
-    logger.info(f"Switch set to {value} for exporters and clients.")
+    logger.info(f"Switch set to {enabled} for exporters and clients.")
 
 
 def _patch_exporters():
@@ -265,8 +265,7 @@ class SwitchableV3Client(V3Client):
 
     def fetch_auth_token(self, *args: Any, **kwargs: Any) -> AuthTokenResponse:
         if _agentops_service_enabled:
-            resp = super().post(*args, **kwargs)
-            return cast(AuthTokenResponse, resp)
+            return super().fetch_auth_token(*args, **kwargs)
         else:
             logger.debug("SwitchableV3Client is switched off, skipping fetch_auth_token request.")
             return AuthTokenResponse(token="dummy", project_id="dummy")

tests/instrumentation/test_agentops.py

@@ -9,7 +9,7 @@
     SwitchableAuthenticatedOTLPExporter,
     SwitchableOTLPMetricExporter,
     SwitchableOTLPSpanExporter,
-    set_switch,
+    enable_agentops_service,
 )
 
 
@@ -19,12 +19,12 @@ def test_switchable_authenticated_exporter():
     with patch.object(
         switchable_authenticated_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
     ) as mock_export:
-        set_switch(True)
+        enable_agentops_service()
         result = switchable_authenticated_exporter.export([])
         assert result == SpanExportResult.SUCCESS
         mock_export.assert_called_once()
 
-        set_switch(False)
+        enable_agentops_service(False)
         result = switchable_authenticated_exporter.export([])
         assert result == SpanExportResult.SUCCESS
         assert mock_export.call_count == 1
@@ -36,12 +36,12 @@ def test_switchable_otlp_metric_exporter():
     with patch.object(
         switchable_otlp_metric_exporter.__class__.__bases__[0], "export", return_value=MetricExportResult.SUCCESS
     ) as mock_export:
-        set_switch(True)
+        enable_agentops_service()
         result = switchable_otlp_metric_exporter.export(metrics_data=MagicMock())
         assert result == MetricExportResult.SUCCESS
         mock_export.assert_called_once()
 
-        set_switch(False)
+        enable_agentops_service(False)
         result = switchable_otlp_metric_exporter.export(metrics_data=MagicMock())
         assert result == MetricExportResult.SUCCESS
         assert mock_export.call_count == 1
@@ -53,12 +53,12 @@ def test_switchable_otlp_span_exporter():
     with patch.object(
         switchable_otlp_span_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
     ) as mock_export:
-        set_switch(True)
+        enable_agentops_service()
         result = switchable_otlp_span_exporter.export([])
         assert result == SpanExportResult.SUCCESS
         mock_export.assert_called_once()
 
-        set_switch(False)
+        enable_agentops_service(False)
         result = switchable_otlp_span_exporter.export([])
         assert result == SpanExportResult.SUCCESS
         assert mock_export.call_count == 1