setup.py 1.1 KB

1234567891011121314151617181920212223242526272829303132
  1. from fastapi import FastAPI
  2. from opentelemetry import trace
  3. from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
  4. from opentelemetry.sdk.resources import SERVICE_NAME, Resource
  5. from opentelemetry.sdk.trace import TracerProvider
  6. from sqlalchemy import Engine
  7. from open_webui.utils.telemetry.exporters import LazyBatchSpanProcessor
  8. from open_webui.utils.telemetry.instrumentors import Instrumentor
  9. from open_webui.utils.telemetry.metrics import setup_metrics
  10. from open_webui.env import (
  11. OTEL_SERVICE_NAME,
  12. OTEL_EXPORTER_OTLP_ENDPOINT,
  13. ENABLE_OTEL_METRICS,
  14. )
  15. def setup(app: FastAPI, db_engine: Engine):
  16. # set up trace
  17. trace.set_tracer_provider(
  18. TracerProvider(
  19. resource=Resource.create(attributes={SERVICE_NAME: OTEL_SERVICE_NAME})
  20. )
  21. )
  22. # otlp export
  23. exporter = OTLPSpanExporter(endpoint=OTEL_EXPORTER_OTLP_ENDPOINT)
  24. trace.get_tracer_provider().add_span_processor(LazyBatchSpanProcessor(exporter))
  25. Instrumentor(app=app, db_engine=db_engine).instrument()
  26. # set up metrics only if enabled
  27. if ENABLE_OTEL_METRICS:
  28. setup_metrics(app)