fix(OPEN-11695): recognize LangChain Gemini provider/model and populate cost + usageDetails

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]

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
 
 

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()