Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/uipath_langchain/_utils/_request_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import os
import time
from typing import Any, AsyncIterator, Dict, Iterator, Mapping
from urllib.parse import quote

import httpx
import openai
Expand Down Expand Up @@ -79,7 +80,7 @@ class UiPathRequestMixin(BaseModel):
default_headers: Mapping[str, str] | None = {
"X-UiPath-Streaming-Enabled": "false",
"X-UiPath-JobKey": os.getenv("UIPATH_JOB_KEY", ""),
"X-UiPath-ProcessKey": os.getenv("UIPATH_PROCESS_KEY", ""),
"X-UiPath-ProcessKey": quote(os.getenv("UIPATH_PROCESS_KEY", ""), safe=""),
}
model_name: str | None = Field(
default_factory=lambda: os.getenv(
Expand Down
3 changes: 2 additions & 1 deletion src/uipath_langchain/chat/bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import os
from collections.abc import Iterator
from typing import Any, Optional
from urllib.parse import quote

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.messages import BaseMessage
Expand Down Expand Up @@ -143,7 +144,7 @@ def _modify_request(self, request, **kwargs):
if job_key:
headers["X-UiPath-JobKey"] = job_key
if process_key:
headers["X-UiPath-ProcessKey"] = process_key
headers["X-UiPath-ProcessKey"] = quote(process_key, safe="")

request.headers.update(headers)

Expand Down
3 changes: 2 additions & 1 deletion src/uipath_langchain/chat/openai.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import logging
import os
from typing import Optional
from urllib.parse import quote

import httpx
from langchain_openai import AzureChatOpenAI
Expand Down Expand Up @@ -162,7 +163,7 @@ def _build_headers(self, token: str) -> dict[str, str]:
if job_key := os.getenv("UIPATH_JOB_KEY"):
headers["X-UiPath-JobKey"] = job_key
if process_key := os.getenv("UIPATH_PROCESS_KEY"):
headers["X-UiPath-ProcessKey"] = process_key
headers["X-UiPath-ProcessKey"] = quote(process_key, safe="")

# Allow extra_headers to override defaults
headers.update(self._extra_headers)
Expand Down
3 changes: 2 additions & 1 deletion src/uipath_langchain/chat/vertex.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import os
from collections.abc import AsyncIterator, Iterator
from typing import Any, Optional
from urllib.parse import quote

import httpx
from langchain_core.callbacks import (
Expand Down Expand Up @@ -266,7 +267,7 @@ def _build_headers(
if job_key := os.getenv("UIPATH_JOB_KEY"):
headers["X-UiPath-JobKey"] = job_key
if process_key := os.getenv("UIPATH_PROCESS_KEY"):
headers["X-UiPath-ProcessKey"] = process_key
headers["X-UiPath-ProcessKey"] = quote(process_key, safe="")
return headers

@staticmethod
Expand Down
133 changes: 133 additions & 0 deletions tests/chat/test_header_encoding.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
import os
from unittest.mock import MagicMock, patch
from urllib.parse import quote

import pytest

from uipath_langchain.chat.openai import UiPathChatOpenAI

NON_ASCII_PROCESS_KEY = "Solution.17.agent.GetCompanyIdAgent-請-test"
ASCII_PROCESS_KEY = "Solution.17.agent.MyAgent-test"

BASE_ENV = {
"UIPATH_URL": "https://cloud.uipath.com/org/tenant",
"UIPATH_ORGANIZATION_ID": "org-id",
"UIPATH_TENANT_ID": "tenant-id",
"UIPATH_ACCESS_TOKEN": "test-token",
}


class TestOpenAIHeaderEncoding:
"""Verify UiPathChatOpenAI percent-encodes non-ASCII header values."""

def _build_headers_with_process_key(self, process_key: str) -> dict[str, str]:
env = {**BASE_ENV, "UIPATH_PROCESS_KEY": process_key}
with patch.dict(os.environ, env, clear=False):
obj = object.__new__(UiPathChatOpenAI)
obj._agenthub_config = None
obj._byo_connection_id = None
obj._extra_headers = {}
return obj._build_headers("fake-token")

def test_ascii_process_key_unchanged(self) -> None:
headers = self._build_headers_with_process_key(ASCII_PROCESS_KEY)
assert headers["X-UiPath-ProcessKey"] == quote(ASCII_PROCESS_KEY, safe="")

def test_non_ascii_process_key_encoded(self) -> None:
headers = self._build_headers_with_process_key(NON_ASCII_PROCESS_KEY)
value = headers["X-UiPath-ProcessKey"]
assert "請" not in value
assert value == quote(NON_ASCII_PROCESS_KEY, safe="")
assert "%E8%AB%8B" in value

def test_header_value_is_ascii_safe(self) -> None:
headers = self._build_headers_with_process_key(NON_ASCII_PROCESS_KEY)
value = headers["X-UiPath-ProcessKey"]
value.encode("ascii")

def test_missing_process_key_omitted(self) -> None:
env = {**BASE_ENV}
env.pop("UIPATH_PROCESS_KEY", None)
with patch.dict(os.environ, env, clear=True):
obj = object.__new__(UiPathChatOpenAI)
obj._agenthub_config = None
obj._byo_connection_id = None
obj._extra_headers = {}
headers = obj._build_headers("fake-token")
assert "X-UiPath-ProcessKey" not in headers


class TestVertexHeaderEncoding:
"""Verify UiPathChatVertex._build_headers percent-encodes non-ASCII header values."""

@pytest.fixture(autouse=True)
def _skip_if_no_google(self) -> None:
pytest.importorskip("google.genai", reason="google-genai not installed")

def test_non_ascii_process_key_encoded(self) -> None:
from uipath_langchain.chat.vertex import UiPathChatVertex

env = {**BASE_ENV, "UIPATH_PROCESS_KEY": NON_ASCII_PROCESS_KEY}
with patch.dict(os.environ, env, clear=False):
headers = UiPathChatVertex._build_headers("fake-token")
value = headers["X-UiPath-ProcessKey"]
assert "請" not in value
assert "%E8%AB%8B" in value
value.encode("ascii")

def test_ascii_process_key_unchanged(self) -> None:
from uipath_langchain.chat.vertex import UiPathChatVertex

env = {**BASE_ENV, "UIPATH_PROCESS_KEY": ASCII_PROCESS_KEY}
with patch.dict(os.environ, env, clear=False):
headers = UiPathChatVertex._build_headers("fake-token")
assert headers["X-UiPath-ProcessKey"] == quote(ASCII_PROCESS_KEY, safe="")


class TestBedrockHeaderEncoding:
"""Verify AwsBedrockCompletionsPassthroughClient percent-encodes non-ASCII header values."""

def test_non_ascii_process_key_encoded(self) -> None:
pytest.importorskip("botocore", reason="botocore not installed")
from uipath_langchain.chat.bedrock import AwsBedrockCompletionsPassthroughClient

env = {**BASE_ENV, "UIPATH_PROCESS_KEY": NON_ASCII_PROCESS_KEY}
with (
patch.dict(os.environ, env, clear=False),
patch(
"uipath_langchain.chat.bedrock.boto3.client", return_value=MagicMock()
),
):
client = AwsBedrockCompletionsPassthroughClient(
model="test-model",
token="fake-token",
api_flavor="converse",
)
request = MagicMock()
request.url = "https://example.com/converse"
request.headers = {}
client._modify_request(request)

value = request.headers["X-UiPath-ProcessKey"]
assert "請" not in value
assert "%E8%AB%8B" in value
value.encode("ascii")


class TestRequestMixinHeaderEncoding:
"""Verify UiPathRequestMixin default_headers percent-encodes non-ASCII values."""

def test_non_ascii_process_key_encoded_in_defaults(self) -> None:
env = {**BASE_ENV, "UIPATH_PROCESS_KEY": NON_ASCII_PROCESS_KEY}
with patch.dict(os.environ, env, clear=False):
import importlib

import uipath_langchain._utils._request_mixin as mod

importlib.reload(mod)
value = mod.UiPathRequestMixin.model_fields["default_headers"].default[
"X-UiPath-ProcessKey"
]
assert "請" not in value
assert "%E8%AB%8B" in value
value.encode("ascii")