23
23
TRUE_LIST = ["true" , "on" , "1" ]
24
24
25
25
_publish = utils .get_env_variable ("OPENLAYER_DISABLE_PUBLISH" ) not in TRUE_LIST
26
- _verify_ssl = (
27
- utils .get_env_variable ("OPENLAYER_VERIFY_SSL" ) or "true"
28
- ).lower () in TRUE_LIST
26
+ _verify_ssl = (utils .get_env_variable ("OPENLAYER_VERIFY_SSL" ) or "true" ).lower () in TRUE_LIST
29
27
_client = None
30
28
29
+ # Configuration variables for programmatic setup
30
+ _configured_api_key : Optional [str ] = None
31
+ _configured_pipeline_id : Optional [str ] = None
32
+ _configured_base_url : Optional [str ] = None
33
+
34
+
35
+ def configure (
36
+ api_key : Optional [str ] = None ,
37
+ inference_pipeline_id : Optional [str ] = None ,
38
+ base_url : Optional [str ] = None ,
39
+ ) -> None :
40
+ """Configure the Openlayer tracer with custom settings.
41
+
42
+ This function allows you to programmatically set the API key, inference pipeline ID,
43
+ and base URL for the Openlayer client, instead of relying on environment variables.
44
+
45
+ Args:
46
+ api_key: The Openlayer API key. If not provided, falls back to OPENLAYER_API_KEY environment variable.
47
+ inference_pipeline_id: The default inference pipeline ID to use for tracing.
48
+ If not provided, falls back to OPENLAYER_INFERENCE_PIPELINE_ID environment variable.
49
+ base_url: The base URL for the Openlayer API. If not provided, falls back to
50
+ OPENLAYER_BASE_URL environment variable or the default.
51
+
52
+ Examples:
53
+ >>> import openlayer.lib.tracing.tracer as tracer
54
+ >>> # Configure with API key and pipeline ID
55
+ >>> tracer.configure(api_key="your_api_key_here", inference_pipeline_id="your_pipeline_id_here")
56
+ >>> # Now use the decorators normally
57
+ >>> @tracer.trace()
58
+ >>> def my_function():
59
+ ... return "result"
60
+ """
61
+ global _configured_api_key , _configured_pipeline_id , _configured_base_url , _client
62
+
63
+ _configured_api_key = api_key
64
+ _configured_pipeline_id = inference_pipeline_id
65
+ _configured_base_url = base_url
66
+
67
+ # Reset the client so it gets recreated with new configuration
68
+ _client = None
69
+
31
70
32
71
def _get_client () -> Optional [Openlayer ]:
33
72
"""Get or create the Openlayer client with lazy initialization."""
@@ -37,13 +76,24 @@ def _get_client() -> Optional[Openlayer]:
37
76
38
77
if _client is None :
39
78
# Lazy initialization - create client when first needed
79
+ client_kwargs = {}
80
+
81
+ # Use configured API key if available, otherwise fall back to environment variable
82
+ if _configured_api_key is not None :
83
+ client_kwargs ["api_key" ] = _configured_api_key
84
+
85
+ # Use configured base URL if available, otherwise fall back to environment variable
86
+ if _configured_base_url is not None :
87
+ client_kwargs ["base_url" ] = _configured_base_url
88
+
40
89
if _verify_ssl :
41
- _client = Openlayer ()
90
+ _client = Openlayer (** client_kwargs )
42
91
else :
43
92
_client = Openlayer (
44
93
http_client = DefaultHttpxClient (
45
94
verify = False ,
46
95
),
96
+ ** client_kwargs ,
47
97
)
48
98
return _client
49
99
@@ -163,9 +213,7 @@ def wrapper(*func_args, **func_kwargs):
163
213
if step_kwargs .get ("name" ) is None :
164
214
step_kwargs ["name" ] = func .__name__
165
215
166
- with create_step (
167
- * step_args , inference_pipeline_id = inference_pipeline_id , ** step_kwargs
168
- ) as step :
216
+ with create_step (* step_args , inference_pipeline_id = inference_pipeline_id , ** step_kwargs ) as step :
169
217
output = exception = None
170
218
try :
171
219
output = func (* func_args , ** func_kwargs )
@@ -252,14 +300,12 @@ async def __anext__(self):
252
300
# Initialize tracing on first iteration only
253
301
if not self ._trace_initialized :
254
302
self ._original_gen = func (* func_args , ** func_kwargs )
255
- self ._step , self ._is_root_step , self ._token = (
256
- _create_and_initialize_step (
257
- step_name = step_name ,
258
- step_type = enums .StepType .USER_CALL ,
259
- inputs = None ,
260
- output = None ,
261
- metadata = None ,
262
- )
303
+ self ._step , self ._is_root_step , self ._token = _create_and_initialize_step (
304
+ step_name = step_name ,
305
+ step_type = enums .StepType .USER_CALL ,
306
+ inputs = None ,
307
+ output = None ,
308
+ metadata = None ,
263
309
)
264
310
self ._inputs = _extract_function_inputs (
265
311
func_signature = func_signature ,
@@ -453,9 +499,7 @@ def _create_and_initialize_step(
453
499
return new_step , is_root_step , token
454
500
455
501
456
- def _handle_trace_completion (
457
- is_root_step : bool , step_name : str , inference_pipeline_id : Optional [str ] = None
458
- ) -> None :
502
+ def _handle_trace_completion (is_root_step : bool , step_name : str , inference_pipeline_id : Optional [str ] = None ) -> None :
459
503
"""Handle trace completion and data streaming."""
460
504
if is_root_step :
461
505
logger .debug ("Ending the trace..." )
@@ -486,8 +530,12 @@ def _handle_trace_completion(
486
530
)
487
531
if _publish :
488
532
try :
489
- inference_pipeline_id = inference_pipeline_id or utils .get_env_variable (
490
- "OPENLAYER_INFERENCE_PIPELINE_ID"
533
+ # Use provided pipeline_id, or fall back to configured default,
534
+ # or finally to environment variable
535
+ inference_pipeline_id = (
536
+ inference_pipeline_id
537
+ or _configured_pipeline_id
538
+ or utils .get_env_variable ("OPENLAYER_INFERENCE_PIPELINE_ID" )
491
539
)
492
540
client = _get_client ()
493
541
if client :
@@ -503,8 +551,7 @@ def _handle_trace_completion(
503
551
except Exception as err : # pylint: disable=broad-except
504
552
logger .error (traceback .format_exc ())
505
553
logger .error (
506
- "Could not stream data to Openlayer (pipeline_id: %s, base_url: %s)"
507
- " Error: %s" ,
554
+ "Could not stream data to Openlayer (pipeline_id: %s, base_url: %s) Error: %s" ,
508
555
inference_pipeline_id ,
509
556
client .base_url ,
510
557
err ,
@@ -536,9 +583,7 @@ def _process_wrapper_inputs_and_outputs(
536
583
func_kwargs = func_kwargs ,
537
584
context_kwarg = context_kwarg ,
538
585
)
539
- _finalize_step_logging (
540
- step = step , inputs = inputs , output = output , start_time = step .start_time
541
- )
586
+ _finalize_step_logging (step = step , inputs = inputs , output = output , start_time = step .start_time )
542
587
543
588
544
589
def _extract_function_inputs (
@@ -606,9 +651,7 @@ def _finalize_async_generator_step(
606
651
) -> None :
607
652
"""Finalize async generator step - called when generator is consumed."""
608
653
_current_step .reset (token )
609
- _finalize_step_logging (
610
- step = step , inputs = inputs , output = output , start_time = step .start_time
611
- )
654
+ _finalize_step_logging (step = step , inputs = inputs , output = output , start_time = step .start_time )
612
655
_handle_trace_completion (
613
656
is_root_step = is_root_step ,
614
657
step_name = step_name ,
0 commit comments