Skip to content

Commit a103ae1

Browse files
fix(closes OPEN-11695): strip Gemini models/ prefix and emit priced usageDetails so LangChain cost estimation stops returning $0
Builds on the v1 callback modernization (OPEN-11315), which resolves the provider from `metadata["ls_provider"]` (ChatGoogleGenerativeAI -> "Google"). Remaining gaps that still produced $0 / empty cost fields: - Provider anchor: also map the `chat-google-generative-ai` _type to "Google" in the legacy _type map, so the exact reported string still resolves when `metadata["ls_provider"]` is absent (older langchain-google-genai, or callers that don't forward it) -- not just via the ls_provider path. - Model: strip the Gemini Developer API "models/" prefix (e.g. "models/gemini-3.5-flash" -> "gemini-3.5-flash"), which the cost table (bare names) otherwise misses -> $0. Runs before the LiteLLM prefix handling ("models" is never a provider, so it's safe). - costDetails/usageDetails: the handler surfaced granular token_details only as informational step metadata, never as a priced column. Emit a per-category `usageDetails` map so the backend prices it into `costDetails`. LangChain reports cache/audio categories as overlapping subsets of the input/output totals, so they are re-partitioned to be non-overlapping and mapped to the cost table's keys (cache_read -> cached_tokens, cache_creation -> cache_creation_tokens, audio -> audio_input_tokens/audio_output_tokens); reasoning stays folded into output_tokens. ChatCompletionStep gains an optional usage_details field, serialized only when set. Verified end-to-end against the live pipeline with the customer's inputs (27131 prompt / 17739 completion): provider=Google, model=gemini-3.5-flash, cost=0.2003475, costDetails populated; and with 10000 cached input tokens the partition prices cached tokens at the cheaper rate (cost=0.1868475). Cross- provider audit confirms the cost table uses these category keys uniformly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018t3JHc1kK6pPMdVBVwMbTs
1 parent 76a798f commit a103ae1

3 files changed

Lines changed: 219 additions & 2 deletions

File tree

src/openlayer/lib/integrations/langchain_callback.py

Lines changed: 90 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,11 @@
3030
"chat-ollama": "Ollama",
3131
"vertexai": "Google",
3232
"amazon_bedrock_converse_chat": "Bedrock",
33+
# ChatGoogleGenerativeAI's _llm_type. Anchored here as well as via
34+
# ls_provider ("google_genai") so the exact reported string (OPEN-11695)
35+
# still resolves when metadata["ls_provider"] is absent (older
36+
# langchain-google-genai, or callers that don't forward it).
37+
"chat-google-generative-ai": "Google",
3338
}
3439

3540
# LangChain v1 injects a standardized ``metadata["ls_provider"]`` (e.g.
@@ -492,6 +497,14 @@ def _extract_model_info(
492497
or serialized.get("name")
493498
)
494499

