From fc3bfd3f788707c10b1792c3548c12f5e23e20e2 Mon Sep 17 00:00:00 2001 From: Vinicius Mello Date: Mon, 13 Jul 2026 11:48:04 -0300 Subject: [PATCH] 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) Claude-Session: https://claude.ai/code/session_018t3JHc1kK6pPMdVBVwMbTs --- .../lib/integrations/langchain_callback.py | 92 +++++++++++++- src/openlayer/lib/tracing/steps.py | 13 +- .../integrations/test_langchain_callback.py | 118 ++++++++++++++++++ 3 files changed, 216 insertions(+), 7 deletions(-) diff --git a/src/openlayer/lib/integrations/langchain_callback.py b/src/openlayer/lib/integrations/langchain_callback.py index 9d995570..d24c686a 100644 --- a/src/openlayer/lib/integrations/langchain_callback.py +++ b/src/openlayer/lib/integrations/langchain_callback.py @@ -30,6 +30,11 @@ "chat-ollama": "Ollama", "vertexai": "Google", "amazon_bedrock_converse_chat": "Bedrock", + # ChatGoogleGenerativeAI's _llm_type. Anchored here as well as via + # ls_provider ("google_genai") so the exact reported string (OPEN-11695) + # still resolves when metadata["ls_provider"] is absent (older + # langchain-google-genai, or callers that don't forward it). + "chat-google-generative-ai": "Google", } # LangChain v1 injects a standardized ``metadata["ls_provider"]`` (e.g. @@ -492,6 +497,14 @@ def _extract_model_info( or serialized.get("name") ) + # Strip the Google Gemini Developer API "models/" prefix + # (e.g. "models/gemini-3.5-flash" -> "gemini-3.5-flash"). The cost table + # stores bare Gemini names and no provider is named "models", so this is + # safe. Runs before the LiteLLM prefix handling below so the remaining + # name has no stray leading segment. + if model and model.startswith("models/"): + model = model[len("models/") :] + # Handle LiteLLM model prefix (e.g. "gemini/gemini-2.5-flash"): # extract the actual provider and strip the prefix from the model name. if model and "/" in model: @@ -581,6 +594,59 @@ def _extract_token_info(self, response: "langchain_schema.LLMResult") -> Dict[st result["token_details"] = token_details return result + @staticmethod + def _build_usage_details( + prompt_tokens: int, + completion_tokens: int, + input_token_details: Optional[Dict[str, int]] = None, + output_token_details: Optional[Dict[str, int]] = None, + ) -> Optional[Dict[str, int]]: + """Build the per-category token map the cost backend prices by exact key. + + The backend prices a **non-overlapping** partition: it sums a price for + each ``usageDetails`` key it recognizes. LangChain, however, reports the + granular categories (``cache_read``, ``cache_creation``, ``audio``) as + **subsets** of the ``input_tokens`` / ``output_tokens`` totals. So we + break each granular category out under the key the cost table uses and + subtract it from the input/output base, keeping the partition + non-overlapping and its sum equal to the total token count. + + Key mapping (LangChain -> Openlayer cost table): + input ``cache_read`` -> ``cached_tokens`` + input ``cache_creation`` -> ``cache_creation_tokens`` + input ``audio`` -> ``audio_input_tokens`` + output ``audio`` -> ``audio_output_tokens`` + + ``reasoning`` is intentionally left folded into ``output_tokens`` -- it + is billed at the output rate and the cost table has no separate price for + it. Zero-valued categories are omitted; returns ``None`` when there are no + tokens so the ``usageDetails`` column is omitted rather than emitted empty. + """ + input_token_details = input_token_details or {} + output_token_details = output_token_details or {} + + cache_read = int(input_token_details.get("cache_read", 0) or 0) + cache_creation = int(input_token_details.get("cache_creation", 0) or 0) + audio_input = int(input_token_details.get("audio", 0) or 0) + audio_output = int(output_token_details.get("audio", 0) or 0) + + # Base = total minus the granular categories broken out below. + input_base = prompt_tokens - cache_read - cache_creation - audio_input + output_base = completion_tokens - audio_output + + usage_details: Dict[str, int] = {} + for key, value in ( + ("input_tokens", input_base), + ("output_tokens", output_base), + ("cached_tokens", cache_read), + ("cache_creation_tokens", cache_creation), + ("audio_input_tokens", audio_input), + ("audio_output_tokens", audio_output), + ): + if value and value > 0: + usage_details[key] = value + return usage_details or None + def _extract_output(self, response: "langchain_schema.LLMResult") -> str: """Extract output text from LLM response. @@ -722,6 +788,18 @@ def _handle_llm_end( if token_details: step.metadata = {**step.metadata, "token_details": token_details} + # Also emit a priced, non-overlapping per-category usageDetails map so + # the backend can populate costDetails (the token_details above are only + # surfaced as informational metadata, not priced). + usage_details = self._build_usage_details( + token_info.get("prompt_tokens", 0), + token_info.get("completion_tokens", 0), + (token_details or {}).get("input_token_details"), + (token_details or {}).get("output_token_details"), + ) + if usage_details: + token_info["usage_details"] = usage_details + self._end_step( run_id=run_id, parent_run_id=parent_run_id, @@ -1104,11 +1182,21 @@ def _handle_llm_new_token(self, token: str, **kwargs: Any) -> Any: run_id = kwargs.get("run_id") if run_id and run_id in self.steps: # Convert usage to the expected format like _extract_token_info does + prompt_tokens = usage.get("input_tokens", 0) + completion_tokens = usage.get("output_tokens", 0) token_info = { - "prompt_tokens": usage.get("input_tokens", 0), - "completion_tokens": usage.get("output_tokens", 0), + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, "tokens": usage.get("total_tokens", 0), } + usage_details = self._build_usage_details( + prompt_tokens, + completion_tokens, + usage.get("input_token_details"), + usage.get("output_token_details"), + ) + if usage_details: + token_info["usage_details"] = usage_details # Update the step with token usage information step = self.steps[run_id] diff --git a/src/openlayer/lib/tracing/steps.py b/src/openlayer/lib/tracing/steps.py index 0d68d7f3..74a7229b 100644 --- a/src/openlayer/lib/tracing/steps.py +++ b/src/openlayer/lib/tracing/steps.py @@ -141,11 +141,7 @@ def to_dict(self) -> Dict[str, Any]: # Include valid attachments only (filter out ones with no data/reference) if self.attachments: - valid_attachments = [ - attachment.to_dict() - for attachment in self.attachments - if attachment.is_valid() - ] + valid_attachments = [attachment.to_dict() for attachment in self.attachments if attachment.is_valid()] if valid_attachments: result["attachments"] = valid_attachments @@ -187,6 +183,11 @@ def __init__( self.model: str = None self.model_parameters: Dict[str, Any] = None self.raw_output: str = None + # Optional per-category token map (e.g. {"input_tokens": ..., + # "output_tokens": ...}). When set, the backend prices each category by + # exact key match to populate ``costDetails``; when unset it is omitted + # so integrations that only report scalar tokens are unaffected. + self.usage_details: Optional[Dict[str, int]] = None def to_dict(self) -> Dict[str, Any]: """Dictionary representation of the ChatCompletionStep.""" @@ -203,6 +204,8 @@ def to_dict(self) -> Dict[str, Any]: "rawOutput": self.raw_output, } ) + if self.usage_details: + step_dict["usageDetails"] = self.usage_details return step_dict diff --git a/tests/lib/integrations/test_langchain_callback.py b/tests/lib/integrations/test_langchain_callback.py index 765c444a..8ce302f7 100644 --- a/tests/lib/integrations/test_langchain_callback.py +++ b/tests/lib/integrations/test_langchain_callback.py @@ -583,3 +583,121 @@ def test_have_langchain_true_and_schema_from_core(self) -> None: assert lc.HAVE_LANGCHAIN is True # The schema alias must resolve to langchain_core, not the legacy path. assert lc.langchain_schema.__name__.startswith("langchain_core") + + +# --------------------------------------------------------------------------- # +# OPEN-11695: model name normalization + priced usageDetails/costDetails +# --------------------------------------------------------------------------- # +class TestModelNameNormalization: + def test_strips_gemini_models_prefix(self) -> None: + # Customer case: ChatGoogleGenerativeAI reports ls_provider=google_genai + # and a model with the Gemini Developer API "models/" prefix. Without the + # strip the cost table lookup misses -> $0. + handler = OpenlayerHandler() + info = handler._extract_model_info( + serialized={}, + invocation_params={ + "_type": "chat-google-generative-ai", + "model": "models/gemini-3.5-flash", + }, + metadata={ + "ls_provider": "google_genai", + "ls_model_name": "models/gemini-3.5-flash", + }, + ) + assert info["provider"] == "Google" + assert info["model"] == "gemini-3.5-flash" + + def test_models_prefix_strip_independent_of_provider(self) -> None: + handler = OpenlayerHandler() + info = handler._extract_model_info( + serialized={}, + invocation_params={"model": "models/gemini-2.5-flash"}, + metadata={}, + ) + assert info["model"] == "gemini-2.5-flash" + + def test_type_anchor_maps_when_ls_provider_absent(self) -> None: + # OPEN-11695: the exact reported _type string must resolve to Google even + # without metadata["ls_provider"] (older langchain-google-genai etc.). + handler = OpenlayerHandler() + info = handler._extract_model_info( + serialized={}, + invocation_params={ + "_type": "chat-google-generative-ai", + "model": "models/gemini-3.5-flash", + }, + metadata={}, + ) + assert info["provider"] == "Google" + assert info["model"] == "gemini-3.5-flash" + + +class TestUsageDetailsPricing: + def test_scalar_partition_when_no_details(self) -> None: + assert OpenlayerHandler._build_usage_details(100, 50) == { + "input_tokens": 100, + "output_tokens": 50, + } + + def test_none_when_no_tokens(self) -> None: + assert OpenlayerHandler._build_usage_details(0, 0) is None + + def test_cached_tokens_partitioned_non_overlapping(self) -> None: + # LangChain reports input_tokens INCLUSIVE of cache_read/cache_creation; + # the backend prices a non-overlapping partition under its own keys. + details = OpenlayerHandler._build_usage_details(350, 100, {"cache_read": 100, "cache_creation": 200}, None) + assert details is not None + assert details == { + "input_tokens": 50, # 350 - 100 - 200 + "output_tokens": 100, + "cached_tokens": 100, + "cache_creation_tokens": 200, + } + assert sum(details.values()) == 350 + 100 # partition conserves total + + def test_audio_broken_out_both_directions(self) -> None: + details = OpenlayerHandler._build_usage_details(300, 120, {"audio": 30}, {"audio": 20}) + assert details == { + "input_tokens": 270, + "output_tokens": 100, + "audio_input_tokens": 30, + "audio_output_tokens": 20, + } + + def test_reasoning_stays_folded_into_output(self) -> None: + # Reasoning is billed at the output rate; must not be split out or + # subtracted from output_tokens. + details = OpenlayerHandler._build_usage_details(100, 240, None, {"reasoning": 200}) + assert details == {"input_tokens": 100, "output_tokens": 240} + + def test_usage_details_set_on_step_via_callbacks(self) -> None: + handler = OpenlayerHandler() + run_id = uuid.uuid4() + handler.on_chat_model_start( + serialized={"name": "gemini"}, + messages=[[HumanMessage(content="hi")]], + run_id=run_id, + invocation_params={"model": "models/gemini-3.5-flash"}, + metadata={"ls_provider": "google_genai"}, + ) + step = handler.steps[run_id] + assert isinstance(step, steps.ChatCompletionStep) + message = _ai_message_with_usage( + input_tokens=27131, + output_tokens=17739, + total_tokens=44870, + input_token_details={"cache_read": 10000}, + ) + handler.on_llm_end(LLMResult(generations=[[ChatGeneration(message=message)]]), run_id=run_id) + # Priced, non-overlapping partition lands on the step (and serializes as + # the "usageDetails" column the backend prices into costDetails). + assert step.usage_details == { + "input_tokens": 17131, + "output_tokens": 17739, + "cached_tokens": 10000, + } + assert step.to_dict()["usageDetails"] == step.usage_details + + def test_step_omits_usage_details_when_unset(self) -> None: + assert "usageDetails" not in steps.ChatCompletionStep(name="x").to_dict()