Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

genai: fix usage metadata #545

Merged
merged 3 commits into from
Oct 11, 2024
Merged

genai: fix usage metadata #545

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
75afc5a
genai: fix usage metadata
efriis Oct 8, 2024
9adb117
x
efriis Oct 8, 2024
8120d57
Merge branch 'main' into erick/genai-fix-usage-metadata
lkuligin Oct 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 55 additions & 5 deletions libs/genai/langchain_google_genai/chat_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -568,10 +568,17 @@ def _parse_response_candidate(
def _response_to_result(
response: GenerateContentResponse,
stream: bool = False,
prev_usage: Optional[UsageMetadata] = None,
) -> ChatResult:
"""Converts a PaLM API response into a LangChain ChatResult."""
llm_output = {"prompt_feedback": proto.Message.to_dict(response.prompt_feedback)}

# previous usage metadata needs to be subtracted because gemini api returns
# already-accumulated token counts with each chunk
prev_input_tokens = prev_usage["input_tokens"] if prev_usage else 0
prev_output_tokens = prev_usage["output_tokens"] if prev_usage else 0
prev_total_tokens = prev_usage["total_tokens"] if prev_usage else 0

# Get usage metadata
try:
input_tokens = response.usage_metadata.prompt_token_count
Expand All @@ -580,9 +587,9 @@ def _response_to_result(
cache_read_tokens = response.usage_metadata.cached_content_token_count
if input_tokens + output_tokens + cache_read_tokens + total_tokens > 0:
lc_usage = UsageMetadata(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=total_tokens,
input_tokens=input_tokens - prev_input_tokens,
output_tokens=output_tokens - prev_output_tokens,
total_tokens=total_tokens - prev_total_tokens,
input_token_details={"cache_read": cache_read_tokens},
)
else:
Expand Down Expand Up @@ -1086,9 +1093,31 @@ def _stream(
**kwargs,
metadata=self.default_metadata,
)

prev_usage_metadata: UsageMetadata | None = None
for chunk in response:
_chat_result = _response_to_result(chunk, stream=True)
_chat_result = _response_to_result(
chunk, stream=True, prev_usage=prev_usage_metadata
)
gen = cast(ChatGenerationChunk, _chat_result.generations[0])
message = cast(AIMessageChunk, gen.message)

curr_usage_metadata: UsageMetadata | dict[str, int] = (
message.usage_metadata or {}
)

prev_usage_metadata = (
message.usage_metadata
if prev_usage_metadata is None
else UsageMetadata(
input_tokens=prev_usage_metadata.get("input_tokens", 0)
+ curr_usage_metadata.get("input_tokens", 0),
output_tokens=prev_usage_metadata.get("output_tokens", 0)
+ curr_usage_metadata.get("output_tokens", 0),
total_tokens=prev_usage_metadata.get("total_tokens", 0)
+ curr_usage_metadata.get("total_tokens", 0),
)
)

if run_manager:
run_manager.on_llm_new_token(gen.text)
Expand Down Expand Up @@ -1134,14 +1163,35 @@ async def _astream(
generation_config=generation_config,
cached_content=cached_content or self.cached_content,
)
prev_usage_metadata: UsageMetadata | None = None
async for chunk in await _achat_with_retry(
request=request,
generation_method=self.async_client.stream_generate_content,
**kwargs,
metadata=self.default_metadata,
):
_chat_result = _response_to_result(chunk, stream=True)
_chat_result = _response_to_result(
chunk, stream=True, prev_usage=prev_usage_metadata
)
gen = cast(ChatGenerationChunk, _chat_result.generations[0])
message = cast(AIMessageChunk, gen.message)

curr_usage_metadata: UsageMetadata | dict[str, int] = (
message.usage_metadata or {}
)

prev_usage_metadata = (
message.usage_metadata
if prev_usage_metadata is None
else UsageMetadata(
input_tokens=prev_usage_metadata.get("input_tokens", 0)
+ curr_usage_metadata.get("input_tokens", 0),
output_tokens=prev_usage_metadata.get("output_tokens", 0)
+ curr_usage_metadata.get("output_tokens", 0),
total_tokens=prev_usage_metadata.get("total_tokens", 0)
+ curr_usage_metadata.get("total_tokens", 0),
)
)

if run_manager:
await run_manager.on_llm_new_token(gen.text)
Expand Down
Loading