Repository: All-Hands-AI/OpenHands
Pull request: Use LLM APIs responses in token counting
**End-user friendly description of the problem this fixes or functionality that this introduces**

- [ ] Include this change in the Release Notes. If checked, you must provide an **end-user friendly** description for your change below

---
**Give a summary of what the PR does, explaining any non-trivial design decisions**

This PR proposes to use the liteLLM token usage data explicitly.
- track it in Metrics
- add two utility methods to link it to events

Please note: technically we also have it in the `tool_call_metadata` optional field in Event, because we have the entire litellm ModelResponse there. This PR proposes to take it out of a ModelResponse, for a few reasons:
- it's not always there: MessageActions, some FinishActions, and user actions don't have a tool. We can fill the data for those from elsewhere (e.g. tokenizer on the corresponding message)
- it seems easier to work with, in a format we need, than the format that litellm returns, though I could be wrong
- we may want to remove the full ModelResponse from events sometime, it contains a lot we don't really need. Events are supposed to be sort of simple and higher level, while the technical details of an API response in litellm are quite complex and of a different nature.


---
Part of https://github.com/All-Hands-AI/OpenHands/issues/6707

Follow-up on https://github.com/All-Hands-AI/OpenHands/pull/5550

**Link of any specific issues this addresses**
Fix #2947 

---

To run this PR locally, use the following command:
```
docker run -it --rm   -p 3000:3000   -v /var/run/docker.sock:/var/run/docker.sock   --add-host host.docker.internal:host-gateway   -e SANDBOX_RUNTIME_CONTAINER_IMAGE=docker.all-hands.dev/all-hands-ai/runtime:9d3f77f-nikolaik   --name openhands-app-9d3f77f   docker.all-hands.dev/all-hands-ai/openhands:9d3f77f
```

Problem statement:
Feat: make use of litellm's response "usage" data
**What problem or use case are you trying to solve?**

**TASK**

We want to enhance our get_token_count() implementation in llm.py, to take advantage of the token counts we are provided from our dependencies, if available, and only fallback to count them if they are not already available.

Read the llm.py file. It uses the litellm library, and you can find some code that uses the Usage object from litellm, too.

This `usage` data, when available, provides live token counts. It will be available for some LLMs, and it will not be available for others.

If these values can be stored and/or attached to the corresponding messages, these values could be calculated and returned in get_token_count() in llm.py and used to live-compute the token counts.

Read also codeact_agent.py file, to see how the LLM is used, and message.py to understand the messages. Read event.py, there will be multiple files with this name, that's okay, read them all, to see Events.

First, read and review the files you need to understand how we can do this.

Then, decide on a solution, and do it.

Unit tests are required.

**/TASK**

**ORIGINAL MESSAGE**

The `usage` data, if filled, provides live token counts. It will work for some LLMs, and it will not be available for others.

If these values could be stored and/or attached to the messages in the codeact_agent.py accordingly, these values could be calculated and returned in get_token_count() in llm.py.


Then this could help a running memory condenser in calculations?

```
ModelResponse(
    id='chatcmpl-47df6b79-cb8d-4f07-920a-8c5551ea7ee4',
    choices=[
        Choices(
            finish_reason='stop',
            index=0,
            message=Message(
                content='{\n  "action": "run",\n  "args": {\n    "command": "ls"\n  }\n}\n',
                role='assistant',
                tool_calls=None,
                function_call=None
            )
        )
    ],
    created=1721067905,
    model='anthropic/claude-3.5-sonnet:beta',
    object='chat.completion',
    system_fingerprint=None,
    usage=Usage(
        prompt_tokens=10,
        completion_tokens=20,
        total_tokens=30
    )
)
```

