Troubleshooting Span Issues in Langchain with Flask and Vertex AI
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.
