Skip to content

Commit 7d84079

Browse files
cosminachocosminacho
committed
fixes
1 parent e5fc320 commit 7d84079

3 files changed

Lines changed: 99 additions & 68 deletions

File tree

‎packages/uipath_langchain_client/src/uipath_langchain_client/clients/azure/chat_models.py‎

Lines changed: 42 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,46 +1,68 @@
11
from typing import Self
22

3-
from pydantic import model_validator
3+
from httpx import URL, Request
4+
from pydantic import Field, model_validator
45

56
from uipath_langchain_client.base_client import UiPathBaseChatModel
67
from uipath_langchain_client.settings import UiPathAPIConfig
78

89
try:
9-
from azure.ai.inference import ChatCompletionsClient
10-
from azure.ai.inference.aio import ChatCompletionsClient as ChatCompletionsClientAsync
11-
from azure.core.credentials import AzureKeyCredential
12-
from langchain_azure_ai.chat_models import AzureAIChatCompletionsModel
10+
from azure.core.credentials import AzureKeyCredential, TokenCredential
11+
from azure.core.credentials_async import AsyncTokenCredential
12+
from langchain_azure_ai.chat_models import AzureAIOpenAIApiChatModel
13+
from openai import AsyncOpenAI, OpenAI
1314
except ImportError as e:
1415
raise ImportError(
1516
"The 'azure' extra is required to use UiPathAzureAIChatCompletionsModel. "
1617
"Install it with: uv add uipath-langchain-client[azure]"
1718
) from e
1819

1920

20-
class UiPathAzureAIChatCompletionsModel(UiPathBaseChatModel, AzureAIChatCompletionsModel): # type: ignore[override]
21+
class UiPathAzureAIChatCompletionsModel(UiPathBaseChatModel, AzureAIOpenAIApiChatModel): # type: ignore[override]
2122
api_config: UiPathAPIConfig = UiPathAPIConfig(
2223
api_type="completions",
2324
client_type="passthrough",
2425
vendor_type="azure",
25-
freeze_base_url=True,
26+
freeze_base_url=False,
2627
)
2728

28-
# Override fields to avoid errors when instantiating the class
29-
endpoint: str | None = "PLACEHOLDER"
29+
# Override fields to avoid env var lookup / validation errors at instantiation
30+
endpoint: str | None = Field(default="PLACEHOLDER")
31+
credential: str | AzureKeyCredential | TokenCredential | AsyncTokenCredential | None = Field(default="PLACEHOLDER")
3032