Hints from issue discussion:
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.
[OpenHands](https://github.com/All-Hands-AI/OpenHands) started fixing the issue! You can monitor the progress [here](https://github.com/All-Hands-AI/OpenHands/actions/runs/12002658456).
An attempt was made to automatically fix this issue, but it was unsuccessful. A branch named 'openhands-fix-issue-2947' has been created with the attempted changes. You can view the branch [here](https://github.com/All-Hands-AI/OpenHands/tree/openhands-fix-issue-2947). Manual intervention may be required.
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.
This issue was closed because it has been stalled for over 30 days with no activity.
I'm working on a fix! enyst can [track my progress at all-hands.dev](https://app.all-hands.dev/conversations/7ab1757d146b45b69ecbac3c9d0dae2d)

Commit message: fix tests

Patch to review:
diff --git a/openhands/core/message_utils.py b/openhands/core/message_utils.py
index 7683c7c4453c..ce8342c6e6b2 100644
--- a/openhands/core/message_utils.py
+++ b/openhands/core/message_utils.py
@@ -29,6 +29,7 @@
 from openhands.events.observation.error import ErrorObservation
 from openhands.events.observation.observation import Observation
 from openhands.events.serialization.event import truncate_content
+from openhands.llm.metrics import Metrics, TokensUsage
 
 
 def events_to_messages(
@@ -362,3 +363,49 @@ def apply_prompt_caching(messages: list[Message]) -> None:
                 -1
             ].cache_prompt = True  # Last item inside the message content
             break
+
+
+def get_single_tokens_usage_for_event(
+    event: Event, metrics: Metrics
+) -> TokensUsage | None:
+    """
+    Returns at most one token usage record for the `model_response.id` in this event's
+    `tool_call_metadata`.
+
+    If no response_id is found, or none match in metrics.tokens_usages, returns [].
+    """
+    if event.tool_call_metadata and event.tool_call_metadata.model_response:
+        response_id = event.tool_call_metadata.model_response.get('id')
+        if response_id:
+            return next(
+                (
+                    usage
+                    for usage in metrics.tokens_usages
+                    if usage.response_id == response_id
+                ),
+                None,
+            )
+    return None
+
+
+def get_tokens_usage_for_event_id(
+    events: list[Event], event_id: int, metrics: Metrics
+) -> TokensUsage | None:
+    """
+    Starting from the event with .id == event_id and moving backwards in `events`,
+    find the first TokensUsage record (if any) associated with a response_id from
+    tool_call_metadata.model_response.id.
+
+    Returns the first match found, or None if none is found.
+    """
+    # find the index of the event with the given id
+    idx = next((i for i, e in enumerate(events) if e.id == event_id), None)
+    if idx is None:
+        return None
+
+    # search backward from idx down to 0
+    for i in range(idx, -1, -1):
+        usage = get_single_tokens_usage_for_event(events[i], metrics)
+        if usage is not None:
+            return usage
+    return None
diff --git a/openhands/llm/llm.py b/openhands/llm/llm.py
index b40f11ca8396..7bd37dc76f79 100644
--- a/openhands/llm/llm.py
+++ b/openhands/llm/llm.py
@@ -497,20 +497,21 @@ def _post_completion(self, response: ModelResponse) -> float:
             stats += 'Response Latency: %.3f seconds\n' % latest_latency.latency
 
         usage: Usage | None = response.get('usage')
+        response_id = response.get('id', 'unknown')
 
         if usage:
             # keep track of the input and output tokens
-            input_tokens = usage.get('prompt_tokens')
-            output_tokens = usage.get('completion_tokens')
+            prompt_tokens = usage.get('prompt_tokens', 0)
+            completion_tokens = usage.get('completion_tokens', 0)
 
-            if input_tokens:
-                stats += 'Input tokens: ' + str(input_tokens)
+            if prompt_tokens:
+                stats += 'Input tokens: ' + str(prompt_tokens)
 
-            if output_tokens:
+            if completion_tokens:
                 stats += (
-                    (' | ' if input_tokens else '')
+                    (' | ' if prompt_tokens else '')
                     + 'Output tokens: '
-                    + str(output_tokens)
+                    + str(completion_tokens)
                     + '\n'
                 )
 
@@ -519,7 +520,7 @@ def _post_completion(self, response: ModelResponse) -> float:
                 'prompt_tokens_details'
             )
             cache_hit_tokens = (
-                prompt_tokens_details.cached_tokens if prompt_tokens_details else None
+                prompt_tokens_details.cached_tokens if prompt_tokens_details else 0
             )
             if cache_hit_tokens:
                 stats += 'Input tokens (cache hit): ' + str(cache_hit_tokens) + '\n'
@@ -528,10 +529,20 @@ def _post_completion(self, response: ModelResponse) -> float:
             # but litellm doesn't separate them in the usage stats
             # so we can read it from the provider-specific extra field
             model_extra = usage.get('model_extra', {})
-            cache_write_tokens = model_extra.get('cache_creation_input_tokens')
+            cache_write_tokens = model_extra.get('cache_creation_input_tokens', 0)
             if cache_write_tokens:
                 stats += 'Input tokens (cache write): ' + str(cache_write_tokens) + '\n'
 
+            # Record in metrics
+            # We'll treat cache_hit_tokens as "cache read" and cache_write_tokens as "cache write"
+            self.metrics.add_tokens_usage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                cache_read_tokens=cache_hit_tokens,
+                cache_write_tokens=cache_write_tokens,
+                response_id=response_id,
+            )
+
         # log the stats
         if stats:
             logger.debug(stats)
