Custom callback for header injection (#47)

Author: cosminacho

.github/workflows/cd-langchain.yml

@@ -34,7 +34,8 @@ jobs:
             exit 0
           fi
 
-          echo "Core package version also changed — waiting for 'cd' workflow to succeed..."
+          CORE_VERSION=$(grep '^__version__' src/uipath/llm_client/__version__.py | sed -n 's/.*"\([^"]*\)".*/\1/p')
+          echo "Core package version also changed ($CORE_VERSION) — waiting for 'cd' workflow to succeed..."
 
           # Poll the cd workflow for up to 15 minutes
           for i in $(seq 1 30); do
@@ -43,9 +44,17 @@ jobs:
 
             case "$STATUS" in
               completed:success)
-                echo "cd workflow succeeded — waiting 120s for PyPI indexing..."
-                sleep 120
-                exit 0
+                echo "cd workflow succeeded — polling PyPI until uipath-llm-client==$CORE_VERSION is available..."
+                for j in $(seq 1 30); do
+                  if uv pip index versions uipath-llm-client 2>/dev/null | grep -q "$CORE_VERSION"; then
+                    echo "uipath-llm-client==$CORE_VERSION is available on PyPI."
+                    exit 0
+                  fi
+                  echo "  Not yet available, retrying in 30s ($j/30)..."
+                  sleep 30
+                done
+                echo "::error::uipath-llm-client==$CORE_VERSION never appeared on PyPI (15 min)"
+                exit 1
                 ;;
               completed:*)
                 echo "::error::cd workflow finished with: $STATUS"

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 All notable changes to `uipath_llm_client` (core package) will be documented in this file.
 
+## [1.5.6] - 2026-03-21
+
+### Feature
+- Added `_DYNAMIC_REQUEST_HEADERS` ContextVar and helper functions (`get_dynamic_request_headers`, `set_dynamic_request_headers`) to `utils/headers.py`
+- Inject dynamic request headers in httpx `send()` for both sync and async clients
+
 ## [1.5.5] - 2026-03-19
 
 ### Fix

packages/uipath_langchain_client/CHANGELOG.md

@@ -2,6 +2,13 @@
 
 All notable changes to `uipath_langchain_client` will be documented in this file.
 
+## [1.5.6] - 2026-03-21
+
+### Feature
+- Added `UiPathDynamicHeadersCallback`: extend and implement `get_headers()` to inject custom headers into each LLM gateway request
+- Uses `run_inline = True` so `on_chat_model_start`/`on_llm_start` run in the caller's coroutine, ensuring ContextVar mutations propagate to `httpx.send()`
+- Cleanup via `on_llm_end`/`on_llm_error`
+
 ## [1.5.5] - 2026-03-19
 
 ### Fix headers

packages/uipath_langchain_client/pyproject.toml

@@ -5,8 +5,8 @@ dynamic = ["version"]
 readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
-    "langchain>=1.2.12",
-    "uipath-llm-client>=1.5.5",
+    "langchain>=1.2.13",
+    "uipath-llm-client>=1.5.6",
 ]
 
 [project.optional-dependencies]
