Merge pull request trustyai-explainability#32 from trustyai-explainability/develop

Sync with trustyai develop

Author: RobGeada

Dockerfile.server

@@ -69,7 +69,8 @@ ENV HF_HOME=/app/.cache/huggingface \
     FASTEMBED_CACHE_PATH=/app/.cache/fastembed \
     PATH="/app/.venv/bin:$PATH" \
     PYTHONUNBUFFERED=1 \
-    PYTHONDONTWRITEBYTECODE=1
+    PYTHONDONTWRITEBYTECODE=1 \
+    DO_NOT_TRACK=1
 
 EXPOSE 8000
 ENTRYPOINT ["./scripts/entrypoint.sh"]

Makefile

@@ -1,3 +1,22 @@
+PLATFORMS ?= linux/$(TARGETARCH)
+TARGETARCH ?= amd64
+
+DOCKER_BUILDX_CMD ?= docker buildx
+IMAGE_BUILD_CMD ?= $(DOCKER_BUILDX_CMD) build
+IMAGE_BUILD_EXTRA_OPTS ?=
+
+IMAGE_REGISTRY ?= quay.io/trustyai
+IMAGE_NAME := nemo-guardrails-server
+IMAGE_REPO ?= $(IMAGE_REGISTRY)/$(IMAGE_NAME)
+IMAGE_TAG ?= $(IMAGE_REPO):$(GIT_TAG)
+
+GIT_COMMIT_SHA ?= $(shell git rev-parse HEAD)
+GIT_TAG ?= $(or $(shell git describe --abbrev=0 2>/dev/null),$(shell git rev-parse --short HEAD))
+BUILD_REF ?= $(or $(shell git describe --abbrev=0 2>/dev/null),$(GIT_TAG)) # release tag or fallback to git commit hash
+
+# The name of the kind cluster to use for the "kind-load" target.
+KIND_CLUSTER ?= kind
+
 .PHONY: all test tests test_watch test_coverage test_profile docs docs-strict docs-serve docs-update-cards docs-check-cards docs-watch-cards pre_commit help
 
 # Default target executed when no specific target is provided to make.
@@ -46,6 +65,35 @@ pre_commit:
 	pre-commit install
 	pre-commit run --all-files
 
+# BUILD
+.PHONY: image-build
+image-build: # Build the image using Docker buildx.
+	$(IMAGE_BUILD_CMD) -t $(IMAGE_TAG) \
+		--file Dockerfile.server \
+		--platform=$(PLATFORMS) \
+		--build-arg COMMIT_SHA=${GIT_COMMIT_SHA} \
+		--build-arg BUILD_REF=${BUILD_REF} \
+		$(PUSH) \
+		$(LOAD) \
+		$(IMAGE_BUILD_EXTRA_OPTS) ./
+
+# Build the container image for the server
+.PHONY: image-local-build
+image-local-build: # Build the image using Docker buildx
+	set -e; \
+	builder=$$($(DOCKER_BUILDX_CMD) create --use); \
+	trap '$(DOCKER_BUILDX_CMD) rm -f "$$builder"' EXIT; \
+	$(MAKE) image-build PUSH="$(PUSH)" LOAD="$(LOAD)"
+
+.PHONY: image-local-push
+image-local-push: # Push the image to the local Docker registry
+image-local-push: PUSH=--push # Build the image for local development and push it to $IMAGE_REPO.
+image-local-push: image-local-build
+
+.PHONY: image-kind
+image-kind: LOAD=--load
+image-kind: image-build # Build the image and load it to kind cluster $KIND_CLUSTER ("kind" by default)
+	kind load docker-image $(IMAGE_TAG) --name $(KIND_CLUSTER)
 
 # HELP
 
@@ -64,3 +112,7 @@ help:
 	@echo 'docs-watch-cards             - watch for file changes and auto-update cards'
 	@echo 'docs-check-redirects         - validate that all redirect targets exist'
 	@echo 'pre_commit                   - run pre-commit hooks'
+	@echo 'image-build                  - build the image using Docker buildx'
+	@echo 'image-local-build            - build the image using Docker buildx for local development'
+	@echo 'image-local-push             - build the image and push it to the local Docker registry'
+	@echo 'image-kind                   - build the image and load it to kind cluster $KIND_CLUSTER ("kind" by default)'

nemoguardrails/actions/llm/utils.py

