setup.py 944 B

1234567891011121314151617181920212223
  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.trace.exporters import LazyBatchSpanProcessor
  8. from open_webui.utils.trace.instrumentors import Instrumentor
  9. from open_webui.env import OTEL_SERVICE_NAME, OTEL_EXPORTER_OTLP_ENDPOINT
  10. def setup(app: FastAPI, db_engine: Engine):
  11. # set up trace
  12. trace.set_tracer_provider(
  13. TracerProvider(
  14. resource=Resource.create(attributes={SERVICE_NAME: OTEL_SERVICE_NAME})
  15. )
  16. )
  17. # otlp export
  18. exporter = OTLPSpanExporter(endpoint=OTEL_EXPORTER_OTLP_ENDPOINT)
  19. trace.get_tracer_provider().add_span_processor(LazyBatchSpanProcessor(exporter))
  20. Instrumentor(app=app, db_engine=db_engine).instrument()