Troubleshooting Span Issues in Langchain with Flask and Vertex AI

·Jul 17, 2024 08:57 AM

Hi all ! I'm having an issue adding some data to my Span. My rag is a Langchain application running on a Flask server, I call Sonnet 3.5 running on vertex.ai. I noticed the token count aren't saved with the default LanchainInstrumentor, so I decided to add them manually. Here's my Custom Callback Handler to get the tokens usage:

class VertexAICallbackHandler(BaseCallbackHandler):

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Collect token usage."""
        prompt_tokens = response.llm_output.get('usage').get('input_tokens')
        response_tokens = response.llm_output.get('usage').get('output_tokens')
        current_span = trace.get_current_span()
        current_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, prompt_tokens)
        current_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, response_tokens)
        current_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, prompt_tokens + response_tokens)

My issue is that the current_span i get from trace.get_current_span() is always invalid. My guess is I can't retrieve the current span because it's not on the same thread. But i'm clueless on how to work around.

from typing import Any from uuid import UUID from langchain_core.outputs import LLMResult from langchain_core.tracers import Run from langchain_openai import ChatOpenAI from openinference.instrumentation.langchain import LangChainInstrumentor, _tracer from opentelemetry.sdk import trace as trace_sdk from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor class _MyTracer(_tracer.OpenInferenceTracer): def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> Run: span = LangChainInstrumentor().get_span(run_id) span.set_attribute("Hello", "world") return super().on_llm_end(response, run_id=run_id, **kwargs) setattr(_tracer, "OpenInferenceTracer", _MyTracer) tracer_provider = trace_sdk.TracerProvider() tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) LangChainInstrumentor().instrument(tracer_provider=tracer_provider) if __name__ == "__main__": for chunk in ChatOpenAI(model_name="gpt-3.5-turbo").stream("Hi"): print(chunk.content, end="")

class _MyTracer(_tracer.OpenInferenceTracer): def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> Run: span = LangChainInstrumentor().get_span(run_id) prompt_tokens = response.llm_output.get('usage').get('input_tokens') response_tokens = response.llm_output.get('usage').get('output_tokens') span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, prompt_tokens) span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, response_tokens) span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, prompt_tokens + response_tokens) return super().on_llm_end(response, run_id=run_id, **kwargs) setattr(_tracer, "OpenInferenceTracer", _MyTracer) tracer_provider = trace_sdk.TracerProvider( resource=Resource({ResourceAttributes.PROJECT_NAME: get_env_project_name()}), span_limits=trace_sdk.SpanLimits(max_attributes=10_000), ) tracer_provider.add_span_processor(SimpleSpanProcessor(_OpenInferenceExporter())) LangChainInstrumentor().instrument(skip_dep_check=True, tracer_provider=tracer_provider)

12 comments