fix(AIC-3382): support Bedrock configs in LangChain handlers

examples/streaming.py

@@ -31,16 +31,21 @@ async def run(key: str, user_input: str) -> None:
             registry=global_registry,
         ).stream(user_input, new_context())
 
+    wrote_chunk = False
     async for event in stream:
         if event["type"] == "chunk":
             sys.stdout.write(event.get("text", ""))
             sys.stdout.flush()
+            wrote_chunk = True
         else:
             # Final event — full response + normalised usage
+            if not wrote_chunk:
+                sys.stdout.write(str(event.get("response") or ""))
             sys.stdout.write("\n\n")
             print("Usage:", json_pretty(event.get("usage")))
             if event.get("judgeResults"):
                 print("Judge results:", json_pretty(event["judgeResults"]))
+            sys.stdout.write("\n")
 
 
 def json_pretty(obj: object) -> str:

packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py

@@ -147,12 +147,37 @@ def _build_initial_messages(
     return messages
 
 
+def _resolved_model_name(config: AiConfigRep, fallback_name: str = "") -> str:
+    """Bedrock ``model.region`` is an inference-profile prefix, prepended once."""
+    model = config.get("model") or {}
+    name = str(model.get("name") or fallback_name)
+    provider = str((config.get("provider") or {}).get("name") or "").lower()
+    if provider != "bedrock":
+        return name
+    prefix = str(model.get("region") or "")
+    if not prefix or name.startswith(f"{prefix}."):
+        return name
+    return f"{prefix}.{name}"
+
+
+def _config_for_model_call(config: AiConfigRep) -> AiConfigRep:
+    """Shallow copy with a resolved Bedrock model name. Does not mutate *config*."""
+    resolved = _resolved_model_name(config)
+    model = dict(config.get("model") or {})
+    if model.get("name") == resolved:
+        return config
+    return {**config, "model": {**model, "name": resolved}}
+
+
 def _model_constructor_kwargs(
     config: AiConfigRep, fallback_name: str
 ) -> dict[str, Any]:
     raw = (config.get("model") or {}).get("parameters")
     parameters = dict(raw) if isinstance(raw, dict) else {}
-    parameters["model"] = (config.get("model") or {}).get("name") or fallback_name
+    provider = str((config.get("provider") or {}).get("name") or "").lower()
+    if provider == "bedrock":
+        parameters.pop("tools", None)
+    parameters["model"] = _resolved_model_name(config, fallback_name)
     return parameters
 
 
@@ -176,15 +201,25 @@ def _make_default_chat_model(config: AiConfigRep) -> Any:
         return lc_anthropic.ChatAnthropic(
             **_model_constructor_kwargs(config, "claude-3-5-sonnet-20241022")
         )
+    if provider == "bedrock":
+        try:
+            lc_aws = importlib.import_module("langchain_aws")
+        except ImportError as exc:
+            raise ImportError(
+                "Using Bedrock models requires langchain-aws. "
+                "Install it with: pip install langchain-aws"
+            ) from exc
+        return lc_aws.ChatBedrockConverse(**_model_constructor_kwargs(config, ""))
     lc_openai = importlib.import_module("langchain_openai")
     return lc_openai.ChatOpenAI(**_model_constructor_kwargs(config, "gpt-4o"))
 
 
 async def _resolve_base_model(config: AiConfigRep, llm: Any) -> Any:
+    invocation = _config_for_model_call(config)
     if llm is None:
-        return _make_default_chat_model(config)
+        return _make_default_chat_model(invocation)
     if _is_model_factory(llm):
-        model = llm(config)
+        model = llm(invocation)
         if asyncio.iscoroutine(model):
             return await model
         return model

packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py

@@ -63,17 +63,13 @@ def model_name(config: AiConfigRep) -> str:
 
 
 def serving_provider(config: AiConfigRep) -> str:
-    """The provider that actually serves the model.
+    """The configured provider, lower-cased, for ``gen_ai.provider.name``.
 
     ``gen_ai.provider.name`` names who served the request, and its semconv enum has no
-    ``langchain`` member: LangChain is the framework, not the provider. This mirrors the choice
-    ``_make_default_chat_model`` makes, so the attribute agrees with the client that is really
-    used. It is a binary choice, not a passthrough of the configured name: the configured name
-    lower-cased if it equals ``anthropic``, otherwise ``openai``, no matter what else the config
-    names (Bedrock, Azure, Cohere, a typo, or nothing at all).
+    ``langchain`` member: LangChain is the framework, not the provider. Empty or missing
+    names fall back to ``openai``. ``gen_ai.system`` stays the literal ``langchain``.
     """
-    provider = ((config.get("provider") or {}).get("name") or "").lower()
-    return "anthropic" if provider == "anthropic" else "openai"
+    return str((config.get("provider") or {}).get("name") or "openai").lower()
 
 
 # ─── Span starts ─────────────────────────────────────────────────────────────

packages/langchain-agents/tests/test_handler.py

@@ -778,18 +778,27 @@ async def test_gen_ai_provider_name_is_anthropic_when_config_names_it(self) -> N
         assert rec.root.attributes["gen_ai.provider.name"] == "anthropic"
         assert rec.root.attributes["gen_ai.system"] == "langchain"
 
+    @pytest.mark.parametrize(
+        ("provider_name", "expected"),
+        [
+            ("OpenAI", "openai"),
+            ("Bedrock", "bedrock"),
+            ("Azure", "azure"),
+            ("Cohere", "cohere"),
+            ("Typo", "typo"),
+            ("", "openai"),
+        ],
+    )
     @pytest.mark.asyncio
-    async def test_gen_ai_provider_name_falls_back_to_openai_for_anything_else(
-        self,
+    async def test_gen_ai_provider_name_is_the_configured_name(
+        self, provider_name: str, expected: str
     ) -> None:
-        # A binary choice, not a passthrough: Bedrock, Azure, Cohere, a typo, or nothing at all all
-        # report `openai`, because that mirrors which chat model class is really instantiated.
         ctx, rec = _recording()
-        cfg = {**BASE_CONFIG, "provider": {"name": "Bedrock"}}
+        cfg = {**BASE_CONFIG, "provider": {"name": provider_name}}
         llm = _FakeToolModel(replies=[_ai_message("hi")])
         with ctx:
             await create_langchain_agents_handler(llm)(cfg, "q")
-        assert rec.root.attributes["gen_ai.provider.name"] == "openai"
+        assert rec.root.attributes["gen_ai.provider.name"] == expected
 
     @pytest.mark.asyncio
     async def test_response_model_is_the_requested_name(self) -> None:
@@ -2160,12 +2169,20 @@ def factory(config: Any) -> Any:
             **BASE_CONFIG,
             "model": {
                 "name": "gpt-4o",
-                "parameters": {"temperature": 0.2, "max_tokens": 512},
+                "parameters": {
+                    "temperature": 0.2,
+                    "max_tokens": 512,
+                    "tools": [{"name": "openai-tool"}],
+                },
             },
         }
         with ctx:
             result = await create_langchain_agents_handler(factory)(cfg, "q")
-        assert seen[0]["model"]["parameters"] == {"temperature": 0.2, "max_tokens": 512}
+        assert seen[0]["model"]["parameters"] == {
+            "temperature": 0.2,
+            "max_tokens": 512,
+            "tools": [{"name": "openai-tool"}],
+        }
         assert result["output"] == "from-factory"
 
     @pytest.mark.asyncio
@@ -2191,7 +2208,11 @@ async def test_default_openai_constructor_receives_parameters(self) -> None:
             **BASE_CONFIG,
             "model": {
                 "name": "gpt-4o",
-                "parameters": {"temperature": 0.2, "max_tokens": 512},
+                "parameters": {
+                    "temperature": 0.2,
+                    "max_tokens": 512,
+                    "tools": [{"name": "openai-tool"}],
+                },
             },
         }
         with (
@@ -2202,6 +2223,7 @@ async def test_default_openai_constructor_receives_parameters(self) -> None:
         assert ctor.call_args.kwargs == {
             "temperature": 0.2,
             "max_tokens": 512,
+            "tools": [{"name": "openai-tool"}],
             "model": "gpt-4o",
         }
 
@@ -2227,6 +2249,146 @@ async def test_default_anthropic_constructor_receives_parameters(self) -> None:
             "model": "claude-sonnet-4-5",
         }
 
+    @pytest.mark.asyncio
+    async def test_bedrock_region_is_prepended_to_the_default_constructor(
+        self,
+    ) -> None:
+        ctx, _rec = _recording()
+        llm = _FakeToolModel(replies=[_ai_message("bedrock")])
+        ctor = MagicMock(return_value=llm)
+        cfg = {
+            **BASE_CONFIG,
+            "provider": {"name": "Bedrock"},
+            "tools": TOOL_CONFIG["tools"],
+            "model": {
+                "name": "anthropic.claude-sonnet-4-5",
+                "region": "us",
+                "parameters": {
+                    "temperature": 0.2,
+                    "tools": [{"name": "duplicated-search"}],
+                },
+            },
+        }
+        with (
+            ctx,
+            patch.dict(
+                "sys.modules",
+                {"langchain_aws": MagicMock(ChatBedrockConverse=ctor)},
+            ),
+        ):
+            await create_langchain_agents_handler()(
+                cfg, "q", {"search": AsyncMock(return_value="result")}
+            )
+        assert ctor.call_args.kwargs == {
+            "temperature": 0.2,
+            "model": "us.anthropic.claude-sonnet-4-5",
+        }
+        assert cfg["model"]["parameters"]["tools"] == [{"name": "duplicated-search"}]
+        assert cfg["model"]["name"] == "anthropic.claude-sonnet-4-5"
+
+    @pytest.mark.asyncio
+    async def test_bedrock_region_prefix_is_idempotent(self) -> None:
+        ctx, _rec = _recording()
+        llm = _FakeToolModel(replies=[_ai_message("bedrock")])
+        ctor = MagicMock(return_value=llm)
+        cfg = {
+            **BASE_CONFIG,
+            "provider": {"name": "Bedrock"},
+            "model": {
+                "name": "us.anthropic.claude-sonnet-4-5",
+                "region": "us",
+            },
+        }
+        with (
+            ctx,
+            patch.dict(
+                "sys.modules",
+                {"langchain_aws": MagicMock(ChatBedrockConverse=ctor)},
+            ),
+        ):
+            await create_langchain_agents_handler()(cfg, "q")
+        assert ctor.call_args.kwargs["model"] == "us.anthropic.claude-sonnet-4-5"
+
+    @pytest.mark.asyncio
+    async def test_bedrock_without_region_keeps_the_model_name(self) -> None:
+        ctx, _rec = _recording()
+        llm = _FakeToolModel(replies=[_ai_message("bedrock")])
+        ctor = MagicMock(return_value=llm)
+        cfg = {
+            **BASE_CONFIG,
+            "provider": {"name": "Bedrock"},
+            "model": {"name": "anthropic.claude-sonnet-4-5"},
+        }
+        with (
+            ctx,
+            patch.dict(
+                "sys.modules",
+                {"langchain_aws": MagicMock(ChatBedrockConverse=ctor)},
+            ),
+        ):
+            await create_langchain_agents_handler()(cfg, "q")
+        assert ctor.call_args.kwargs["model"] == "anthropic.claude-sonnet-4-5"
+
+    def test_bedrock_without_langchain_aws_has_a_clear_error(self) -> None:
+        cfg = {
+            **BASE_CONFIG,
+            "provider": {"name": "Bedrock"},
+            "model": {"name": "anthropic.claude-sonnet-4-5"},
+        }
+        with patch(
+            "importlib.import_module",
+            side_effect=ModuleNotFoundError("No module named 'langchain_aws'"),
+        ):
+            with pytest.raises(
+                ImportError,
+                match=r"pip install langchain-aws",
+            ):
+                handler_mod._make_default_chat_model(cfg)
+
+    @pytest.mark.asyncio
+    async def test_non_bedrock_ignores_model_region(self) -> None:
+        ctx, _rec = _recording()
+        llm = _FakeToolModel(replies=[_ai_message("openai")])
+        ctor = MagicMock(return_value=llm)
+        cfg = {
+            **BASE_CONFIG,
+            "provider": {"name": "OpenAI"},
+            "model": {"name": "gpt-4o", "region": "us"},
+        }
+        with (
+            ctx,
+            patch.dict("sys.modules", {"langchain_openai": MagicMock(ChatOpenAI=ctor)}),
+        ):
+            await create_langchain_agents_handler()(cfg, "q")
+        assert ctor.call_args.kwargs["model"] == "gpt-4o"
+
+    @pytest.mark.asyncio
+    async def test_factory_receives_prefixed_bedrock_name_without_mutating_config(
+        self,
+    ) -> None:
+        ctx, _rec = _recording()
+        llm = _FakeToolModel(replies=[_ai_message("from-factory")])
+        seen: list[Any] = []
+
+        def factory(config: Any) -> Any:
+            seen.append(config)
+            return llm
+
+        cfg = {
+            **BASE_CONFIG,
+            "provider": {"name": "Bedrock"},
+            "model": {
+                "name": "anthropic.claude-sonnet-4-5",
+                "region": "us",
+                "parameters": {"temperature": 0.2},
+            },
+        }
+        with ctx:
+            await create_langchain_agents_handler(factory)(cfg, "q")
+        assert seen[0]["model"]["name"] == "us.anthropic.claude-sonnet-4-5"
+        assert cfg["model"]["name"] == "anthropic.claude-sonnet-4-5"
+        assert seen[0] is not cfg
+
     @pytest.mark.asyncio
     async def test_factory_is_resolved_on_the_streaming_path(self) -> None:
         ctx, _rec = _recording()