[Opik-2168] Add cost tracking support for ChatBedrock and ChatBedrockConverse in langchain, expose bedrock as an officially supported provider (#2862)

* Draft

* Start using new usage extractors in OpikTracer

* Fix lint errors

* Fix lint errors

* Refactor usage extractors, add new mappers to LangChainUsage

* Simplified if statements in usage extractors

* Fix lint errors

* Add default values to langchain usage structure

* Change default values from int 0 to None

* Add map_to_groq_completions_usage

* Update langchain openai test to better reflect the expected usage structure

* Add docstrings to LLMUsageInfo model

* Expose bedrock provider in the LLMProvider enum as officially supported one

* Draft implementation

* Add streaming mode test

* Add fixture to ensure bedrock is configured, add more tests for async use case.

* Remove unnecessary comments

* Move boto3 import to ensure_aws_bedrock_configured fixture

* Rename tests

* Update env var name for aws region in langchain tests workflow file

* Propagate AWS creds from secrets to env vars

* assume correct role for bedrock

* skip tag session

* Update extractor logic, add support for ChatBedrockConverse

* Fix lint errors

* Update parameter settings for chat models in tests

---------

Co-authored-by: Alex Bezpalko <[email protected]>

Author: alexkuzmik

.github/workflows/lib-langchain-tests.yml

@@ -12,6 +12,9 @@ env:
   GOOGLE_CLOUD_PROJECT: opik-sdk-tests
   GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
   ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+  AWS_DEFAULT_REGION: us-east-1
+  AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+  AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
   OPIK_ENABLE_LITELLM_MODELS_MONITORING: False
   OPIK_SENTRY_ENABLE: False
 on:
@@ -52,7 +55,17 @@ jobs:
           cd ./tests
           pip install --no-cache-dir --disable-pip-version-check -r library_integration/langchain/requirements.txt
 
+      - name: change aws role
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          role-to-assume: ${{ vars.BEDROCK_ROLE }}
+          aws-region: us-east-1
+          role-chaining: true
+          role-skip-session-tagging: true
+
       - name: Run tests
         run: |
           cd ./tests/library_integration/langchain/
-          python -m pytest  -vv .
+          python -m pytest  -vv .

sdks/python/src/opik/integrations/bedrock/converse_decorator.py

@@ -2,6 +2,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 from typing_extensions import override
 
+import opik
 import opik.dict_utils as dict_utils
 import opik.llm_usage as llm_usage
 from opik.api_objects import span
@@ -63,7 +64,7 @@ def _end_span_inputs_preprocessor(
     ) -> arguments_helpers.EndSpanParameters:
         usage = output["usage"]
         usage_in_openai_format = llm_usage.try_build_opik_usage_or_log_error(
-            provider="_bedrock",
+            provider=opik.LLMProvider.BEDROCK,
             usage=usage,
             logger=LOGGER,
             error_message="Failed to log token usage from bedrock LLM call",

sdks/python/src/opik/integrations/langchain/provider_usage_extractors/bedrock_usage_extractor.py

@@ -0,0 +1,76 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Optional, List
+
+import opik
+from opik import llm_usage
+from . import provider_usage_extractor_protocol
+from . import langchain_run_helpers
+
+if TYPE_CHECKING:
+    pass
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BedrockUsageExtractor(
+    provider_usage_extractor_protocol.ProviderUsageExtractorProtocol
+):
+    PROVIDER = opik.LLMProvider.BEDROCK
+
+    def is_provider_run(self, run_dict: Dict[str, Any]) -> bool:
+        try:
+            if run_dict.get("serialized") is None:
+                return False
+
+            class_id: List[str] = run_dict.get("serialized", {}).get("id", [])
+            if len(class_id) == 0:
+                return False
+
+            class_name = class_id[-1]
+            is_bedrock = "ChatBedrock" in class_name
+
+            return is_bedrock
+
+        except Exception:
+            LOGGER.debug(
+                "Failed to check if Run instance is from Bedrock LLM, returning False.",
+                exc_info=True,
+            )
+            return False
+
+    def get_llm_usage_info(self, run_dict: Dict[str, Any]) -> llm_usage.LLMUsageInfo:
+        usage_dict = _try_get_token_usage(run_dict)
+        model = _try_get_model_name(run_dict)
+
+        return llm_usage.LLMUsageInfo(
+            provider=self.PROVIDER, model=model, usage=usage_dict
+        )
+
+
+def _try_get_token_usage(run_dict: Dict[str, Any]) -> Optional[llm_usage.OpikUsage]:
+    try:
+        langchain_usage = langchain_run_helpers.try_get_token_usage(run_dict)
+        bedrock_usage_dict = langchain_usage.map_to_bedrock_usage()
+
+        opik_usage = llm_usage.OpikUsage.from_bedrock_dict(bedrock_usage_dict)
+        return opik_usage
+    except Exception:
+        LOGGER.warning(
+            "Failed to extract token usage from presumably Bedrock LLM langchain run.",
+            exc_info=True,
+        )
+        return None
+
+
+def _try_get_model_name(run_dict: Dict[str, Any]) -> Optional[str]:
+    MODEL_NAME_KEY = "model_id"
+
+    model = run_dict.get("serialized", {}).get("kwargs", {}).get(MODEL_NAME_KEY, None)
+
+    if model is None:
+        LOGGER.error(
+            "Failed to extract model name from presumably Bedrock LLM langchain Run object: %s",
+            run_dict,
+        )
+
+    return model

sdks/python/src/opik/integrations/langchain/provider_usage_extractors/langchain_run_helpers/langchain_usage.py

@@ -108,6 +108,13 @@ def map_to_anthropic_usage(self) -> Dict[str, Any]:
 
         return anthropic_usage
 
+    def map_to_bedrock_usage(self) -> Dict[str, Any]:
+        bedrock_usage: Dict[str, Any] = {
+            "inputTokens": self.input_tokens,
+            "outputTokens": self.output_tokens,
+        }
+        return bedrock_usage
+
     def map_to_openai_completions_usage(self) -> Dict[str, Any]:
         openai_usage: Dict[str, Any] = {
             "prompt_tokens": self.input_tokens,

sdks/python/src/opik/integrations/langchain/provider_usage_extractors/usage_extractor.py

@@ -9,6 +9,7 @@
     vertexai_usage_extractor,
     groq_usage_extractor,
     anthropic_vertexai_usage_extractor,
+    bedrock_usage_extractor,
 )
 from . import provider_usage_extractor_protocol
 
@@ -23,6 +24,7 @@
     vertexai_usage_extractor.VertexAIUsageExtractor(),
     groq_usage_extractor.GroqUsageExtractor(),
     anthropic_vertexai_usage_extractor.AnthropicVertexAIUsageExtractor(),
+    bedrock_usage_extractor.BedrockUsageExtractor(),
 ]
 
 

sdks/python/src/opik/llm_usage/bedrock_usage.py

@@ -3,7 +3,7 @@
 
 
 class BedrockUsage(base_original_provider_usage.BaseOriginalProviderUsage):
-    """Anthropic calls token usage data. Updated 11.03.2025"""
+    """Bedrock calls token usage data. Updated 11.03.2025"""
 
     inputTokens: int
     """The number of input tokens which were used."""

sdks/python/src/opik/llm_usage/opik_usage_factory.py

@@ -5,9 +5,7 @@
 from . import opik_usage
 
 
-# One provider can have multiple formats of usage dicts, so it could be many build functions
-# if provider's name specified as string and not as LLMProvider enum value -
-# it means that we do not support cost tracking for this provider (but support usage info)
+# One provider can have multiple formats of usage dicts, so it can have more than 1 build function
 _PROVIDER_TO_OPIK_USAGE_BUILDERS: Dict[
     Union[str, LLMProvider],
     List[Callable[[Dict[str, Any]], opik_usage.OpikUsage]],
@@ -19,7 +17,7 @@
     LLMProvider.GOOGLE_VERTEXAI: [opik_usage.OpikUsage.from_google_dict],
     LLMProvider.GOOGLE_AI: [opik_usage.OpikUsage.from_google_dict],
     LLMProvider.ANTHROPIC: [opik_usage.OpikUsage.from_anthropic_dict],
-    "_bedrock": [opik_usage.OpikUsage.from_bedrock_dict],
+    LLMProvider.BEDROCK: [opik_usage.OpikUsage.from_bedrock_dict],
 }
 
 

sdks/python/src/opik/types.py

@@ -34,6 +34,9 @@ class LLMProvider(str, enum.Enum):
     GROQ = "groq"
     """Used for models hosted by Groq. https://groq.com"""
 
+    BEDROCK = "bedrock"
+    """Used for models hosted by AWS Bedrock. https://aws.amazon.com/bedrock"""
+
     @classmethod
     def has_value(cls, value: str) -> bool:
         return value in [enum_item.value for enum_item in cls]

sdks/python/tests/conftest.py

@@ -190,14 +190,14 @@ def temp_file_15mb():
         yield f
 
 
-@pytest.fixture()
+@pytest.fixture(scope="session")
 def ensure_openai_configured():
     # don't use assertion here to prevent printing os.environ with all env variables
     if not ("OPENAI_API_KEY" in os.environ and "OPENAI_ORG_ID" in os.environ):
         raise Exception("OpenAI not configured!")
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def ensure_google_project_and_location_configured():
     if not (
         "GOOGLE_CLOUD_PROJECT" in os.environ and "GOOGLE_CLOUD_LOCATION" in os.environ
@@ -207,15 +207,15 @@ def ensure_google_project_and_location_configured():
         )
 
 
-@pytest.fixture()
+@pytest.fixture(scope="session")
 def ensure_anthropic_configured():
     # don't use assertion here to prevent printing os.environ with all env variables
 
     if "ANTHROPIC_API_KEY" not in os.environ:
         raise Exception("Anthropic not configured!")
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def ensure_vertexai_configured(ensure_google_project_and_location_configured):
     GOOGLE_APPLICATION_CREDENTIALS_PATH = "gcp_credentials.json"
 
@@ -246,8 +246,23 @@ def ensure_vertexai_configured(ensure_google_project_and_location_configured):
             os.remove(GOOGLE_APPLICATION_CREDENTIALS_PATH)
 
 
-@pytest.fixture()
+@pytest.fixture(scope="session")
 def ensure_google_api_configured():
     GOOGLE_API_KEY = "GOOGLE_API_KEY"
     if GOOGLE_API_KEY not in os.environ:
         raise Exception(f"{GOOGLE_API_KEY} env var must be set")
+
+
+@pytest.fixture(scope="session")
+def ensure_aws_bedrock_configured():
+    import boto3
+
+    session = boto3.Session()
+
+    bedrock_client = session.client(service_name="bedrock")
+    try:
+        available_models = bedrock_client.list_foundation_models()["modelSummaries"]
+        if not available_models:
+            raise Exception("AWS Bedrock not configured! No models available")
+    except Exception as e:
+        raise Exception(f"AWS Bedrock not configured! {e}")

sdks/python/tests/library_integration/langchain/constants.py

@@ -27,3 +27,13 @@
     "original_usage.prompt_tokens_details.audio_tokens": ANY_BUT_NONE,
     "original_usage.prompt_tokens_details.cached_tokens": ANY_BUT_NONE,
 }
+
+BEDROCK_MODEL_FOR_TESTS = "us.anthropic.claude-sonnet-4-20250514-v1:0"
+
+EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT = {
+    "prompt_tokens": ANY_BUT_NONE,
+    "completion_tokens": ANY_BUT_NONE,
+    "total_tokens": ANY_BUT_NONE,
+    "original_usage.inputTokens": ANY_BUT_NONE,
+    "original_usage.outputTokens": ANY_BUT_NONE,
+}