From ef309c0e330ae3c5b0cc004d4484b54057ae78d1 Mon Sep 17 00:00:00 2001 From: "L.B." Date: Sat, 9 Nov 2024 15:48:05 -0500 Subject: [PATCH] Fix callback handlers in tracing.py The imported observability tools don't seem to have been included and cause the `chain-server` container to crash stating that the module can't be found. As a workaround, I've added locally defined callbacks to unblock the use of the "Try it Now!" demo experience. ``` from RAG.tools.observability.langchain import opentelemetry_callback as langchain_otel_cb from RAG.tools.observability.llamaindex import opentelemetry_callback as llama_index_otel_cb ``` Error in container: ``` from RAG.tools.observability.langchain import opentelemetry_callback as langchain_otel_cb 2024-11-09 11:38:58 ModuleNotFoundError: No module named 'RAG.tools.observability' ``` --- RAG/src/chain_server/tracing.py | 39 ++++++++++++++++++++++++++------- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/RAG/src/chain_server/tracing.py b/RAG/src/chain_server/tracing.py index 5a1607b3..eb6eeb69 100644 --- a/RAG/src/chain_server/tracing.py +++ b/RAG/src/chain_server/tracing.py @@ -29,9 +29,6 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from RAG.tools.observability.langchain import opentelemetry_callback as langchain_otel_cb -from RAG.tools.observability.llamaindex import opentelemetry_callback as llama_index_otel_cb - # Configure tracer used by the Chain Server to create spans resource = Resource.create({SERVICE_NAME: "chain-server"}) provider = TracerProvider(resource=resource) @@ -45,11 +42,37 @@ # Configure Propagator used for processing trace context received by the Chain Server propagator = TraceContextTextMapPropagator() - # Configure Langchain OpenTelemetry callback handler - langchain_cb_handler = langchain_otel_cb.OpenTelemetryCallbackHandler(tracer) - - # Configure LlamaIndex OpenTelemetry callback handler - llama_index_cb_handler = llama_index_otel_cb.OpenTelemetryCallbackHandler(tracer) + # Define Langchain OpenTelemetry callback handler + class LangchainOpenTelemetryCallbackHandler(langchain_base_cb_handler): + def __init__(self, tracer): + super().__init__() + self.tracer = tracer + + def on_start(self, *args, **kwargs): + with self.tracer.start_as_current_span("LangchainOperation") as span: + span.set_attributes({"operation": "start"}) + + def on_end(self, *args, **kwargs): + with self.tracer.start_as_current_span("LangchainOperationEnd") as span: + span.set_attributes({"operation": "end"}) + + langchain_cb_handler = LangchainOpenTelemetryCallbackHandler(tracer) + + # Define LlamaIndex OpenTelemetry callback handler + class LlamaIndexOpenTelemetryCallbackHandler(llama_index_base_cb_handler): + def __init__(self, tracer): + super().__init__() + self.tracer = tracer + + def on_start(self, *args, **kwargs): + with self.tracer.start_as_current_span("LlamaIndexOperation") as span: + span.set_attributes({"operation": "start"}) + + def on_end(self, *args, **kwargs): + with self.tracer.start_as_current_span("LlamaIndexOperationEnd") as span: + span.set_attributes({"operation": "end"}) + + llama_index_cb_handler = LlamaIndexOpenTelemetryCallbackHandler(tracer) else: propagator = CompositePropagator([]) # No-op propagator