268 lines
10 KiB
Python
268 lines
10 KiB
Python
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from typing import Callable, List, Optional, Set, Union
|
|
from colorama import Fore
|
|
from opentelemetry.sdk.trace import SpanProcessor, ReadableSpan
|
|
from opentelemetry.sdk.trace.sampling import Sampler
|
|
from opentelemetry.sdk.trace.export import SpanExporter
|
|
from opentelemetry.sdk.metrics.export import MetricExporter
|
|
from opentelemetry.sdk._logs.export import LogExporter
|
|
from opentelemetry.sdk.resources import SERVICE_NAME
|
|
from opentelemetry.propagators.textmap import TextMapPropagator
|
|
from opentelemetry.util.re import parse_env_headers
|
|
|
|
from traceloop.sdk.images.image_uploader import ImageUploader
|
|
from traceloop.sdk.metrics.metrics import MetricsWrapper
|
|
from traceloop.sdk.logging.logging import LoggerWrapper
|
|
from traceloop.sdk.instruments import Instruments
|
|
from traceloop.sdk.config import (
|
|
is_content_tracing_enabled,
|
|
is_tracing_enabled,
|
|
is_metrics_enabled,
|
|
is_logging_enabled,
|
|
)
|
|
from traceloop.sdk.fetcher import Fetcher
|
|
from traceloop.sdk.tracing.tracing import (
|
|
TracerWrapper,
|
|
set_association_properties,
|
|
set_external_prompt_tracing_context,
|
|
)
|
|
from typing import Dict
|
|
from traceloop.sdk.client.client import Client
|
|
|
|
|
|
class Traceloop:
|
|
AUTO_CREATED_KEY_PATH = str(
|
|
Path.home() / ".cache" / "traceloop" / "auto_created_key"
|
|
)
|
|
AUTO_CREATED_URL = str(Path.home() / ".cache" / "traceloop" / "auto_created_url")
|
|
|
|
__tracer_wrapper: TracerWrapper
|
|
__fetcher: Optional[Fetcher] = None
|
|
__app_name: Optional[str] = None
|
|
__client: Optional[Client] = None
|
|
|
|
@staticmethod
|
|
def init(
|
|
app_name: str = sys.argv[0],
|
|
api_endpoint: str = "https://api.traceloop.com",
|
|
api_key: Optional[str] = None,
|
|
enabled: bool = True,
|
|
headers: Dict[str, str] = {},
|
|
disable_batch=False,
|
|
telemetry_enabled: bool = True,
|
|
exporter: Optional[SpanExporter] = None,
|
|
metrics_exporter: MetricExporter = None,
|
|
metrics_headers: Dict[str, str] = None,
|
|
logging_exporter: LogExporter = None,
|
|
logging_headers: Dict[str, str] = None,
|
|
processor: Optional[Union[SpanProcessor, List[SpanProcessor]]] = None,
|
|
propagator: TextMapPropagator = None,
|
|
sampler: Optional[Sampler] = None,
|
|
traceloop_sync_enabled: bool = False,
|
|
should_enrich_metrics: bool = True,
|
|
resource_attributes: dict = {},
|
|
instruments: Optional[Set[Instruments]] = None,
|
|
block_instruments: Optional[Set[Instruments]] = None,
|
|
image_uploader: Optional[ImageUploader] = None,
|
|
span_postprocess_callback: Optional[Callable[[ReadableSpan], None]] = None,
|
|
) -> Optional[Client]:
|
|
if not enabled:
|
|
TracerWrapper.set_disabled(True)
|
|
print(
|
|
Fore.YELLOW
|
|
+ "Traceloop instrumentation is disabled via init flag"
|
|
+ Fore.RESET
|
|
)
|
|
return
|
|
|
|
api_endpoint = os.getenv("TRACELOOP_BASE_URL") or api_endpoint
|
|
api_key = os.getenv("TRACELOOP_API_KEY") or api_key
|
|
Traceloop.__app_name = app_name
|
|
|
|
if not is_tracing_enabled():
|
|
print(Fore.YELLOW + "Tracing is disabled" + Fore.RESET)
|
|
return
|
|
|
|
enable_content_tracing = is_content_tracing_enabled()
|
|
|
|
if exporter or processor:
|
|
print(Fore.GREEN + "Traceloop exporting traces to a custom exporter")
|
|
|
|
headers = os.getenv("TRACELOOP_HEADERS") or headers
|
|
|
|
if isinstance(headers, str):
|
|
headers = parse_env_headers(headers)
|
|
|
|
if (
|
|
not exporter
|
|
and not processor
|
|
and api_endpoint == "https://api.traceloop.com"
|
|
and not api_key
|
|
):
|
|
print(
|
|
Fore.RED
|
|
+ "Error: Missing Traceloop API key,"
|
|
+ " go to https://app.traceloop.com/settings/api-keys to create one"
|
|
)
|
|
print("Set the TRACELOOP_API_KEY environment variable to the key")
|
|
print(Fore.RESET)
|
|
return
|
|
|
|
if not exporter and not processor and headers:
|
|
print(
|
|
Fore.GREEN
|
|
+ f"Traceloop exporting traces to {api_endpoint}, authenticating with custom headers"
|
|
)
|
|
|
|
if api_key and not exporter and not processor and not headers:
|
|
print(
|
|
Fore.GREEN
|
|
+ f"Traceloop exporting traces to {api_endpoint} authenticating with bearer token"
|
|
)
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
print(Fore.RESET)
|
|
|
|
# Tracer init
|
|
resource_attributes.update({SERVICE_NAME: app_name})
|
|
TracerWrapper.set_static_params(
|
|
resource_attributes, enable_content_tracing, api_endpoint, headers
|
|
)
|
|
Traceloop.__tracer_wrapper = TracerWrapper(
|
|
disable_batch=disable_batch,
|
|
processor=processor,
|
|
propagator=propagator,
|
|
exporter=exporter,
|
|
sampler=sampler,
|
|
should_enrich_metrics=should_enrich_metrics,
|
|
image_uploader=image_uploader or ImageUploader(api_endpoint, api_key),
|
|
instruments=instruments,
|
|
block_instruments=block_instruments,
|
|
span_postprocess_callback=span_postprocess_callback,
|
|
)
|
|
|
|
metrics_disabled_by_config = not is_metrics_enabled()
|
|
has_custom_spans_pipeline = processor or exporter
|
|
custom_trace_without_custom_metrics = has_custom_spans_pipeline and not metrics_exporter
|
|
|
|
if metrics_disabled_by_config or custom_trace_without_custom_metrics:
|
|
print(Fore.YELLOW + "Metrics are disabled" + Fore.RESET)
|
|
else:
|
|
metrics_endpoint = os.getenv("TRACELOOP_METRICS_ENDPOINT") or api_endpoint
|
|
metrics_headers = (
|
|
os.getenv("TRACELOOP_METRICS_HEADERS") or metrics_headers or headers
|
|
)
|
|
if metrics_exporter or processor:
|
|
print(Fore.GREEN + "Traceloop exporting metrics to a custom exporter")
|
|
|
|
MetricsWrapper.set_static_params(
|
|
resource_attributes, metrics_endpoint, metrics_headers
|
|
)
|
|
Traceloop.__metrics_wrapper = MetricsWrapper(exporter=metrics_exporter)
|
|
|
|
if is_logging_enabled() and (logging_exporter or not exporter):
|
|
logging_endpoint = os.getenv("TRACELOOP_LOGGING_ENDPOINT") or api_endpoint
|
|
logging_headers = (
|
|
os.getenv("TRACELOOP_LOGGING_HEADERS") or logging_headers or headers
|
|
)
|
|
if logging_exporter or processor:
|
|
print(Fore.GREEN + "Traceloop exporting logs to a custom exporter")
|
|
|
|
LoggerWrapper.set_static_params(
|
|
resource_attributes, logging_endpoint, logging_headers
|
|
)
|
|
Traceloop.__logger_wrapper = LoggerWrapper(exporter=logging_exporter)
|
|
|
|
if (
|
|
api_endpoint.find("traceloop.com") != -1
|
|
and api_key
|
|
and (exporter is None)
|
|
and (processor is None)
|
|
):
|
|
if traceloop_sync_enabled:
|
|
Traceloop.__fetcher = Fetcher(base_url=api_endpoint, api_key=api_key)
|
|
Traceloop.__fetcher.run()
|
|
print(
|
|
Fore.GREEN
|
|
+ "Traceloop syncing configuration and prompts"
|
|
+ Fore.RESET
|
|
)
|
|
Traceloop.__client = Client(
|
|
api_key=api_key, app_name=app_name, api_endpoint=api_endpoint
|
|
)
|
|
return Traceloop.__client
|
|
|
|
@staticmethod
|
|
def set_association_properties(properties: dict) -> None:
|
|
set_association_properties(properties)
|
|
|
|
def set_prompt(template: str, variables: dict, version: int):
|
|
set_external_prompt_tracing_context(template, variables, version)
|
|
|
|
@staticmethod
|
|
def get_default_span_processor(
|
|
disable_batch: bool = False,
|
|
api_endpoint: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
exporter: Optional[SpanExporter] = None
|
|
) -> SpanProcessor:
|
|
"""
|
|
Creates and returns the default Traceloop span processor.
|
|
|
|
This function allows users to get the default Traceloop span processor
|
|
to combine it with their custom processors when using the processors parameter.
|
|
|
|
Args:
|
|
disable_batch: If True, uses SimpleSpanProcessor, otherwise BatchSpanProcessor
|
|
api_endpoint: The endpoint URL for the exporter (uses current config if None)
|
|
headers: Headers for the exporter (uses current config if None)
|
|
exporter: Custom exporter to use (creates default if None)
|
|
|
|
Returns:
|
|
SpanProcessor: The default Traceloop span processor
|
|
|
|
Example:
|
|
# Get the default processor and combine with custom one
|
|
default_processor = Traceloop.get_default_span_processor()
|
|
custom_processor = MyCustomSpanProcessor()
|
|
|
|
Traceloop.init(
|
|
processors=[default_processor, custom_processor]
|
|
)
|
|
"""
|
|
from traceloop.sdk.tracing.tracing import get_default_span_processor
|
|
if headers is None:
|
|
if api_key is None:
|
|
api_key = os.getenv("TRACELOOP_API_KEY")
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
if api_endpoint is None:
|
|
api_endpoint = os.getenv("TRACELOOP_BASE_URL") or "https://api.traceloop.com"
|
|
return get_default_span_processor(disable_batch, api_endpoint, headers, exporter)
|
|
|
|
@staticmethod
|
|
def get():
|
|
"""
|
|
Returns the shared SDK client instance, using the current global configuration.
|
|
|
|
To use the SDK as a singleton, first make sure you have called :func:`Traceloop.init()`
|
|
at startup time. Then ``get()`` will return the same shared :class:`Traceloop.client.Client`
|
|
instance each time. The client will be initialized if it has not been already.
|
|
|
|
If you need to create multiple client instances with different configurations, instead of this
|
|
singleton approach you can call the :class:`Traceloop.client.Client` constructor directly instead.
|
|
"""
|
|
if not Traceloop.__client:
|
|
raise Exception(
|
|
"Client not initialized, you should call Traceloop.init() first. "
|
|
"If you are still getting this error - you are missing the api key"
|
|
)
|
|
return Traceloop.__client
|