Token Usage Service

Utilities for extracting and tracking LLM token usage from pipeline executions.

Key Functions

• extract_token_usage() - Extract token counts from pipeline outputs

• record_usage() - Record token usage to database

ai4drpm.services.engine.token_usage_service.extract_token_usage_from_pipeline_result(pipeline_result, component_names=None)

Extract token usage information from a Haystack pipeline execution result.

Haystack generators (OpenAIGenerator, etc.) return token usage in their output’s meta list. Each meta dict contains a usage dict with: - prompt_tokens - completion_tokens - total_tokens

Parameters:

• pipeline_result (Dict[str, Any]) – The result dict from Pipeline.run()

• component_names (Optional[List[str]]) – Optional list of component names to extract from. If None, extracts from all components with usage data.

Returns:

• component_name: Name of the generator component

• model: Model name (if available)

• prompt_tokens: Input tokens

• completion_tokens: Output tokens

• total_tokens: Total tokens

Return type:

List of dicts, each containing

Example

pipeline_result = {

“generator”: {

“replies”: [”…”], “meta”: [{

“model”: “mistral-large-latest”, “usage”: {

“prompt_tokens”: 1500, “completion_tokens”: 500, “total_tokens”: 2000

}

}]

}

}

usage_list = extract_token_usage_from_pipeline_result(pipeline_result) # Returns: [{ # “component_name”: “generator”, # “model”: “mistral-large-latest”, # “prompt_tokens”: 1500, # “completion_tokens”: 500, # “total_tokens”: 2000 # }]

ai4drpm.services.engine.token_usage_service.aggregate_token_usage(usage_records)

Aggregate multiple token usage records into totals.

Parameters:

usage_records (List[Dict[str, Any]]) – List of usage dicts from extract_token_usage_from_pipeline_result

Returns:

• prompt_tokens: Total prompt tokens

• completion_tokens: Total completion tokens

• total_tokens: Total tokens

• models: Set of unique models used

• component_count: Number of components with usage

Return type:

Dict with aggregated totals

ai4drpm.services.engine.token_usage_service.infer_provider_from_model(model)

Infer the provider name from the model name.

Parameters:

model (str) – Model name (e.g., “mistral-large-latest”, “gpt-4”)

Return type:

str

Returns:

Provider name (e.g., “mistral”, “openai”)