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,