github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-08-12 14:23:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

community: support usage_metadata for litellm (#30625)

Browse Source 
Support "usage_metadata" for LiteLLM. 

If no one reviews your PR within a few days, please @-mention one of
baskaryan, eyurtsev, ccurme, vbarda, hwchase17.
This commit is contained in:
Alejandro Rodríguez 2025-04-02 19:45:15 -04:00 committed by GitHub
parent 01d0cfe450
commit 884125e129
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 19 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
libs/community/langchain_community/chat_models/litellm.py
View File

@ -48,6 +48,7 @@ from langchain_core.messages import (
ToolCallChunk,
ToolMessage,
)
from langchain_core.messages.ai import UsageMetadata
from langchain_core.outputs import (
ChatGeneration,
ChatGenerationChunk,
@ -410,14 +411,19 @@ class ChatLiteLLM(BaseChatModel):
def _create_chat_result(self, response: Mapping[str, Any]) -> ChatResult:
generations = []
token_usage = response.get("usage", {})
for res in response["choices"]:
message = _convert_dict_to_message(res["message"])
if isinstance(message, AIMessage):
message.response_metadata = {
"model_name": self.model_name or self.model
}
message.usage_metadata = _create_usage_metadata(token_usage)
gen = ChatGeneration(
message=message,
generation_info=dict(finish_reason=res.get("finish_reason")),
)
generations.append(gen)
token_usage = response.get("usage", {})
set_model_value = self.model
if self.model_name is not None:
set_model_value = self.model_name
@ -585,3 +591,13 @@ class ChatLiteLLM(BaseChatModel):
@property
def _llm_type(self) -> str:
return "litellm-chat"
def _create_usage_metadata(token_usage: Mapping[str, Any]) -> UsageMetadata:
input_tokens = token_usage.get("prompt_tokens", 0)
output_tokens = token_usage.get("completion_tokens", 0)
return UsageMetadata(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
)

4
libs/community/tests/integration_tests/chat_models/test_litellm_standard.py
View File

@ -19,5 +19,5 @@ class TestLiteLLMStandard(ChatModelIntegrationTests):
return {"model": "ollama/mistral"}
@pytest.mark.xfail(reason="Not yet implemented.")
def test_usage_metadata(self, model: BaseChatModel) -> None:
super().test_usage_metadata(model)
def test_usage_metadata_streaming(self, model: BaseChatModel) -> None:
super().test_usage_metadata_streaming(model)