File size: 1,574 Bytes
944b1b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fc62e60
944b1b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import os
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult


class FilteredConsoleExporter(SpanExporter):
    WATCHED_SPANS = {
        "call_llm",
        "vector_store",
        "ChampService",
        "convert_messages_langchain",
        "invoke",
        "retrieving documents",
        "PromptSanitizer",
        "sanitize docs_content",
        "sanitize retrieval_query",
        "sanitize_document",
    }

    def export(self, spans):
        for span in spans:
            if (
                span.name in self.WATCHED_SPANS
                and span.end_time is not None
                and span.start_time is not None
            ):
                print(
                    f"[{span.name}] duration: {(span.end_time - span.start_time) / 1e6:.2f}ms"
                )
        return SpanExportResult.SUCCESS

    def shutdown(self):
        pass


def setup_telemetry(app):
    if os.getenv("ENV", "dev").lower() != "dev":
        return

    provider = TracerProvider()
    provider.add_span_processor(BatchSpanProcessor(FilteredConsoleExporter()))
    trace.set_tracer_provider(provider)

    FastAPIInstrumentor.instrument_app(app)
    HTTPXClientInstrumentor().instrument()