1
import time
2
import random
3
from typing import Dict, Any, List
4
from contextlib import contextmanager
5
6
import opik

1
print(f"Opik version: {opik.__version__}")
2
print(f"Initial tracing state: {opik.is_tracing_active()}") # the opposite of the track_disable

1
print(f"Current tracing state: {opik.is_tracing_active()}")
2
3
opik.set_tracing_active(False)
4
print(f"After disabling: {opik.is_tracing_active()}")
5
6
opik.set_tracing_active(True)
7
print(f"After enabling: {opik.is_tracing_active()}")

1
@contextmanager
2
def tracing_enabled(enabled: bool):
3
    """Context manager for temporary tracing control."""
4
    original_state = opik.is_tracing_active()
5
    try:
6
        opik.set_tracing_active(enabled)
7
        yield
8
    finally:
9
        opik.set_tracing_active(original_state)
10
11
print(f"Before context: {opik.is_tracing_active()}")
12
13
with tracing_enabled(False):
14
    print(f"Inside context (disabled): {opik.is_tracing_active()}")
15
    # ny traced functions here won't create spans
16
17
print(f"After context: {opik.is_tracing_active()}")

1
class TracingStrategy:
2
    """Base class for tracing strategies."""
3
4
    def should_trace(self, **kwargs) -> bool:
5
        """Determine if tracing should be enabled for this request."""
6
        raise NotImplementedError
7
8
9
class UserTierStrategy(TracingStrategy):
10
    """Trace only premium users."""
11
12
    def __init__(self, premium_tiers: List[str] = None):
13
        self.premium_tiers = premium_tiers or ["premium", "enterprise"]
14
15
    def should_trace(self, user_tier: str = None, **kwargs) -> bool:
16
        return user_tier in self.premium_tiers
17
18
19
class SamplingStrategy(TracingStrategy):
20
    """Trace a percentage of requests."""
21
22
    def __init__(self, sample_rate: float = 0.1):
23
        self.sample_rate = sample_rate
24
25
    def should_trace(self, **kwargs) -> bool:
26
        return random.random() < self.sample_rate
27
28
29
class DebugModeStrategy(TracingStrategy):
30
    """Trace when in debug mode or for specific users."""
31
32
    def __init__(self, debug_users: List[str] = None):
33
        self.debug_users = debug_users or []
34
        self.debug_mode = False
35
36
    def should_trace(self, user_id: str = None, **kwargs) -> bool:
37
        return self.debug_mode or (user_id in self.debug_users)
38
39
    def enable_debug(self):
40
        self.debug_mode = True
41
42
    def disable_debug(self):
43
        self.debug_mode = False
44
45
strategies = {
46
    "premium_only": UserTierStrategy(),
47
    "10_percent_sample": SamplingStrategy(0.1),
48
    "debug_mode": DebugModeStrategy(["debug_user_1", "debug_user_2"]),
49
}

1
@opik.track(name="user_query")
2
def handle_user_query(query: str, user_id: str, user_tier: str) -> Dict[str, Any]:
3
    """Simulate handling a user query with LLM processing."""
4
    time.sleep(0.01)
5
    
6
    if user_tier in ["premium", "enterprise"]:
7
        response = f"Premium response to: {query}"
8
        tokens_used = random.randint(150, 300)
9
    else:
10
        response = f"Basic response to: {query}"
11
        tokens_used = random.randint(50, 100)
12
    
13
    return {
14
        "query": query,
15
        "response": response,
16
        "user_id": user_id,
17
        "user_tier": user_tier,
18
        "tokens_used": tokens_used,
19
        "processing_time_ms": random.randint(100, 500)
20
    }
21
22
23
24
class SmartRequestHandler:
25
    """Request handler with configurable tracing strategy."""
26
27
    def __init__(self, strategy: TracingStrategy):
28
        self.strategy = strategy
29
30
    def handle_request(
31
        self, query: str, user_id: str, user_tier: str
32
    ) -> Dict[str, Any]:
33
        """Handle request with conditional tracing."""
34
35
        should_trace = self.strategy.should_trace(
36
            user_id=user_id, user_tier=user_tier, query=query
37
        )
38
39
        opik.set_tracing_active(should_trace)
40
41
42
        response_and_metadata = handle_user_query(query, user_id, user_tier)
43
44
        return {
45
            "traced": should_trace,
46
            "result": response_and_metadata,
47
        }
48
49
50
print("=== Testing Premium-Only Strategy ===")
51
handler = SmartRequestHandler(UserTierStrategy())
52
53
requests = [
54
    ("What is AI?", "user1", "free"),
55
    ("Explain ML", "user2", "premium"),
56
    ("How does it work?", "user3", "free"),
57
    ("Advanced question", "user4", "enterprise"),
58
]
59
60
for query, user_id, tier in requests:
61
    result = handler.handle_request(query, user_id, tier)
62
    print(f"User {user_id} ({tier}): Traced = {result['traced']}")