@@ -25,6 +25,8 @@
 from nemoguardrails.colang.v2_x.lang.colang_ast import Flow
 from nemoguardrails.colang.v2_x.runtime.flows import InternalEvent, InternalEvents
 from nemoguardrails.context import (
+    api_request_headers_var,
+    get_llm_needs_runtime_auth,
     llm_call_info_var,
     llm_response_metadata_var,
     reasoning_trace_var,
@@ -54,6 +56,45 @@
 ]
 
 
+_INFRA_PREFIXES = ("x-forwarded", "x-real-", "x-request-id", "x-remote-")
+
+
+def get_extra_headers_from_request(forward_auth: bool = True) -> Optional[Dict[str, str]]:
+    """Forward X-* headers from the incoming request to the LLM call.
+
+    Excludes proxy/infra headers and the inbound Authorization header (which
+    typically carries K8s/proxy auth and must never be forwarded to the LLM).
+    When forward_auth is True, forwards X-Authorization as Authorization to
+    the LLM provider.
+    """
+    request_headers = api_request_headers_var.get()
+    if not request_headers:
+        return None
+
+    extra_headers = {}
+
+    for k, v in request_headers.items():
+        lower = k.lower()
+        if lower in ("authorization", "x-authorization"):
+            continue
+        if lower.startswith("x-") and not lower.startswith(_INFRA_PREFIXES):
+            extra_headers[k] = v
+
+    if forward_auth:
+        auth = request_headers.get("x-authorization")
+        if auth:
+            extra_headers["Authorization"] = auth
+
+    return extra_headers or None
+
+
+def get_extra_headers_for_llm(llm: "BaseLanguageModel") -> Dict[str, str]:
+    """Return extra_headers dict for an LLM based on its runtime auth needs."""
+    needs_runtime_auth = get_llm_needs_runtime_auth(llm)
+    extra_headers = get_extra_headers_from_request(forward_auth=needs_runtime_auth)
+    return extra_headers or {}
+
+
 def _infer_provider_from_module(llm: BaseLanguageModel) -> Optional[str]:
     """Infer provider name from the LLM's module path.
 
@@ -224,6 +265,15 @@ async def llm_call(
         llm_params_with_stop = llm_params
 
     filtered_params = _filter_params_for_openai_reasoning_models(llm, llm_params_with_stop)
+
+    # Forward X-* headers from the incoming request. Only forward Authorization
+    # to models that need runtime auth (i.e. had no valid static API key at init).
+    needs_runtime_auth = get_llm_needs_runtime_auth(llm)
+    extra_headers = get_extra_headers_from_request(forward_auth=needs_runtime_auth)
+    if extra_headers:
+        filtered_params = filtered_params.copy() if filtered_params else {}
+        filtered_params["extra_headers"] = extra_headers
+
     generation_llm: Union[BaseLanguageModel, Runnable] = llm.bind(**filtered_params) if filtered_params else llm
 
     if streaming_handler:

nemoguardrails/context.py

@@ -62,3 +62,24 @@
 llm_response_metadata_var: contextvars.ContextVar[Optional[dict]] = contextvars.ContextVar(
     "llm_response_metadata", default=None
 )
+
+# The request headers from the incoming API request (set by the server layer).
+# Used to forward Authorization and X-* headers to outgoing LLM calls.
+api_request_headers_var: contextvars.ContextVar[Optional[Dict[str, str]]] = contextvars.ContextVar(
+    "api_request_headers", default=None
+)
+
+# Registry tracking which LLM instances need runtime auth (i.e. had no valid
+# static API key at init time). Keyed by id(llm) — safe because LLM objects
+# are server-lifetime singletons created once in _init_llms().
+_llm_needs_runtime_auth: Dict[int, bool] = {}
+
+
+def set_llm_needs_runtime_auth(llm: Any, needs_auth: bool) -> None:
+    """Register whether an LLM instance needs runtime auth from request headers."""
+    _llm_needs_runtime_auth[id(llm)] = needs_auth
+
+
+def get_llm_needs_runtime_auth(llm: Any) -> bool:
+    """Check whether an LLM instance needs runtime auth from request headers."""
+    return _llm_needs_runtime_auth.get(id(llm), False)

nemoguardrails/library/hallucination/actions.py

@@ -22,6 +22,7 @@
 from nemoguardrails import RailsConfig
 from nemoguardrails.actions import action
 from nemoguardrails.actions.llm.utils import (
+    get_extra_headers_for_llm,
     get_multiline_response,
     llm_call,
     strip_quotes,
@@ -79,7 +80,11 @@ async def self_check_hallucination(
 
         # Generate multiple responses with temperature 1.
         # Bind the config parameters to the LLM for this call
-        llm_with_config = llm.bind(temperature=1.0, n=num_responses)
+        bind_kwargs = {"temperature": 1.0, "n": num_responses}
+        extra_headers = get_extra_headers_for_llm(llm)
+        if extra_headers:
+            bind_kwargs["extra_headers"] = extra_headers
+        llm_with_config = llm.bind(**bind_kwargs)
         extra_llm_response = await llm_with_config.agenerate(
             [formatted_prompt],
             callbacks=logging_callback_manager_for_chain.handlers,

nemoguardrails/logging/callbacks.py

@@ -12,7 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import copy
 import logging
+import re
 import uuid
 from time import time
 from typing import Any, Dict, List, Optional, cast
@@ -36,6 +38,20 @@
 
 log = logging.getLogger(__name__)
 
+_BEARER_RE = re.compile(r"(Bearer\s+)\S+", re.IGNORECASE)
+
+
+def _redact_invocation_params(params: Dict[str, Any]) -> Dict[str, Any]:
+    """Return a copy of invocation params with auth tokens redacted."""
+    if "extra_headers" not in params:
+        return params
+    params = copy.copy(params)
+    headers = params["extra_headers"]
+    params["extra_headers"] = {
+        k: _BEARER_RE.sub(r"\1[REDACTED]", v) if "authorization" in k.lower() else v for k, v in headers.items()
+    }
+    return params
+
 
 class LoggingCallbackHandler(AsyncCallbackHandler):
     """Async callback handler that can be used to handle callbacks from langchain."""
@@ -64,7 +80,7 @@ async def on_llm_start(
         if explain_info:
             explain_info.llm_calls.append(llm_call_info)
 
-        log.info("Invocation Params :: %s", kwargs.get("invocation_params", {}))
+        log.info("Invocation Params :: %s", _redact_invocation_params(kwargs.get("invocation_params", {})))
         log.info(
             "Prompt :: %s",
             prompts[0],
@@ -123,7 +139,7 @@ async def on_chat_model_start(
             ]
         )
 
-        log.info("Invocation Params :: %s", kwargs.get("invocation_params", {}))
+        log.info("Invocation Params :: %s", _redact_invocation_params(kwargs.get("invocation_params", {})))
         log.info(
             "Prompt Messages :: %s",
             prompt,

nemoguardrails/rails/llm/llmrails.py

@@ -66,6 +66,7 @@
     generation_options_var,
     llm_stats_var,
     raw_llm_request,
+    set_llm_needs_runtime_auth,
     streaming_handler_var,
 )
 from nemoguardrails.embeddings.index import EmbeddingsIndex
@@ -369,19 +370,31 @@ def _prepare_model_kwargs(self, model_config):
             model_config: The model configuration object
 
         Returns:
-            dict: The prepared kwargs for model initialization
+            dict: kwargs dict for model initialization
         """
         # Make a copy to avoid modifying the original model config
         kwargs = dict(model_config.parameters) if model_config.parameters else {}
 
