feat(haystack): add token count tracking for embeddings

[Hansehart]

Fixes #2322

Summary

Adds token usage tracking for embedding operations in the Haystack instrumentation package. This enables cost monitoring and usage analytics for embedding API calls when using observability platforms like Langfuse.

Changes

• Added _get_embedding_token_count_attributes() function to extract token counts from embedder responses
• Integrated token tracking into the EMBEDDER component handler in _ComponentRunWrapper
• Added test assertions in test_openai_document_embedder_embedding_span_has_expected_attributes to verify token count attributes
• Follows the same implementation pattern as existing _get_llm_token_count_attributes()

Implementation Details

The new function supports two response formats:

1. Standard embedding API format: response["usage"]
2. Custom components format: response["meta"]["usage"]

Token counts tracked:

• LLM_TOKEN_COUNT_PROMPT - Input tokens consumed
• LLM_TOKEN_COUNT_TOTAL - Total tokens used
• Note: LLM_TOKEN_COUNT_COMPLETION is not tracked for embeddings (no output generation)

Test Coverage

Updated existing test to verify token count attributes are correctly captured from OpenAI embedding responses.

---

Note

Adds token usage extraction for embedding responses and records model/tokens on EMBEDDER spans, with tests.

• Instrumentation (Haystack):
  • EMBEDDER spans now include token usage and model details.
    • New _get_embedding_token_count_attributes() extracts prompt_tokens, total_tokens, and model from response["usage"] or response["meta"]["usage"].
    • Integrated into _ComponentRunWrapper for ComponentType.EMBEDDER, merging with existing embedding attributes and setting both LLM_MODEL_NAME and EMBEDDING_MODEL_NAME when available.
• Tests:
  • Update test_openai_document_embedder_embedding_span_has_expected_attributes to assert LLM_MODEL_NAME, LLM_TOKEN_COUNT_PROMPT, and LLM_TOKEN_COUNT_TOTAL for OpenAI embeddings.

^{Written by Cursor Bugbot for commit 835cec0. This will update automatically on new commits. Configure here.}

[github-actions]

CLA Assistant Lite bot All contributors have signed the CLA ✍️ ✅

[Hansehart]

I have read the CLA Document and I hereby sign the CLA

[cursor]

Bug: Model Name Overwritten by Metadata

The EMBEDDING_MODEL_NAME attribute is set twice for embedder components: once from the component's configuration and again from the response metadata. The response metadata value overwrites the component's configured value, potentially causing inconsistent model name reporting.

Additional Locations (1)

• python/instrumentation/openinference-instrumentation-haystack/src/openinference/instrumentation/haystack/_wrappers.py#L556-L561

 

[Hansehart]

I configured response metadata as primary and component config as fallback! Is This acceptable?

[nate-mar]

@Hansehart Thanks so much for this! We very much appreciate it. Let us know if you need help resolving the failed CI checks. My colleague @axiomofjoy will take a look at this when he gets a chance.

[Hansehart]

I’ll try to get this fixed in the next few days. Hints or pointers are very welcome though.

[caroger]

I’ll try to get this fixed in the next few days. Hints or pointers are very welcome though.

Hi @Hansehart , you can see details of the CI/testing result by clicking into the checks. For instance - https://github.com/Arize-ai/openinference/actions/runs/18636229373/job/53302442562?pr=2323

It seems the CI failed at the mypy type checking steps.

For fast iterations - you can re-run the CI locally with uvx --with tox-uv tox run -e py313-ci-haystack-latest