diff --git a/openhands/llm/metrics.py b/openhands/llm/metrics.py
index a010bb26916d..30a23158d1b9 100644
--- a/openhands/llm/metrics.py
+++ b/openhands/llm/metrics.py
@@ -17,11 +17,23 @@ class ResponseLatency(BaseModel):
     response_id: str
 
 
+class TokensUsage(BaseModel):
+    """Metric tracking detailed token usage per completion call."""
+
+    model: str
+    prompt_tokens: int
+    completion_tokens: int
+    cache_read_tokens: int
+    cache_write_tokens: int
+    response_id: str
+
+
 class Metrics:
     """Metrics class can record various metrics during running and evaluation.
-    Currently, we define the following metrics:
-        accumulated_cost: the total cost (USD $) of the current LLM.
-        response_latency: the time taken for each LLM completion call.
+    We track:
+      - accumulated_cost and costs
+      - A list of ResponseLatency
+      - A list of TokensUsage (one per call).
     """
 
     def __init__(self, model_name: str = 'default') -> None:
@@ -29,6 +41,7 @@ def __init__(self, model_name: str = 'default') -> None:
         self._costs: list[Cost] = []
         self._response_latencies: list[ResponseLatency] = []
         self.model_name = model_name
+        self._tokens_usages: list[TokensUsage] = []
 
     @property
     def accumulated_cost(self) -> float:
@@ -54,6 +67,16 @@ def response_latencies(self) -> list[ResponseLatency]:
     def response_latencies(self, value: list[ResponseLatency]) -> None:
         self._response_latencies = value
 
+    @property
+    def tokens_usages(self) -> list[TokensUsage]:
+        if not hasattr(self, '_tokens_usages'):
+            self._tokens_usages = []
+        return self._tokens_usages
+
+    @tokens_usages.setter
+    def tokens_usages(self, value: list[TokensUsage]) -> None:
+        self._tokens_usages = value
+
     def add_cost(self, value: float) -> None:
         if value < 0:
             raise ValueError('Added cost cannot be negative.')
@@ -67,10 +90,33 @@ def add_response_latency(self, value: float, response_id: str) -> None:
             )
         )
 
+    def add_tokens_usage(
+        self,
+        prompt_tokens: int,
+        completion_tokens: int,
+        cache_read_tokens: int,
+        cache_write_tokens: int,
+        response_id: str,
+    ) -> None:
+        """Add a single usage record."""
+        self._tokens_usages.append(
+            TokensUsage(
+                model=self.model_name,
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                cache_read_tokens=cache_read_tokens,
+                cache_write_tokens=cache_write_tokens,
+                response_id=response_id,
+            )
+        )
+
     def merge(self, other: 'Metrics') -> None:
+        """Merge 'other' metrics into this one."""
         self._accumulated_cost += other.accumulated_cost
         self._costs += other._costs
-        self._response_latencies += other._response_latencies
+        # use the property so older picked objects that lack the field won't crash
+        self.tokens_usages += other.tokens_usages
+        self.response_latencies += other.response_latencies
 
     def get(self) -> dict:
         """Return the metrics in a dictionary."""
@@ -80,12 +126,14 @@ def get(self) -> dict:
             'response_latencies': [
                 latency.model_dump() for latency in self._response_latencies
             ],
+            'tokens_usages': [usage.model_dump() for usage in self._tokens_usages],
         }
 
     def reset(self):
         self._accumulated_cost = 0.0
         self._costs = []
         self._response_latencies = []
+        self._tokens_usages = []
 
     def log(self):
         """Log the metrics."""
diff --git a/tests/unit/test_llm.py b/tests/unit/test_llm.py
index 1bfee8550698..0f9a8c7362a7 100644
--- a/tests/unit/test_llm.py
+++ b/tests/unit/test_llm.py
@@ -2,6 +2,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from litellm import PromptTokensDetails
 from litellm.exceptions import (
     RateLimitError,
 )