-        # If the optional API Key Environment Variable is set, add it to kwargs
         if model_config.api_key_env_var:
             api_key = os.environ.get(model_config.api_key_env_var)
             if api_key:
                 kwargs["api_key"] = api_key
+                kwargs["openai_api_key"] = api_key
+        elif "api_key" not in kwargs:
+            # Placeholder to satisfy LangChain constructors. Real auth arrives
+            # via forwarded request headers at call time. If no auth header is
+            # provided, the LLM call will fail with a 401 — this is the intended
+            # fail-safe: requests without valid auth are rejected by the provider.
+            kwargs["api_key"] = "runtime-provided"
+            kwargs["openai_api_key"] = "runtime-provided"
 
         return kwargs
 
+    @staticmethod
+    def _needs_runtime_auth(kwargs):
+        """Check if the model needs auth provided at request time."""
+        return kwargs.get("api_key") == "runtime-provided"
+
     def _init_llms(self):
         """
         Initializes the right LLM engines based on the configuration.
@@ -421,6 +434,7 @@ def _init_llms(self):
                     mode="chat",
                     kwargs=kwargs,
                 )
+                set_llm_needs_runtime_auth(self.llm, self._needs_runtime_auth(kwargs))
                 self.runtime.register_action_param("llm", self.llm)
 
             else:
@@ -453,6 +467,7 @@ def _init_llms(self):
                     mode=mode,
                     kwargs=kwargs,
                 )
+                set_llm_needs_runtime_auth(llm_model, self._needs_runtime_auth(kwargs))
 
                 # Configure the model based on its type
                 if llm_config.type == "main":