feat: inject VEADK source header into migrated models (#182)

Author: warm-wm

agentkit/frameworks/model_replacement.py

@@ -7,10 +7,12 @@
 import os
 from typing import Any
 
-
 logger = logging.getLogger(__name__)
 ARK_DEFAULT_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
+VEADK_SOURCE_HEADER = "veadk-source"
+VEADK_SOURCE_VALUE = "veadk"
 _DISABLED_VALUES = {"0", "false", "off", "none", "disabled"}
+_HEADER_PATCH_MARKER = "__agentkit_veadk_source_header__"
 _MODEL_PARAM_KEYS = {
     "frequency_penalty",
     "max_tokens",
@@ -49,6 +51,55 @@ def _replacement_enabled() -> bool:
     return value not in _DISABLED_VALUES
 
 
+def _merge_veadk_source_header(headers: Any = None) -> dict[str, Any]:
+    merged = {
+        key: value
+        for key, value in dict(headers or {}).items()
+        if key.lower() != VEADK_SOURCE_HEADER
+    }
+    merged[VEADK_SOURCE_HEADER] = VEADK_SOURCE_VALUE
+    return merged
+
+
+def _header_kwarg_model_cls(base_cls: type[Any], header_kwarg: str) -> type[Any]:
+    if getattr(base_cls, _HEADER_PATCH_MARKER, False):
+        return base_cls
+
+    class AgentKitHeaderModel(base_cls):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            kwargs[header_kwarg] = _merge_veadk_source_header(kwargs.get(header_kwarg))
+            super().__init__(*args, **kwargs)
+
+    AgentKitHeaderModel.__name__ = base_cls.__name__
+    AgentKitHeaderModel.__qualname__ = base_cls.__qualname__
+    AgentKitHeaderModel.__module__ = base_cls.__module__
+    setattr(AgentKitHeaderModel, _HEADER_PATCH_MARKER, True)
+    return AgentKitHeaderModel
+
+
+def _openai_client_args_model_cls(base_cls: type[Any]) -> type[Any]:
+    if getattr(base_cls, _HEADER_PATCH_MARKER, False):
+        return base_cls
+
+    class AgentKitHeaderModel(base_cls):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            if args or kwargs.get("client") is not None:
+                super().__init__(*args, **kwargs)
+                return
+            client_args = dict(kwargs.get("client_args") or {})
+            client_args["default_headers"] = _merge_veadk_source_header(
+                client_args.get("default_headers")
+            )
+            kwargs["client_args"] = client_args
+            super().__init__(*args, **kwargs)
+
+    AgentKitHeaderModel.__name__ = base_cls.__name__
+    AgentKitHeaderModel.__qualname__ = base_cls.__qualname__
+    AgentKitHeaderModel.__module__ = base_cls.__module__
+    setattr(AgentKitHeaderModel, _HEADER_PATCH_MARKER, True)
+    return AgentKitHeaderModel
+
+
 def _replacement_model_config(config: dict[str, Any]) -> dict[str, Any]:
     model_id = _target_model_id()
     if not model_id:
@@ -90,6 +141,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                 client_args={
                     "base_url": _target_base_url(),
                     "api_key": api_key,
+                    "default_headers": _merge_veadk_source_header(),
                 },
                 **_replacement_model_config(kwargs),
             )
@@ -101,14 +153,57 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 def _patch_model_attr(module_name: str, attr: str, replacement: type[Any]) -> bool:
     try:
         module = importlib.import_module(module_name)
-    except Exception:
+    except Exception:  # noqa: BLE001 - optional framework imports are best effort
         return False
     if hasattr(module, attr):
         setattr(module, attr, replacement)
         return True
     return False
 
 
+def _patch_langchain_model_headers() -> bool:
+    try:
+        module = importlib.import_module("langchain_openai.chat_models.base")
+        replacement = _header_kwarg_model_cls(module.ChatOpenAI, "default_headers")
+    except Exception:  # noqa: BLE001 - optional framework imports are best effort
+        return False
+
+    patched = False
+    for module_name in (
+        "langchain_openai",
+        "langchain_openai.chat_models",
+        "langchain_openai.chat_models.base",
+    ):
+        patched = _patch_model_attr(module_name, "ChatOpenAI", replacement) or patched
+    return patched
+
+
+def _patch_adk_model_headers() -> bool:
+    try:
+        module = importlib.import_module("google.adk.models.lite_llm")
+        replacement = _header_kwarg_model_cls(module.LiteLlm, "extra_headers")
+    except Exception:  # noqa: BLE001 - optional framework imports are best effort
+        return False
+
+    patched = False
+    for module_name in ("google.adk.models", "google.adk.models.lite_llm"):
+        patched = _patch_model_attr(module_name, "LiteLlm", replacement) or patched
+    return patched
+
+
+def _patch_strands_openai_model_headers() -> bool:
+    try:
+        module = importlib.import_module("strands.models.openai")
+        replacement = _openai_client_args_model_cls(module.OpenAIModel)
+    except Exception:  # noqa: BLE001 - optional framework imports are best effort
+        return False
+
+    patched = False
+    for module_name in ("strands.models", "strands.models.openai"):
+        patched = _patch_model_attr(module_name, "OpenAIModel", replacement) or patched
+    return patched
+
+
 def _apply_env_aliases() -> bool:
     changed = False
     model_id = _target_model_id()
@@ -138,7 +233,7 @@ def _apply_env_aliases() -> bool:
 def _patch_strands_model_classes() -> bool:
     try:
         replacement = _agentkit_openai_model_cls()
-    except Exception:
+    except Exception:  # noqa: BLE001 - optional framework imports are best effort
         return False
 
     patched = False
@@ -155,21 +250,33 @@ def apply_agentkit_model_replacement() -> bool:
     """Apply explicit target-model replacement for generated migration wrappers.
 
     The helper intentionally stays low-intrusion: it normalizes common model
-    environment variables and patches only stable Strands Bedrock/Anthropic
-    model classes before user code constructs them. It does not rewrite project
-    source code or chase framework-specific model SDK constructor details.
+    environment variables and patches stable framework model constructors before
+    user code constructs them. It does not rewrite project source code or patch
+    global HTTP/OpenAI clients.
     """
 
     if not _replacement_enabled():
         logger.info("AgentKit model replacement disabled by ARK_MODEL_REPLACEMENT.")
         return False
     env_changed = _apply_env_aliases()
+    langchain_headers_patched = _patch_langchain_model_headers()
+    adk_headers_patched = _patch_adk_model_headers()
+    strands_headers_patched = _patch_strands_openai_model_headers()
     patched = _patch_strands_model_classes()
     logger.info(
-        "AgentKit model replacement enabled: env_aliases=%s strands_model_patch=%s target_model_configured=%s base_url_configured=%s",
+        "AgentKit model replacement enabled: env_aliases=%s langchain_headers=%s adk_headers=%s strands_headers=%s strands_model_patch=%s target_model_configured=%s base_url_configured=%s",
         env_changed,
+        langchain_headers_patched,
+        adk_headers_patched,
+        strands_headers_patched,
         patched,
         bool(_target_model_id()),
         bool(_target_base_url()),
     )
-    return env_changed or patched
+    return (
+        env_changed
+        or langchain_headers_patched
+        or adk_headers_patched
+        or strands_headers_patched
+        or patched
+    )

agentkit/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-VERSION = "0.8.5"
+VERSION = "0.8.6"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "agentkit-sdk-python"
-version = "0.8.5"
+version = "0.8.6"
 description = "Python SDK for transforming any AI agent into a production-ready application. Framework-agnostic primitives for runtime, memory, authentication, and tools with volcengine-managed infrastructure."
 readme = "README.md"
 requires-python = ">=3.10"