Skip to content

Commit

Permalink
fix(langchain): token usage reporting
Browse files Browse the repository at this point in the history
  • Loading branch information
nirga committed Oct 2, 2024
1 parent a30bb8c commit 1f525fe
Show file tree
Hide file tree
Showing 5 changed files with 113 additions and 37 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -153,9 +153,30 @@ def _set_chat_response(span: Span, response: LLMResult) -> None:
if not should_send_prompts():
return

input_tokens = 0
output_tokens = 0
total_tokens = 0

i = 0
for generations in response.generations:
for generation in generations:
if (
hasattr(generation, "message")
and hasattr(generation.message, "usage_metadata")
and generation.message.usage_metadata is not None
):
input_tokens += (
generation.message.usage_metadata.get("input_tokens")
or generation.message.usage_metadata.get("prompt_tokens")
or 0
)
output_tokens += (
generation.message.usage_metadata.get("output_tokens")
or generation.message.usage_metadata.get("completion_tokens")
or 0
)
total_tokens += input_tokens + output_tokens

prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{i}"
if hasattr(generation, "text") and generation.text != "":
span.set_attribute(
Expand Down Expand Up @@ -201,6 +222,20 @@ def _set_chat_response(span: Span, response: LLMResult) -> None:
)
i += 1

if input_tokens > 0 or output_tokens > 0 or total_tokens > 0:
span.set_attribute(
SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
input_tokens,
)
span.set_attribute(
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
output_tokens,
)
span.set_attribute(
SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
generation.usage_metadata.get("total_tokens"),
)


class TraceloopCallbackHandler(BaseCallbackHandler):
def __init__(self, tracer: Tracer) -> None:
Expand Down Expand Up @@ -481,13 +516,19 @@ def on_llm_end(

span = self._get_span(run_id)

token_usage = (response.llm_output or {}).get("token_usage")
token_usage = (response.llm_output or {}).get("token_usage") or (
response.llm_output or {}
).get("usage")
if token_usage is not None:
prompt_tokens = token_usage.get("prompt_tokens") or token_usage.get(
"input_token_count"
prompt_tokens = (
token_usage.get("prompt_tokens")
or token_usage.get("input_token_count")
or token_usage.get("input_tokens")
)
completion_tokens = token_usage.get("completion_tokens") or token_usage.get(
"generated_token_count"
completion_tokens = (
token_usage.get("completion_tokens")
or token_usage.get("generated_token_count")
or token_usage.get("output_tokens")
)
total_tokens = token_usage.get("total_tokens") or (
prompt_tokens + completion_tokens
Expand Down
14 changes: 7 additions & 7 deletions packages/opentelemetry-instrumentation-langchain/poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -53,15 +53,32 @@ def test_sequential_chain(exporter):
"SequentialChain.workflow",
] == [span.name for span in spans]

workflow_span = next(span for span in spans if span.name == "SequentialChain.workflow")
task_spans = [span for span in spans if span.name in ["synopsis.task", "LLMChain.task"]]
workflow_span = next(
span for span in spans if span.name == "SequentialChain.workflow"
)
task_spans = [
span for span in spans if span.name in ["synopsis.task", "LLMChain.task"]
]
llm_spans = [span for span in spans if span.name == "OpenAI.completion"]

assert workflow_span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "workflow"
assert workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] == "SequentialChain"
assert all(span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "task" for span in task_spans)
assert all(span.attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "SequentialChain" for span in spans)
assert all(span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH] in ["synopsis", "LLMChain"] for span in llm_spans)
assert (
workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME]
== "SequentialChain"
)
assert all(
span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "task"
for span in task_spans
)
assert all(
span.attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "SequentialChain"
for span in spans
)
assert all(
span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH]
in ["synopsis", "LLMChain"]
for span in llm_spans
)

synopsis_span = next(span for span in spans if span.name == "synopsis.task")
review_span = next(span for span in spans if span.name == "LLMChain.task")
Expand Down Expand Up @@ -200,12 +217,14 @@ def test_stream(exporter):
chunks = list(runnable.stream({"product": "colorful socks"}))
spans = exporter.get_finished_spans()

assert [
"PromptTemplate.task",
"ChatCohere.chat",
"StrOutputParser.task",
"RunnableSequence.workflow",
] == [span.name for span in spans]
assert set(
[
"PromptTemplate.task",
"StrOutputParser.task",
"ChatCohere.chat",
"RunnableSequence.workflow",
]
) == set([span.name for span in spans])
assert len(chunks) == 62


Expand All @@ -223,10 +242,12 @@ async def test_astream(exporter):
chunks.append(chunk)
spans = exporter.get_finished_spans()

assert [
"PromptTemplate.task",
"ChatCohere.chat",
"StrOutputParser.task",
"RunnableSequence.workflow",
] == [span.name for span in spans]
assert set(
[
"PromptTemplate.task",
"ChatCohere.chat",
"StrOutputParser.task",
"RunnableSequence.workflow",
]
) == set([span.name for span in spans])
assert len(chunks) == 144
Original file line number Diff line number Diff line change
Expand Up @@ -36,12 +36,14 @@ class Joke(BaseModel):

spans = exporter.get_finished_spans()

assert [
"ChatPromptTemplate.task",
"ChatOpenAI.chat",
"JsonOutputFunctionsParser.task",
"ThisIsATestChain.workflow",
] == [span.name for span in spans]
assert set(
[
"ChatPromptTemplate.task",
"JsonOutputFunctionsParser.task",
"ChatOpenAI.chat",
"ThisIsATestChain.workflow",
]
) == set([span.name for span in spans])

workflow_span = next(
span for span in spans if span.name == "ThisIsATestChain.workflow"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,9 @@ def test_openai(exporter):
assert (
openai_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.role"]
) == "assistant"
assert openai_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 24
assert openai_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 26
assert openai_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 50


@pytest.mark.vcr
Expand All @@ -115,8 +118,8 @@ class Joke(BaseModel):

assert [
"ChatPromptTemplate.task",
"ChatOpenAI.chat",
"JsonOutputFunctionsParser.task",
"ChatOpenAI.chat",
"RunnableSequence.workflow",
] == [span.name for span in spans]

Expand Down Expand Up @@ -169,6 +172,9 @@ class Joke(BaseModel):
)
== response
)
assert openai_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 76
assert openai_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 35
assert openai_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 111


@pytest.mark.vcr
Expand Down Expand Up @@ -214,6 +220,9 @@ def test_anthropic(exporter):
assert (
anthropic_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.role"]
) == "assistant"
assert anthropic_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 19
assert anthropic_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 22
assert anthropic_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 41
output = json.loads(
workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT]
)
Expand Down Expand Up @@ -286,6 +295,9 @@ def test_bedrock(exporter):
assert (
bedrock_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.role"]
) == "assistant"
assert bedrock_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 16
assert bedrock_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 27
assert bedrock_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 43
output = json.loads(
workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT]
)
Expand Down

0 comments on commit 1f525fe

Please sign in to comment.