3133
@model_validator(mode="after")
3234
def setup_uipath_client(self) -> Self:
33-
# TODO: finish implementation once we have a proper model in UiPath API
34-
self._client = ChatCompletionsClient(
35-
endpoint="PLACEHOLDER",
36-
credential=AzureKeyCredential("PLACEHOLDER"),
37-
model=self.model_name,
38-
**self.client_kwargs,
35+
base_url = str(self.uipath_sync_client.base_url).rstrip("/")
36+
37+
def fix_url_and_api_flavor_header(request: Request):
38+
url_suffix = str(request.url).split(base_url)[-1]
39+
if "responses" in url_suffix:
40+
request.headers["X-UiPath-LlmGateway-ApiFlavor"] = "responses"
41+
else:
42+
request.headers["X-UiPath-LlmGateway-ApiFlavor"] = "chat-completions"
43+
request.url = URL(base_url)
44+
45+
async def fix_url_and_api_flavor_header_async(request: Request):
46+
url_suffix = str(request.url).split(base_url)[-1]
47+
if "responses" in url_suffix:
48+
request.headers["X-UiPath-LlmGateway-ApiFlavor"] = "responses"
49+
else:
50+
request.headers["X-UiPath-LlmGateway-ApiFlavor"] = "chat-completions"
51+
request.url = URL(base_url)
52+
53+
self.uipath_sync_client.event_hooks["request"].append(fix_url_and_api_flavor_header)
54+
self.uipath_async_client.event_hooks["request"].append(fix_url_and_api_flavor_header_async)
55+
56+
self.root_client = OpenAI(
57+
api_key="PLACEHOLDER",
58+
max_retries=0, # handled by the UiPath client
59+
http_client=self.uipath_sync_client,
3960
)
40-
self._async_client = ChatCompletionsClientAsync(
41-
endpoint="PLACEHOLDER",
42-
credential=AzureKeyCredential("PLACEHOLDER"),
43-
model=self.model_name,
44-
**self.client_kwargs,
61+
self.root_async_client = AsyncOpenAI(
62+
api_key="PLACEHOLDER",
63+
max_retries=0, # handled by the UiPath client
64+
http_client=self.uipath_async_client,
4565
)
66+
self.client = self.root_client.chat.completions
67+
self.async_client = self.root_async_client.chat.completions
4668
return self

‎packages/uipath_langchain_client/src/uipath_langchain_client/clients/azure/embeddings.py‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,46 +1,45 @@
11
from typing import Self
22

3-
from pydantic import model_validator
3+
from pydantic import Field, model_validator
44

55
from uipath_langchain_client.base_client import UiPathBaseEmbeddings
66
from uipath_langchain_client.settings import UiPathAPIConfig
77

88
try:
9-
from azure.ai.inference import EmbeddingsClient
10-
from azure.ai.inference.aio import EmbeddingsClient as EmbeddingsClientAsync
11-
from langchain_azure_ai.embeddings import AzureAIEmbeddingsModel
9+
from azure.core.credentials import AzureKeyCredential, TokenCredential
10+
from azure.core.credentials_async import AsyncTokenCredential
11+
from langchain_azure_ai.embeddings import AzureAIOpenAIApiEmbeddingsModel
12+
from openai import AsyncOpenAI, OpenAI
1213
except ImportError as e:
1314
raise ImportError(
1415
"The 'azure' extra is required to use UiPathAzureAIEmbeddingsModel. "
1516
"Install it with: uv add uipath-langchain-client[azure]"
1617
) from e
1718

1819

19-
class UiPathAzureAIEmbeddingsModel(UiPathBaseEmbeddings, AzureAIEmbeddingsModel): # type: ignore[override]
20+
class UiPathAzureAIEmbeddingsModel(UiPathBaseEmbeddings, AzureAIOpenAIApiEmbeddingsModel): # type: ignore[override]
2021
api_config: UiPathAPIConfig = UiPathAPIConfig(
2122
api_type="embeddings",
2223
client_type="passthrough",
2324
vendor_type="azure",
2425
freeze_base_url=True,
2526
)
2627

27-
# Override fields to avoid errors when instantiating the class
28-
endpoint: str | None = "PLACEHOLDER"
29-
credentials: str | None = "PLACEHOLDER"
28+
# Override fields to avoid env var lookup / validation errors at instantiation
29+
model: str = Field(default="", alias="model_name")
30+
endpoint: str | None = Field(default="PLACEHOLDER")
31+
credential: str | AzureKeyCredential | TokenCredential | AsyncTokenCredential | None = Field(default="PLACEHOLDER")
3032

3133
@model_validator(mode="after")
3234
def setup_uipath_client(self) -> Self:
33-
# TODO: finish implementation once we have a proper model in UiPath API
34-
self._client = EmbeddingsClient(
35-
endpoint="PLACEHOLDER",
36-
credentials="PLACEHOLDER",
37-
model=self.model_name,
38-
**self.client_kwargs,
39-
)
40-
self._async_client = EmbeddingsClientAsync(
41-
endpoint="PLACEHOLDER",
42-
credentials="PLACEHOLDER",
43-
model=self.model_name,
44-
**self.client_kwargs,
45-
)
35+
self.client = OpenAI(
36+
api_key="PLACEHOLDER",
37+
max_retries=0, # handled by the UiPath client
38+
http_client=self.uipath_sync_client,
39+
).embeddings
40+
self.async_client = AsyncOpenAI(
41+
api_key="PLACEHOLDER",
42+
max_retries=0, # handled by the UiPath client
43+
http_client=self.uipath_async_client,
44+
).embeddings
4645
return self

0 commit comments

Comments
 (0)