kalibr 1.1.3a0__py3-none-any.whl → 1.4.0__py3-none-any.whl

kalibr/instrumentation/registry.py

@@ -3,18 +3,25 @@ Instrumentation Registry
 
 Handles auto-discovery and registration of LLM SDK instrumentations.
 Provides a central place to manage which SDKs are instrumented.
+
+Thread-safe registry using locks to protect shared state.
 """
 
 import os
+import threading
 from typing import Dict, List, Set
 
 # Track which providers have been instrumented
 _instrumented_providers: Set[str] = set()
+# Lock to protect concurrent access to the registry
+_registry_lock = threading.Lock()
 
 
 def auto_instrument(providers: List[str] = None) -> Dict[str, bool]:
     """
     Auto-discover and instrument LLM SDKs
+    
+    Thread-safe: Uses internal lock to protect registry state.
 
     Args:
         providers: List of provider names to instrument.
@@ -35,10 +42,11 @@ def auto_instrument(providers: List[str] = None) -> Dict[str, bool]:
     for provider in providers:
         provider_lower = provider.lower()
 
-        # Skip if already instrumented
-        if provider_lower in _instrumented_providers:
-            results[provider_lower] = True
-            continue
+        # Check if already instrumented (thread-safe read)
+        with _registry_lock:
+            if provider_lower in _instrumented_providers:
+                results[provider_lower] = True
+                continue
 
         try:
             if provider_lower == "openai":
@@ -47,7 +55,8 @@ def auto_instrument(providers: List[str] = None) -> Dict[str, bool]:
                 success = openai_instr.instrument()
                 results[provider_lower] = success
                 if success:
-                    _instrumented_providers.add(provider_lower)
+                    with _registry_lock:
+                        _instrumented_providers.add(provider_lower)
                     print(f"✅ Instrumented OpenAI SDK")
 
             elif provider_lower == "anthropic":
@@ -56,7 +65,8 @@ def auto_instrument(providers: List[str] = None) -> Dict[str, bool]:
                 success = anthropic_instr.instrument()
                 results[provider_lower] = success
                 if success:
-                    _instrumented_providers.add(provider_lower)
+                    with _registry_lock:
+                        _instrumented_providers.add(provider_lower)
                     print(f"✅ Instrumented Anthropic SDK")
 
             elif provider_lower == "google":
@@ -65,7 +75,8 @@ def auto_instrument(providers: List[str] = None) -> Dict[str, bool]:
                 success = google_instr.instrument()
                 results[provider_lower] = success
                 if success:
-                    _instrumented_providers.add(provider_lower)
+                    with _registry_lock:
+                        _instrumented_providers.add(provider_lower)
                     print(f"✅ Instrumented Google Generative AI SDK")
 
             else:
@@ -85,6 +96,8 @@ def auto_instrument(providers: List[str] = None) -> Dict[str, bool]:
 def uninstrument_all() -> Dict[str, bool]:
     """
     Remove instrumentation from all previously instrumented SDKs
+    
+    Thread-safe: Uses internal lock to protect registry state.
 
     Returns:
         Dictionary mapping provider names to uninstrumentation success status
@@ -92,7 +105,10 @@ def uninstrument_all() -> Dict[str, bool]:
     global _instrumented_providers
 
     results = {}
-    providers_to_uninstrument = list(_instrumented_providers)
+    
+    # Get snapshot of providers to uninstrument (thread-safe)
+    with _registry_lock:
+        providers_to_uninstrument = list(_instrumented_providers)
 
     for provider in providers_to_uninstrument:
         try:
@@ -102,7 +118,8 @@ def uninstrument_all() -> Dict[str, bool]:
                 success = openai_instr.uninstrument()
                 results[provider] = success
                 if success:
-                    _instrumented_providers.discard(provider)
+                    with _registry_lock:
+                        _instrumented_providers.discard(provider)
                     print(f"✅ Uninstrumented OpenAI SDK")
 
             elif provider == "anthropic":
@@ -111,7 +128,8 @@ def uninstrument_all() -> Dict[str, bool]:
                 success = anthropic_instr.uninstrument()
                 results[provider] = success
                 if success:
-                    _instrumented_providers.discard(provider)
+                    with _registry_lock:
+                        _instrumented_providers.discard(provider)
                     print(f"✅ Uninstrumented Anthropic SDK")
 
             elif provider == "google":
@@ -120,7 +138,8 @@ def uninstrument_all() -> Dict[str, bool]:
                 success = google_instr.uninstrument()
                 results[provider] = success
                 if success:
-                    _instrumented_providers.discard(provider)
+                    with _registry_lock:
+                        _instrumented_providers.discard(provider)
                     print(f"✅ Uninstrumented Google Generative AI SDK")
 
         except Exception as e:
@@ -133,16 +152,21 @@ def uninstrument_all() -> Dict[str, bool]:
 def get_instrumented_providers() -> List[str]:
     """
     Get list of currently instrumented providers
+    
+    Thread-safe: Returns a snapshot of the current state.
 
     Returns:
         List of provider names that are currently instrumented
     """
-    return list(_instrumented_providers)
+    with _registry_lock:
+        return list(_instrumented_providers)
 
 
 def is_instrumented(provider: str) -> bool:
     """
     Check if a specific provider is instrumented
+    
+    Thread-safe: Protected by internal lock.
 
     Args:
         provider: Provider name to check
@@ -150,4 +174,5 @@ def is_instrumented(provider: str) -> bool:
     Returns:
         True if provider is instrumented, False otherwise
     """
-    return provider.lower() in _instrumented_providers
+    with _registry_lock:
+        return provider.lower() in _instrumented_providers