500+
# Strip the Google Gemini Developer API "models/" prefix
501+
# (e.g. "models/gemini-3.5-flash" -> "gemini-3.5-flash"). The cost table
502+
# stores bare Gemini names and no provider is named "models", so this is
503+
# safe. Runs before the LiteLLM prefix handling below so the remaining
504+
# name has no stray leading segment.
505+
if model and model.startswith("models/"):
506+
model = model[len("models/") :]
507+
495508
# Handle LiteLLM model prefix (e.g. "gemini/gemini-2.5-flash"):
496509
# extract the actual provider and strip the prefix from the model name.
497510
if model and "/" in model:
@@ -581,6 +594,59 @@ def _extract_token_info(self, response: "langchain_schema.LLMResult") -> Dict[st
581594
result["token_details"] = token_details
582595
return result
583596

597+
@staticmethod
598+
def _build_usage_details(
599+
prompt_tokens: int,
600+
completion_tokens: int,
601+
input_token_details: Optional[Dict[str, int]] = None,
602+
output_token_details: Optional[Dict[str, int]] = None,
603+
) -> Optional[Dict[str, int]]:
604+
"""Build the per-category token map the cost backend prices by exact key.
605+
606+
The backend prices a **non-overlapping** partition: it sums a price for
607+
each ``usageDetails`` key it recognizes. LangChain, however, reports the
608+
granular categories (``cache_read``, ``cache_creation``, ``audio``) as
609+
**subsets** of the ``input_tokens`` / ``output_tokens`` totals. So we
610+
break each granular category out under the key the cost table uses and
611+
subtract it from the input/output base, keeping the partition
612+
non-overlapping and its sum equal to the total token count.
613+
614+
Key mapping (LangChain -> Openlayer cost table):
615+
input ``cache_read`` -> ``cached_tokens``
616+
input ``cache_creation`` -> ``cache_creation_tokens``
617+
input ``audio`` -> ``audio_input_tokens``
618+
output ``audio`` -> ``audio_output_tokens``
619+
620+
``reasoning`` is intentionally left folded into ``output_tokens`` -- it
621+
is billed at the output rate and the cost table has no separate price for
622+
it. Zero-valued categories are omitted; returns ``None`` when there are no
623+
tokens so the ``usageDetails`` column is omitted rather than emitted empty.
624+
"""
625+
input_token_details = input_token_details or {}
626+
output_token_details = output_token_details or {}
627+
628+
cache_read = int(input_token_details.get("cache_read", 0) or 0)
629+
cache_creation = int(input_token_details.get("cache_creation", 0) or 0)
630+
audio_input = int(input_token_details.get("audio", 0) or 0)
631+
audio_output = int(output_token_details.get("audio", 0) or 0)
632+
633+
# Base = total minus the granular categories broken out below.
634+
input_base = prompt_tokens - cache_read - cache_creation - audio_input
635+
output_base = completion_tokens - audio_output
636+
637+
usage_details: Dict[str, int] = {}
638+
for key, value in (
639+
("input_tokens", input_base),
640+
("output_tokens", output_base),
641+
("cached_tokens", cache_read),
642+
("cache_creation_tokens", cache_creation),
643+
("audio_input_tokens", audio_input),
644+
("audio_output_tokens", audio_output),
645+
):
646+
if value and value > 0:
647+
usage_details[key] = value
648+
return usage_details or None
649+
584650
def _extract_output(self, response: "langchain_schema.LLMResult") -> str:
585651
"""Extract output text from LLM response.
586652
@@ -722,6 +788,18 @@ def _handle_llm_end(
722788
if token_details:
723789
step.metadata = {**step.metadata, "token_details": token_details}
724790

791+
# Also emit a priced, non-overlapping per-category usageDetails map so
792+
# the backend can populate costDetails (the token_details above are only
793+
# surfaced as informational metadata, not priced).
794+
usage_details = self._build_usage_details(
795+
token_info.get("prompt_tokens", 0),
796+
token_info.get("completion_tokens", 0),
797+
(token_details or {}).get("input_token_details"),
798+
(token_details or {}).get("output_token_details"),
799+
)
800+
if usage_details:
801+
token_info["usage_details"] = usage_details
802+
725803
self._end_step(
726804
run_id=run_id,
727805
parent_run_id=parent_run_id,
@@ -1104,11 +1182,21 @@ def _handle_llm_new_token(self, token: str, **kwargs: Any) -> Any:
11041182
run_id = kwargs.get("run_id")
11051183
if run_id and run_id in self.steps:
11061184
# Convert usage to the expected format like _extract_token_info does
1185+
prompt_tokens = usage.get("input_tokens", 0)
1186+
completion_tokens = usage.get("output_tokens", 0)
11071187
token_info = {
1108-
"prompt_tokens": usage.get("input_tokens", 0),
1109-
"completion_tokens": usage.get("output_tokens", 0),
1188+
"prompt_tokens": prompt_tokens,
1189+
"completion_tokens": completion_tokens,
11101190
"tokens": usage.get("total_tokens", 0),
11111191
}
1192+
usage_details = self._build_usage_details(
1193+
prompt_tokens,
1194+
completion_tokens,
1195+
usage.get("input_token_details"),
1196+
usage.get("output_token_details"),
1197+
)
1198+
if usage_details:
1199+
token_info["usage_details"] = usage_details
11121200

11131201
# Update the step with token usage information
11141202
step = self.steps[run_id]

src/openlayer/lib/tracing/steps.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,11 @@ def __init__(
187187
self.model: str = None
188188
self.model_parameters: Dict[str, Any] = None
189189
self.raw_output: str = None
190+
# Optional per-category token map (e.g. {"input_tokens": ...,
191+
# "output_tokens": ...}). When set, the backend prices each category by
192+
# exact key match to populate ``costDetails``; when unset it is omitted
193+
# so integrations that only report scalar tokens are unaffected.
194+
self.usage_details: Optional[Dict[str, int]] = None
190195

191196
def to_dict(self) -> Dict[str, Any]:
192197
"""Dictionary representation of the ChatCompletionStep."""
@@ -203,6 +208,8 @@ def to_dict(self) -> Dict[str, Any]:
203208
"rawOutput": self.raw_output,
204209
}
205210
)
211+
if self.usage_details:
212+
step_dict["usageDetails"] = self.usage_details
206213
return step_dict
207214

208215

tests/lib/integrations/test_langchain_callback.py

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -583,3 +583,125 @@ def test_have_langchain_true_and_schema_from_core(self) -> None:
583583
assert lc.HAVE_LANGCHAIN is True
584584
# The schema alias must resolve to langchain_core, not the legacy path.
585585
assert lc.langchain_schema.__name__.startswith("langchain_core")
586+
587+
588+
# --------------------------------------------------------------------------- #
589+
# OPEN-11695: model name normalization + priced usageDetails/costDetails
590+
# --------------------------------------------------------------------------- #
591+
class TestModelNameNormalization:
592+
def test_strips_gemini_models_prefix(self) -> None:
593+
# Customer case: ChatGoogleGenerativeAI reports ls_provider=google_genai
594+
# and a model with the Gemini Developer API "models/" prefix. Without the
595+
# strip the cost table lookup misses -> $0.
596+
handler = OpenlayerHandler()
597+
info = handler._extract_model_info(
598+
serialized={},
599+
invocation_params={
600+
"_type": "chat-google-generative-ai",
601+
"model": "models/gemini-3.5-flash",
602+
},
603+
metadata={
604+
"ls_provider": "google_genai",
605+
"ls_model_name": "models/gemini-3.5-flash",
606+
},
607+
)
608+
assert info["provider"] == "Google"
609+
assert info["model"] == "gemini-3.5-flash"
610+
611+
def test_models_prefix_strip_independent_of_provider(self) -> None:
612+
handler = OpenlayerHandler()
613+
info = handler._extract_model_info(
614+
serialized={},
615+
invocation_params={"model": "models/gemini-2.5-flash"},
616+
metadata={},
617+
)
618+
assert info["model"] == "gemini-2.5-flash"
619+
620+
def test_type_anchor_maps_when_ls_provider_absent(self) -> None:
621+
# OPEN-11695: the exact reported _type string must resolve to Google even
622+
# without metadata["ls_provider"] (older langchain-google-genai etc.).
623+
handler = OpenlayerHandler()
624+
info = handler._extract_model_info(
625+
serialized={},
626+
invocation_params={
627+
"_type": "chat-google-generative-ai",
628+
"model": "models/gemini-3.5-flash",
629+
},
630+
metadata={},
631+
)
632+
assert info["provider"] == "Google"
633+
assert info["model"] == "gemini-3.5-flash"
634+
635+
636+
class TestUsageDetailsPricing:
637+
def test_scalar_partition_when_no_details(self) -> None:
638+
assert OpenlayerHandler._build_usage_details(100, 50) == {
639+
"input_tokens": 100,
640+
"output_tokens": 50,
641+
}
642+
643+
def test_none_when_no_tokens(self) -> None:
644+
assert OpenlayerHandler._build_usage_details(0, 0) is None
645+
646+
def test_cached_tokens_partitioned_non_overlapping(self) -> None:
647+
# LangChain reports input_tokens INCLUSIVE of cache_read/cache_creation;
648+
# the backend prices a non-overlapping partition under its own keys.
649+
details = OpenlayerHandler._build_usage_details(
650+
350, 100, {"cache_read": 100, "cache_creation": 200}, None
651+
)
652+
assert details == {
653+
"input_tokens": 50, # 350 - 100 - 200
654+
"output_tokens": 100,
655+
"cached_tokens": 100,
656+
"cache_creation_tokens": 200,
657+
}
658+
assert sum(details.values()) == 350 + 100 # partition conserves total
659+
660+
def test_audio_broken_out_both_directions(self) -> None:
661+
details = OpenlayerHandler._build_usage_details(
662+
300, 120, {"audio": 30}, {"audio": 20}
663+
)
664+
assert details == {
665+
"input_tokens": 270,
666+
"output_tokens": 100,
667+
"audio_input_tokens": 30,
668+
"audio_output_tokens": 20,
669+
}
670+
671+
def test_reasoning_stays_folded_into_output(self) -> None:
672+
# Reasoning is billed at the output rate; must not be split out or
673+
# subtracted from output_tokens.
674+
details = OpenlayerHandler._build_usage_details(100, 240, None, {"reasoning": 200})
675+
assert details == {"input_tokens": 100, "output_tokens": 240}
676+
677+
def test_usage_details_set_on_step_via_callbacks(self) -> None:
678+
handler = OpenlayerHandler()
679+
run_id = uuid.uuid4()
680+
handler.on_chat_model_start(
681+
serialized={"name": "gemini"},
682+
messages=[[HumanMessage(content="hi")]],
683+
run_id=run_id,
684+
invocation_params={"model": "models/gemini-3.5-flash"},
685+
metadata={"ls_provider": "google_genai"},
686+
)
687+
step = handler.steps[run_id]
688+
message = _ai_message_with_usage(
689+
input_tokens=27131,
690+
output_tokens=17739,
691+
total_tokens=44870,
692+
input_token_details={"cache_read": 10000},
693+
)
694+
handler.on_llm_end(
695+
LLMResult(generations=[[ChatGeneration(message=message)]]), run_id=run_id
696+
)
697+
# Priced, non-overlapping partition lands on the step (and serializes as
698+
# the "usageDetails" column the backend prices into costDetails).
699+
assert step.usage_details == {
700+
"input_tokens": 17131,
701+
"output_tokens": 17739,
702+
"cached_tokens": 10000,
703+
}
704+
assert step.to_dict()["usageDetails"] == step.usage_details
705+
706+
def test_step_omits_usage_details_when_unset(self) -> None:
707+
assert "usageDetails" not in steps.ChatCompletionStep(name="x").to_dict()

0 commit comments

Comments
 (0)