fix: resolve type mismatch in usage summary propagation (#55)

Update BaseLM.get_last_usage type hint to return ModelUsageSummary, reflecting actual client implementations.
Wrap atomic usage stats in UsageSummary within LMHandler to satisfy RLMChatCompletion type expectations and ensure consistent serialization structure.

Author: thomasgauthier

rlm/clients/base_lm.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Any
 
-from rlm.core.types import UsageSummary
+from rlm.core.types import ModelUsageSummary, UsageSummary
 
 
 class BaseLM(ABC):
@@ -28,6 +28,6 @@ def get_usage_summary(self) -> UsageSummary:
         raise NotImplementedError
 
     @abstractmethod
-    def get_last_usage(self) -> UsageSummary:
+    def get_last_usage(self) -> ModelUsageSummary:
         """Get the last cost summary of the model."""
         raise NotImplementedError

rlm/core/lm_handler.py

@@ -51,10 +51,12 @@ def _handle_single(self, request: LMRequest, handler: "LMHandler") -> LMResponse
         content = client.completion(request.prompt)
         end_time = time.perf_counter()
 
-        usage_summary = client.get_last_usage()
+        model_usage = client.get_last_usage()
+        root_model = request.model or client.model_name
+        usage_summary = UsageSummary(model_usage_summaries={root_model: model_usage})
         return LMResponse.success_response(
             chat_completion=RLMChatCompletion(
-                root_model=request.model or client.model_name,
+                root_model=root_model,
                 prompt=request.prompt,
                 response=content,
                 usage_summary=usage_summary,
@@ -76,11 +78,13 @@ async def run_all():
         end_time = time.perf_counter()
 
         total_time = end_time - start_time
-        usage_summary = client.get_last_usage()
+        model_usage = client.get_last_usage()
+        root_model = request.model or client.model_name
+        usage_summary = UsageSummary(model_usage_summaries={root_model: model_usage})
 
         chat_completions = [
             RLMChatCompletion(
-                root_model=request.model or client.model_name,
+                root_model=root_model,
                 prompt=prompt,
                 response=content,
                 usage_summary=usage_summary,