From 2447b80ba807912b19550c012b89c5b29b0cfc00 Mon Sep 17 00:00:00 2001
From: Teo <teocns@gmail.com>
Date: Tue, 7 Jan 2025 15:57:35 +0100
Subject: [PATCH] refactor(client): pass client to ClientTelemetry constructor

---
 agentops/client.py           |  2 +-
 agentops/telemetry/client.py | 19 ++++++++++++-------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/agentops/client.py b/agentops/client.py
index 7463e6dd..911b89f2 100644
--- a/agentops/client.py
+++ b/agentops/client.py
@@ -48,7 +48,7 @@ def __init__(self):
         self._config = Configuration()
         self._pre_init_queue = {"agents": []}
         self._host_env = None  # Cache host env data
-        self.telemetry = ClientTelemetry()
+        self.telemetry = ClientTelemetry(self)
 
         self.configure(
             api_key=os.environ.get("AGENTOPS_API_KEY"),
diff --git a/agentops/telemetry/client.py b/agentops/telemetry/client.py
index d6f687ea..36ef722c 100644
--- a/agentops/telemetry/client.py
+++ b/agentops/telemetry/client.py
@@ -15,16 +15,18 @@
 
 if TYPE_CHECKING:
     from agentops.session import Session
+    from agentops.client import Client
 
 
 class ClientTelemetry:
     """Manages telemetry at the agentops.Client level, shared across sessions"""
 
-    def __init__(self):
+    def __init__(self,client: "Client"):
         self._otel_manager: Optional[OTELManager] = None
         self._tracer_provider: Optional[TracerProvider] = None
         self._session_exporters: Dict[UUID, ExportManager] = {}
         self.config: Optional[OTELConfig] = None
+        self.client = client
 
     def initialize(self, config: OTELConfig) -> None:
         """Initialize telemetry components"""
@@ -64,12 +66,15 @@ def initialize(self, config: OTELConfig) -> None:
 
     def get_session_tracer(self, session_id: UUID, jwt: str):
         """Get or create a tracer for a specific session"""
+        if not self.client:
+            raise RuntimeError("Client not initialized")
+        
         # Create session-specific exporter
         exporter = ExportManager(
             session_id=session_id,
-            endpoint=self.config.endpoint,
+            endpoint=self.client._config.endpoint,
             jwt=jwt,
-            api_key=self.config.api_key,
+            api_key=self.client._config.api_key,
             retry_config=self.config.retry_config if self.config else None,
             custom_formatters=self.config.custom_formatters if self.config else None,
         )
@@ -80,11 +85,11 @@ def get_session_tracer(self, session_id: UUID, jwt: str):
         # Add both batch and in-flight processors
         batch_processor = BatchSpanProcessor(
             exporter,
-            max_queue_size=self.config.max_queue_size,
-            schedule_delay_millis=config.max_wait_time,
+            max_queue_size=self.client._config.max_queue_size,
+            schedule_delay_millis=self.client._config.max_wait_time,
             max_export_batch_size=min(
-                max(self.config.max_queue_size // 20, 1),
-                min(self.config.max_queue_size, 32),
+                max(self.client._config.max_queue_size // 20, 1),
+                min(self.client._config.max_queue_size, 32),
             ),
             export_timeout_millis=20000,
         )