enrich llm.log with token counts

Everyone can contribute. Help move this issue forward while earning points, leveling up and collecting rewards.

Customers have requested token counts as part of the logs for self-hosted models. 1 2

Other work - https://gitlab.com/groups/gitlab-org/-/epics/16748+s

Implementation details

  1. We can use litellm to estimate the number of tokens for an input/output
  2. These can be computed right before we request prompt to litellm, and right after it returns responses. This happens in our litellm client
  3. We should include this as part of the telemetry. For example, for code completions we are pass the prompt input size and model output length (but not the tokens) https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/blob/main/ai_gateway/instrumentators/base.py#L99

Definition of Done

  • Token counts (in/out) have been added to the AIGW logs
Edited by 🤖 GitLab Bot 🤖