Add token usage to gemini (streaming), litellm and anthropic

Also included a token_usage sample that showcases the token usage of subagents with different models under a parent agent.

PiperOrigin-RevId: 759347015
This commit is contained in:
Selcuk Gun
2025-05-15 16:22:36 -07:00
committed by Copybara-Service
parent 4d5760917d
commit 509db3f9fb
8 changed files with 515 additions and 45 deletions
+2 -1
View File
@@ -121,6 +121,7 @@ class Gemini(BaseLlm):
content=types.ModelContent(
parts=[types.Part.from_text(text=text)],
),
usage_metadata=llm_response.usage_metadata,
)
text = ''
yield llm_response
@@ -174,7 +175,7 @@ class Gemini(BaseLlm):
@cached_property
def _live_api_client(self) -> Client:
if self._api_backend == 'vertex':
#use beta version for vertex api
# use beta version for vertex api
api_version = 'v1beta1'
# use default api version for vertex
return Client(