comet-ml
[FR]: DSPy Integration - Add token usage to trace.
Proposal summary
I would like token consumption to be added to DSPy integration. I followed the guide to integrate Opik by using the Opik callback in DSPy settings configuration and it worked well but it is only registering provider and model while it does not track token consumption.
The trace i get is as follows:
function (span, general, tagged)
|--- Module (span, general, with graph)
|--- Predict (span, llm, with graph)
|-- LLM (span, llm, with provider and model)
Token consumption is nowhere to be seen.
I am using openrouter as provider and i am receiving bodies like this as answer:
{
"id":"identifier",
"provider":"provider_different_than_openrouter", # Because of openrouters internal routing
"model":"selected_model",
"object":"chat.completion",
"created":1234567890,
"choices":[
{
"logprobs":null,
"finish_reason":"stop",
"native_finish_reason":"stop",
"index":0,
"message":{
"role":"assistant",
# I've seen that this content is the value of outputs inside on_lm_end callback
"content":"response",
"refusal":null,
"reasoning":null
}
}
],
# Usage info
"usage":{
"prompt_tokens":123,
"completion_tokens":11,
"total_tokens":134,
"prompt_tokens_details":null
}
}
Motivation
As a workaround, I managed to add token consumption using this code right after calling the module:
opik_context.update_current_span(
usage=usage, # I manually recover token_usage from module call
)
But is is being registered here:
function (span, general, tagged) <------------
|--- Module (span, general, with graph)
|--- Predict (span, llm, with graph)
|-- LLM (span, llm, with provider and model)
Which is not ideal.
I think this feature is necessary in order to track token consumption and cost.
Hey @ralsofias thanks for reporting We’re taking a look Can you confirm opik version used on the SDK
I am using opik version 1.8.75 and dspy version 3.0.3.
In order to be able to check that body i went to .venv/lib/python3.12/site-packages/litellm/llms/custom_httpx/llm_http_handler.py and added some prints.
Excuse me @idoberko2, but I don't know if you saw my last message in this other issue (https://github.com/comet-ml/opik/issues/3680) and I have no permission to reopen it.
The problem with that other issue is that while it is true that DSPy integration registers the right model and openrouter as provider, the response received from openrouter contains the "real" provider (with variable cost) selected inside openrouter and that is not reflected on the trace (don't know if it should be reflected or if it is doable).