@@ -429,3 +430,62 @@ def test_get_token_count_error_handling(
     mock_logger.error.assert_called_once_with(
         'Error getting token count for\n model gpt-4o\nToken counting failed'
     )
+
+
+@patch('openhands.llm.llm.litellm_completion')
+def test_llm_token_usage(mock_litellm_completion, default_config):
+    # This mock response includes usage details with prompt_tokens,
+    # completion_tokens, prompt_tokens_details.cached_tokens, and model_extra.cache_creation_input_tokens
+    mock_response_1 = {
+        'id': 'test-response-usage',
+        'choices': [{'message': {'content': 'Usage test response'}}],
+        'usage': {
+            'prompt_tokens': 12,
+            'completion_tokens': 3,
+            'prompt_tokens_details': PromptTokensDetails(cached_tokens=2),
+            'model_extra': {'cache_creation_input_tokens': 5},
+        },
+    }
+
+    # Create a second usage scenario to test accumulation and a different response_id
+    mock_response_2 = {
+        'id': 'test-response-usage-2',
+        'choices': [{'message': {'content': 'Second usage test response'}}],
+        'usage': {
+            'prompt_tokens': 7,
+            'completion_tokens': 2,
+            'prompt_tokens_details': PromptTokensDetails(cached_tokens=1),
+            'model_extra': {'cache_creation_input_tokens': 3},
+        },
+    }
+
+    # We'll make mock_litellm_completion return these responses in sequence
+    mock_litellm_completion.side_effect = [mock_response_1, mock_response_2]
+
+    llm = LLM(config=default_config)
+
+    # First call
+    llm.completion(messages=[{'role': 'user', 'content': 'Hello usage!'}])
+
+    # Verify we have exactly one usage record after first call
+    tokens_usage_list = llm.metrics.get()['tokens_usages']
+    assert len(tokens_usage_list) == 1
+    usage_entry_1 = tokens_usage_list[0]
+    assert usage_entry_1['prompt_tokens'] == 12
+    assert usage_entry_1['completion_tokens'] == 3
+    assert usage_entry_1['cache_read_tokens'] == 2
+    assert usage_entry_1['cache_write_tokens'] == 5
+    assert usage_entry_1['response_id'] == 'test-response-usage'
+
+    # Second call
+    llm.completion(messages=[{'role': 'user', 'content': 'Hello again!'}])
+
+    # Now we expect two usage records total
+    tokens_usage_list = llm.metrics.get()['tokens_usages']
+    assert len(tokens_usage_list) == 2
+    usage_entry_2 = tokens_usage_list[-1]
+    assert usage_entry_2['prompt_tokens'] == 7
+    assert usage_entry_2['completion_tokens'] == 2
+    assert usage_entry_2['cache_read_tokens'] == 1
+    assert usage_entry_2['cache_write_tokens'] == 3
+    assert usage_entry_2['response_id'] == 'test-response-usage-2'
diff --git a/tests/unit/test_message_utils.py b/tests/unit/test_message_utils.py
index d3114519c82b..e50ed08f9513 100644
--- a/tests/unit/test_message_utils.py
+++ b/tests/unit/test_message_utils.py
@@ -3,13 +3,18 @@
 import pytest
 
 from openhands.core.message import ImageContent, TextContent
-from openhands.core.message_utils import get_action_message, get_observation_message
+from openhands.core.message_utils import (
+    get_action_message,
+    get_observation_message,
+    get_single_tokens_usage_for_event,
+    get_tokens_usage_for_event_id,
+)
 from openhands.events.action import (
     AgentFinishAction,
     CmdRunAction,
     MessageAction,
 )
-from openhands.events.event import EventSource, FileEditSource, FileReadSource
+from openhands.events.event import Event, EventSource, FileEditSource, FileReadSource
 from openhands.events.observation.browse import BrowserOutputObservation
 from openhands.events.observation.commands import (
     CmdOutputMetadata,
@@ -21,6 +26,7 @@
 from openhands.events.observation.files import FileEditObservation, FileReadObservation
 from openhands.events.observation.reject import UserRejectObservation
 from openhands.events.tool import ToolCallMetadata
+from openhands.llm.metrics import Metrics, TokensUsage
 
 
 def test_cmd_output_observation_message():
@@ -269,3 +275,113 @@ def test_agent_finish_action_with_tool_metadata():
     assert len(result.content) == 1
     assert isinstance(result.content[0], TextContent)
     assert 'Initial thought\nTask completed' in result.content[0].text
+
+
+def test_get_single_tokens_usage_for_event():
+    """Test that we get the single matching usage record (if any) based on the event's model_response.id."""
+    metrics = Metrics(model_name='test-model')
+    usage_record = TokensUsage(
+        model='test-model',
+        prompt_tokens=10,
+        completion_tokens=5,
+        cache_read_tokens=2,
+        cache_write_tokens=1,
+        response_id='test-response-id',
+    )
+    metrics.add_tokens_usage(
+        prompt_tokens=usage_record.prompt_tokens,
+        completion_tokens=usage_record.completion_tokens,
+        cache_read_tokens=usage_record.cache_read_tokens,
+        cache_write_tokens=usage_record.cache_write_tokens,
+        response_id=usage_record.response_id,
+    )
+
+    # Create an event referencing that response_id
+    event = Event()
+    mock_tool_call_metadata = ToolCallMetadata(
+        tool_call_id='test-tool-call',
+        function_name='fake_function',
+        model_response={'id': 'test-response-id'},
+        total_calls_in_response=1,
+    )
+    event._tool_call_metadata = (
+        mock_tool_call_metadata  # normally you'd do event.tool_call_metadata = ...
+    )
+
+    # We should find that usage record
+    found = get_single_tokens_usage_for_event(event, metrics)
+    assert found is not None
+    assert found.prompt_tokens == 10
+    assert found.response_id == 'test-response-id'
+
+    # If we change the event's response ID, we won't find anything
+    mock_tool_call_metadata.model_response.id = 'some-other-id'
+    found2 = get_single_tokens_usage_for_event(event, metrics)
+    assert found2 is None
+
+    # If the event has no tool_call_metadata, also returns None
+    event._tool_call_metadata = None
+    found3 = get_single_tokens_usage_for_event(event, metrics)
+    assert found3 is None
+
+
+def test_get_tokens_usage_for_event_id():
+    """
+    Test that we search backward from the event with the given id,
+    finding the first usage record that matches a response_id in that or previous events.
+    """
+    metrics = Metrics(model_name='test-model')
+    usage_1 = TokensUsage(
+        model='test-model',
+        prompt_tokens=12,
+        completion_tokens=3,
+        cache_read_tokens=2,
+        cache_write_tokens=5,
+        response_id='resp-1',
+    )
+    usage_2 = TokensUsage(
+        model='test-model',
+        prompt_tokens=7,
+        completion_tokens=2,
+        cache_read_tokens=1,
+        cache_write_tokens=3,
+        response_id='resp-2',
+    )
+    metrics._tokens_usages.append(usage_1)
+    metrics._tokens_usages.append(usage_2)
+
+    # Build a list of events
+    events = []
+    for i in range(5):
+        e = Event()
+        e._id = i
+        # We'll attach usage_1 to event 1, usage_2 to event 3
+        if i == 1:
+            e._tool_call_metadata = ToolCallMetadata(
+                tool_call_id='tid1',
+                function_name='fn1',
+                model_response={'id': 'resp-1'},
+                total_calls_in_response=1,
+            )
+        elif i == 3:
+            e._tool_call_metadata = ToolCallMetadata(
+                tool_call_id='tid2',
+                function_name='fn2',
+                model_response={'id': 'resp-2'},
+                total_calls_in_response=1,
+            )
+        events.append(e)
+
+    # If we ask for event_id=3, we find usage_2 immediately
+    found_3 = get_tokens_usage_for_event_id(events, 3, metrics)
+    assert found_3 is not None
+    assert found_3.response_id == 'resp-2'
+
+    # If we ask for event_id=2, no usage in event2, so we check event1 -> usage_1 found
+    found_2 = get_tokens_usage_for_event_id(events, 2, metrics)
+    assert found_2 is not None
+    assert found_2.response_id == 'resp-1'
+
+    # If we ask for event_id=0, no usage in event0 or earlier, so return None
+    found_0 = get_tokens_usage_for_event_id(events, 0, metrics)
+    assert found_0 is None


Code review comments:
diff_hunk: @@ -17,18 +17,31 @@ class ResponseLatency(BaseModel):
     response_id: str
 
 
+class TokensUsage(BaseModel):
path: openhands/llm/metrics.py
text: @user1:
Minor nit, but I'd expect this to be called `TokenUsage` instead of the pluralized form.

@author:
Hah, I felt the same, it reads strange 😅 , o3-mini wanted it for some mysterious reason. Fixed!

They are going to align humans to their preferences, aren't they? 😂

diff_hunk: @@ -362,3 +363,49 @@ def apply_prompt_caching(messages: list[Message]) -> None:
                 -1
             ].cache_prompt = True  # Last item inside the message content
             break
+
+
+def get_single_tokens_usage_for_event(
+    event: Event, metrics: Metrics
+) -> TokensUsage | None:
+    """
+    Returns at most one token usage record for the `model_response.id` in this event's
+    `tool_call_metadata`.
+
+    If no response_id is found, or none match in metrics.tokens_usages, returns [].
path: openhands/core/message_utils.py
text: @user1:
```suggestion
    If no response_id is found, or none match in metrics.tokens_usages, returns None.
```