From 5b2834fc7a46427b22fd923dfeb2a537ab697b40 Mon Sep 17 00:00:00 2001 From: Thomas Gauthier-Caron Date: Sun, 18 Jan 2026 12:14:26 -0500 Subject: [PATCH] fix: resolve type mismatch in usage summary propagation 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. --- rlm/clients/base_lm.py | 4 ++-- rlm/core/lm_handler.py | 12 ++++++++---- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/rlm/clients/base_lm.py b/rlm/clients/base_lm.py index d742650ab..365536621 100644 --- a/rlm/clients/base_lm.py +++ b/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 diff --git a/rlm/core/lm_handler.py b/rlm/core/lm_handler.py index 46d20339b..9eb751011 100644 --- a/rlm/core/lm_handler.py +++ b/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,