@@ -17,17 +17,17 @@ google = [
     "langchain-google-genai>=4.2.1",
 ]
 anthropic = [
-    "langchain-anthropic>=1.3.5",
-    "anthropic[bedrock,vertex]>=0.85.0",
+    "langchain-anthropic>=1.4.0",
+    "anthropic[bedrock,vertex]>=0.86.0",
 ]
 aws = [
-    "langchain-aws[anthropic]>=1.4.0",
+    "langchain-aws[anthropic]>=1.4.1",
 ]
 vertexai = [
     "langchain-google-vertexai>=3.2.2",
 ]
 azure = [
-    "langchain-azure-ai>=1.1.0",
+    "langchain-azure-ai>=1.1.1",
 ]
 fireworks = [
     "langchain-fireworks>=1.1.0",

packages/uipath_langchain_client/src/uipath_langchain_client/__init__.py

@@ -33,6 +33,7 @@
 """
 
 from uipath_langchain_client.__version__ import __version__
+from uipath_langchain_client.callbacks import UiPathDynamicHeadersCallback
 from uipath_langchain_client.clients import UiPathChat, UiPathEmbeddings
 from uipath_langchain_client.factory import get_chat_model, get_embedding_model
 from uipath_langchain_client.settings import (
@@ -47,6 +48,7 @@
     "get_embedding_model",
     "UiPathChat",
     "UiPathEmbeddings",
+    "UiPathDynamicHeadersCallback",
     "get_default_client_settings",
     "LLMGatewaySettings",
     "PlatformSettings",

packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py

@@ -1,3 +1,3 @@
 __title__ = "UiPath LangChain Client"
 __description__ = "A Python client for interacting with UiPath's LLM services via LangChain."
-__version__ = "1.5.5"
+__version__ = "1.5.6"

packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py

@@ -30,6 +30,10 @@
 from typing import Any, Literal
 
 from httpx import URL, Response
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import BaseMessage
@@ -322,72 +326,128 @@ class UiPathBaseChatModel(UiPathBaseLLMClient, BaseChatModel):
     from the ContextVar (populated by the httpx client's send()) and inject them into
     the AIMessage's response_metadata under the 'uipath_llmgateway_headers' key.
 
+    Dynamic request headers are injected via UiPathDynamicHeadersCallback: set
+    ``run_inline = True`` (already the default) so LangChain calls
+    ``on_chat_model_start`` in the same coroutine as ``_agenerate``, ensuring the
+    ContextVar is visible when ``httpx.send()`` fires.
+
     Passthrough clients that delegate to vendor SDKs should inherit from this class
     so that headers are captured transparently.
     """
 
     def _generate(
         self,
         messages: list[BaseMessage],
-        *args: Any,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         set_captured_response_headers({})
         try:
-            result = super()._generate(messages, *args, **kwargs)
+            result = self._uipath_generate(messages, stop=stop, run_manager=run_manager, **kwargs)
             self._inject_gateway_headers(result.generations)
             return result
         finally:
             set_captured_response_headers({})
 
+    def _uipath_generate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        """Override in subclasses to provide the core (non-wrapped) generate logic."""
+        return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
+
     async def _agenerate(
         self,
         messages: list[BaseMessage],
-        *args: Any,
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         set_captured_response_headers({})
         try:
-            result = await super()._agenerate(messages, *args, **kwargs)
+            result = await self._uipath_agenerate(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
             self._inject_gateway_headers(result.generations)
             return result
         finally:
             set_captured_response_headers({})
 
+    async def _uipath_agenerate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        """Override in subclasses to provide the core (non-wrapped) async generate logic."""
+        return await super()._agenerate(messages, stop=stop, run_manager=run_manager, **kwargs)
+
     def _stream(
         self,
         messages: list[BaseMessage],
-        *args: Any,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
         set_captured_response_headers({})
         try:
             first = True
-            for chunk in super()._stream(messages, *args, **kwargs):
+            for chunk in self._uipath_stream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            ):
                 if first:
                     self._inject_gateway_headers([chunk])
                     first = False
                 yield chunk
         finally:
             set_captured_response_headers({})
 
+    def _uipath_stream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        """Override in subclasses to provide the core (non-wrapped) stream logic."""
+        yield from super()._stream(messages, stop=stop, run_manager=run_manager, **kwargs)
+
     async def _astream(
         self,
         messages: list[BaseMessage],
-        *args: Any,
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:
         set_captured_response_headers({})
         try:
             first = True
-            async for chunk in super()._astream(messages, *args, **kwargs):
+            async for chunk in self._uipath_astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            ):
                 if first:
                     self._inject_gateway_headers([chunk])
                     first = False
                 yield chunk
         finally:
             set_captured_response_headers({})
 
+    async def _uipath_astream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        """Override in subclasses to provide the core (non-wrapped) async stream logic."""
+        async for chunk in super()._astream(messages, stop=stop, run_manager=run_manager, **kwargs):
+            yield chunk
+
     def _inject_gateway_headers(self, generations: Sequence[ChatGeneration]) -> None:
         """Inject captured gateway headers into each generation's response_metadata."""
         if not self.captured_headers:

packages/uipath_langchain_client/src/uipath_langchain_client/callbacks.py

@@ -0,0 +1,60 @@
+"""LangChain callbacks for dynamic per-request header injection."""
+
+from abc import abstractmethod
+from typing import Any
+
+from langchain_core.callbacks import BaseCallbackHandler
+
+from uipath.llm_client.utils.headers import set_dynamic_request_headers
+
+
+class UiPathDynamicHeadersCallback(BaseCallbackHandler):
+    """Base callback for injecting dynamic headers into each LLM gateway request.
+
+    Extend this class and implement ``get_headers()`` to return the headers to
+    inject. ``run_inline = True`` ensures ``on_chat_model_start`` is called
+    directly in the caller's coroutine (not via ``asyncio.gather``), so the
+    ContextVar mutation is visible when ``httpx.send()`` fires.
+
+    Example (OTEL trace propagation)::
+
+        from opentelemetry import trace, propagate
+
+        class OtelHeadersCallback(UiPathDynamicHeadersCallback):
+            def get_headers(self) -> dict[str, str]:
+                carrier: dict[str, str] = {}
+                propagate.inject(carrier)
+                return carrier
+
+        chat = get_chat_model(model_name="gpt-4o", client_settings=settings)
+        response = chat.invoke("Hello!", config={"callbacks": [OtelHeadersCallback()]})
+    """
+
+    run_inline: bool = True  # dispatch in the caller's coroutine, not via asyncio.gather
+
+    @abstractmethod
+    def get_headers(self) -> dict[str, str]:
+        """Return headers to inject into the next LLM gateway request."""
+        ...
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[Any]],
+        **kwargs: Any,
+    ) -> None:
+        set_dynamic_request_headers(self.get_headers())
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        **kwargs: Any,
+    ) -> None:
+        set_dynamic_request_headers(self.get_headers())
+
+    def on_llm_end(self, response: Any, **kwargs: Any) -> None:
+        set_dynamic_request_headers({})
+
+    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
+        set_dynamic_request_headers({})