Skip to content
23 changes: 23 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,29 @@ On timeout, an `APITimeoutError` is thrown.

Note that requests that time out are [retried twice by default](#retries).

#### Tracing Configuration

When using the tracing decorators (`@trace()` and `@trace_async()`), you can configure timeouts and retries for the underlying API calls to Openlayer using the `configure()` function:

```python
from openlayer.lib import trace, configure

# Configure timeout and retries for tracing API calls
configure(
api_key="your_api_key_here",
inference_pipeline_id="your_pipeline_id_here",
timeout=30.0, # 30 seconds timeout for tracing API calls (int or float)
max_retries=5 # Maximum number of retries for failed requests (default: 2)
)

# Now use the decorators normally
@trace()
def my_function():
return "result"
```

These settings apply to all API calls made by the tracer when streaming trace data to Openlayer. If not specified, the defaults are 60 seconds timeout and 2 retries.

## Advanced

### Logging
Expand Down
20 changes: 17 additions & 3 deletions src/openlayer/lib/tracing/tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import traceback
from contextlib import contextmanager
from functools import wraps
from typing import Any, Awaitable, Dict, Generator, List, Optional, Tuple
from typing import Any, Awaitable, Dict, Generator, List, Optional, Tuple, Union

from ..._base_client import DefaultHttpxClient
from ..._client import Openlayer
Expand All @@ -34,24 +34,30 @@
_configured_api_key: Optional[str] = None
_configured_pipeline_id: Optional[str] = None
_configured_base_url: Optional[str] = None
_configured_timeout: Optional[Union[int, float]] = None
_configured_max_retries: Optional[int] = None


def configure(
api_key: Optional[str] = None,
inference_pipeline_id: Optional[str] = None,
base_url: Optional[str] = None,
timeout: Optional[Union[int, float]] = None,
max_retries: Optional[int] = None,
) -> None:
"""Configure the Openlayer tracer with custom settings.

This function allows you to programmatically set the API key, inference pipeline ID,
and base URL for the Openlayer client, instead of relying on environment variables.
base URL, timeout, and retry settings for the Openlayer client, instead of relying on environment variables.

Args:
api_key: The Openlayer API key. If not provided, falls back to OPENLAYER_API_KEY environment variable.
inference_pipeline_id: The default inference pipeline ID to use for tracing.
If not provided, falls back to OPENLAYER_INFERENCE_PIPELINE_ID environment variable.
base_url: The base URL for the Openlayer API. If not provided, falls back to
OPENLAYER_BASE_URL environment variable or the default.
timeout: The timeout for the Openlayer API in seconds (int or float). Defaults to 60 seconds.
max_retries: The maximum number of retries for failed API requests. Defaults to 2.

Examples:
>>> import openlayer.lib.tracing.tracer as tracer
Expand All @@ -62,11 +68,13 @@ def configure(
>>> def my_function():
... return "result"
"""
global _configured_api_key, _configured_pipeline_id, _configured_base_url, _client
global _configured_api_key, _configured_pipeline_id, _configured_base_url, _configured_timeout, _configured_max_retries, _client

_configured_api_key = api_key
_configured_pipeline_id = inference_pipeline_id
_configured_base_url = base_url
_configured_timeout = timeout
_configured_max_retries = max_retries

# Reset the client so it gets recreated with new configuration
_client = None
Expand All @@ -90,6 +98,12 @@ def _get_client() -> Optional[Openlayer]:
if _configured_base_url is not None:
client_kwargs["base_url"] = _configured_base_url

if _configured_timeout is not None:
client_kwargs["timeout"] = _configured_timeout

if _configured_max_retries is not None:
client_kwargs["max_retries"] = _configured_max_retries

if _verify_ssl:
_client = Openlayer(**client_kwargs)
else:
Expand Down
54 changes: 51 additions & 3 deletions tests/test_tracer_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,25 @@ def teardown_method(self):
tracer._configured_api_key = None
tracer._configured_pipeline_id = None
tracer._configured_base_url = None
tracer._configured_timeout = None
tracer._configured_max_retries = None
tracer._client = None

def test_configure_sets_global_variables(self):
"""Test that configure() sets the global configuration variables."""
api_key = "test_api_key"
pipeline_id = "test_pipeline_id"
base_url = "https://test.api.com"
timeout = 30.5
max_retries = 5

tracer.configure(api_key=api_key, inference_pipeline_id=pipeline_id, base_url=base_url)
tracer.configure(api_key=api_key, inference_pipeline_id=pipeline_id, base_url=base_url, timeout=timeout, max_retries=max_retries)

assert tracer._configured_api_key == api_key
assert tracer._configured_pipeline_id == pipeline_id
assert tracer._configured_base_url == base_url
assert tracer._configured_timeout == timeout
assert tracer._configured_max_retries == max_retries

def test_configure_resets_client(self):
"""Test that configure() resets the client to force recreation."""
Expand Down Expand Up @@ -77,6 +83,42 @@ def test_get_client_uses_both_configured_values(self, mock_openlayer: Any) -> No

mock_openlayer.assert_called_once_with(api_key=api_key, base_url=base_url)

@patch("openlayer.lib.tracing.tracer.Openlayer")
def test_get_client_uses_configured_timeout(self, mock_openlayer: Any) -> None:
"""Test that _get_client() uses the configured timeout."""
with patch.object(tracer, "_publish", True):
timeout = 45.5
tracer.configure(timeout=timeout)

tracer._get_client()

mock_openlayer.assert_called_once_with(timeout=timeout)

@patch("openlayer.lib.tracing.tracer.Openlayer")
def test_get_client_uses_configured_max_retries(self, mock_openlayer: Any) -> None:
"""Test that _get_client() uses the configured max_retries."""
with patch.object(tracer, "_publish", True):
max_retries = 10
tracer.configure(max_retries=max_retries)

tracer._get_client()

mock_openlayer.assert_called_once_with(max_retries=max_retries)

@patch("openlayer.lib.tracing.tracer.Openlayer")
def test_get_client_uses_all_configured_values(self, mock_openlayer: Any) -> None:
"""Test that _get_client() uses all configured values together."""
with patch.object(tracer, "_publish", True):
api_key = "configured_api_key"
base_url = "https://configured.api.com"
timeout = 25
max_retries = 3
tracer.configure(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries)

tracer._get_client()

mock_openlayer.assert_called_once_with(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries)

@patch("openlayer.lib.tracing.tracer.DefaultHttpxClient")
@patch("openlayer.lib.tracing.tracer.Openlayer")
def test_get_client_with_ssl_disabled_and_config(self, mock_openlayer: Any, mock_http_client: Any) -> None:
Expand Down Expand Up @@ -150,13 +192,19 @@ def test_configure_with_none_values(self):
"""Test that configure() with None values doesn't overwrite existing config."""
# Set initial configuration
tracer.configure(
api_key="initial_key", inference_pipeline_id="initial_pipeline", base_url="https://initial.com"
api_key="initial_key",
inference_pipeline_id="initial_pipeline",
base_url="https://initial.com",
timeout=60.0,
max_retries=5
)

# Configure with None values
tracer.configure(api_key=None, inference_pipeline_id=None, base_url=None)
tracer.configure(api_key=None, inference_pipeline_id=None, base_url=None, timeout=None, max_retries=None)

# Values should be set to None (this is the expected behavior)
assert tracer._configured_api_key is None
assert tracer._configured_pipeline_id is None
assert tracer._configured_base_url is None
assert tracer._configured_timeout is None
assert tracer._configured_max_retries is None