Fix/code review bugs

packages/uipath_langchain_client/demo.py

@@ -17,9 +17,7 @@
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.tools import tool
 from uipath_langchain_client import get_chat_model, get_embedding_model
-from uipath_langchain_client.settings import get_default_client_settings
-
-from uipath.llm_client.settings.constants import RoutingMode
+from uipath_langchain_client.settings import RoutingMode, get_default_client_settings
 
 
 def demo_basic_chat():
@@ -136,8 +134,31 @@ def calculate(expression: str) -> str:
         Args:
             expression: A mathematical expression to evaluate (e.g., "2 + 2").
         """
+        import ast
+
         try:
-            result = eval(expression)
+            # Restrict to a safe subset: only literals and basic arithmetic operators.
+            # This prevents arbitrary code execution via eval().
+            tree = ast.parse(expression, mode="eval")
+            allowed_node_types = (
+                ast.Expression,
+                ast.BinOp,
+                ast.UnaryOp,
+                ast.Constant,
+                ast.Add,
+                ast.Sub,
+                ast.Mult,
+                ast.Div,
+                ast.FloorDiv,
+                ast.Mod,
+                ast.Pow,
+                ast.USub,
+                ast.UAdd,
+            )
+            for node in ast.walk(tree):
+                if not isinstance(node, allowed_node_types):
+                    return "Error: unsupported operation in expression"
+            result = eval(compile(tree, "<string>", "eval"), {"__builtins__": {}})
             return str(result)
         except Exception as e:
             return f"Error: {e}"

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.10"
+__version__ = "1.6.0"

packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py

@@ -25,7 +25,7 @@
 
 import logging
 from abc import ABC
-from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
+from collections.abc import AsyncGenerator, Generator, Mapping, Sequence
 from functools import cached_property
 from typing import Any, Literal
 
@@ -239,7 +239,7 @@ def uipath_stream(
         stream_type: Literal["text", "bytes", "lines", "raw"] = "lines",
         raise_status_error: bool = False,
         **kwargs: Any,
-    ) -> Iterator[str | bytes]:
+    ) -> Generator[str | bytes, None, None]:
         """Make a synchronous streaming HTTP request to the UiPath API.
 
         Args:
@@ -282,7 +282,7 @@ async def uipath_astream(
         stream_type: Literal["text", "bytes", "lines", "raw"] = "lines",
         raise_status_error: bool = False,
         **kwargs: Any,
-    ) -> AsyncIterator[str | bytes]:
+    ) -> AsyncGenerator[str | bytes, None]:
         """Make an asynchronous streaming HTTP request to the UiPath API.
 
         Args:
@@ -393,7 +393,7 @@ def _stream(
         stop: list[str] | None = None,
         run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
-    ) -> Iterator[ChatGenerationChunk]:
+    ) -> Generator[ChatGenerationChunk, None, None]:
         set_captured_response_headers({})
         try:
             first = True
@@ -413,7 +413,7 @@ def _uipath_stream(
         stop: list[str] | None = None,
         run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
-    ) -> Iterator[ChatGenerationChunk]:
+    ) -> Generator[ChatGenerationChunk, None, None]:
         """Override in subclasses to provide the core (non-wrapped) stream logic."""
         yield from super()._stream(messages, stop=stop, run_manager=run_manager, **kwargs)
 
@@ -423,7 +423,7 @@ async def _astream(
         stop: list[str] | None = None,
         run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
-    ) -> AsyncIterator[ChatGenerationChunk]:
+    ) -> AsyncGenerator[ChatGenerationChunk, None]:
         set_captured_response_headers({})
         try:
             first = True
@@ -443,7 +443,7 @@ async def _uipath_astream(
         stop: list[str] | None = None,
         run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
-    ) -> AsyncIterator[ChatGenerationChunk]:
+    ) -> AsyncGenerator[ChatGenerationChunk, None]:
         """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

packages/uipath_langchain_client/src/uipath_langchain_client/callbacks.py

@@ -1,14 +1,14 @@
 """LangChain callbacks for dynamic per-request header injection."""
 
-from abc import abstractmethod
+from abc import ABC, 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):
+class UiPathDynamicHeadersCallback(BaseCallbackHandler, ABC):
     """Base callback for injecting dynamic headers into each LLM gateway request.
 
     Extend this class and implement ``get_headers()`` to return the headers to

packages/uipath_langchain_client/src/uipath_langchain_client/clients/anthropic/chat_models.py

@@ -45,15 +45,17 @@ class UiPathChatAnthropic(UiPathBaseChatModel, ChatAnthropic):
     def setup_api_flavor_and_version(self) -> Self:
         self.api_config.vendor_type = self.vendor_type
         match self.vendor_type:
+            case VendorType.ANTHROPIC:
+                self.api_config.api_flavor = None
+            case VendorType.AZURE:
+                self.api_config.api_flavor = None
             case VendorType.VERTEXAI:
                 self.api_config.api_flavor = ApiFlavor.ANTHROPIC_CLAUDE
                 self.api_config.api_version = "v1beta1"
             case VendorType.AWSBEDROCK:
                 self.api_config.api_flavor = ApiFlavor.INVOKE
             case _:
-                raise ValueError(
-                    "anthropic and azure vendors are currently not supported by UiPath"
-                )
+                raise ValueError(f"Unsupported vendor_type: {self.vendor_type}")
         return self
 
     # Override fields to avoid typing issues and fix stuff

packages/uipath_langchain_client/src/uipath_langchain_client/clients/bedrock/chat_models.py

@@ -48,7 +48,7 @@ def _patched_format_data_content_block(block: dict) -> dict:
     ) from e
 
 
-class UiPathChatBedrockConverse(UiPathBaseChatModel, ChatBedrockConverse):
+class UiPathChatBedrockConverse(UiPathBaseChatModel, ChatBedrockConverse):  # type: ignore[override]
     api_config: UiPathAPIConfig = UiPathAPIConfig(
         api_type=ApiType.COMPLETIONS,
         routing_mode=RoutingMode.PASSTHROUGH,
@@ -77,7 +77,7 @@ def setup_uipath_client(self) -> Self:
         return self
 
 
-class UiPathChatBedrock(UiPathBaseChatModel, ChatBedrock):
+class UiPathChatBedrock(UiPathBaseChatModel, ChatBedrock):  # type: ignore[override]
     api_config: UiPathAPIConfig = UiPathAPIConfig(
         api_type=ApiType.COMPLETIONS,
         routing_mode=RoutingMode.PASSTHROUGH,

packages/uipath_langchain_client/src/uipath_langchain_client/clients/bedrock/utils.py

@@ -1,6 +1,7 @@
 import base64
 import json
-from typing import Any, Iterator
+from collections.abc import Generator
+from typing import Any
 
 from httpx import Client
 
@@ -50,7 +51,9 @@ def __init__(self, httpx_client: Client | None = None, region_name: str = "PLACE
         self.httpx_client = httpx_client
         self.meta = _MockClientMeta(region_name=region_name)
 
-    def _stream_generator(self, request_body: dict[str, Any]) -> Iterator[dict[str, Any]]:
+    def _stream_generator(
+        self, request_body: dict[str, Any]
+    ) -> Generator[dict[str, Any], None, None]:
         if self.httpx_client is None:
             raise ValueError("httpx_client is not set")
         with self.httpx_client.stream("POST", "/", json=_serialize_bytes(request_body)) as response:
@@ -71,15 +74,19 @@ def invoke_model(self, **kwargs: Any) -> Any:
         return {
             "body": self.httpx_client.post(
                 "/",
-                json=json.loads(kwargs.get("body", {})),
+                json=json.loads(kwargs.get("body", "{}")),
             )
         }
 
     def invoke_model_with_response_stream(self, **kwargs: Any) -> Any:
-        return {"body": self._stream_generator(json.loads(kwargs.get("body", {})))}
+        return {"body": self._stream_generator(json.loads(kwargs.get("body", "{}")))}
 
     def converse(
-        self, *, messages: list[dict[str, Any]], system: str | None = None, **params: Any
+        self,
+        *,
+        messages: list[dict[str, Any]],
+        system: list[dict[str, Any]] | None = None,
+        **params: Any,
     ) -> Any:
         if self.httpx_client is None:
             raise ValueError("httpx_client is not set")
@@ -95,7 +102,11 @@ def converse(
         ).json()
 
     def converse_stream(
-        self, *, messages: list[dict[str, Any]], system: str | None = None, **params: Any
+        self,
+        *,
+        messages: list[dict[str, Any]],
+        system: list[dict[str, Any]] | None = None,
+        **params: Any,
     ) -> Any:
         return {
             "stream": self._stream_generator(

packages/uipath_langchain_client/src/uipath_langchain_client/clients/fireworks/embeddings.py

@@ -1,6 +1,6 @@
 from typing import Self
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
 from uipath_langchain_client.base_client import UiPathBaseEmbeddings
 from uipath_langchain_client.settings import (
@@ -31,6 +31,8 @@ class UiPathFireworksEmbeddings(UiPathBaseEmbeddings, FireworksEmbeddings):
         freeze_base_url=True,
     )
 
+    model: str = Field(default="", alias="model_name")
+
     @model_validator(mode="after")
     def setup_uipath_client(self) -> Self:
         self.client = OpenAI(
@@ -48,7 +50,8 @@ def setup_uipath_client(self) -> Self:
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
         """Embed search docs."""
         return [
-            i.embedding for i in self.client.embeddings.create(input=texts, model=self.model).data
+            i.embedding
+            for i in self.client.embeddings.create(input=texts, model=self.model_name).data
         ]
 
     def embed_query(self, text: str) -> list[float]:
@@ -59,7 +62,9 @@ async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
         """Embed search docs asynchronously."""
         return [
             i.embedding
-            for i in (await self.async_client.embeddings.create(input=texts, model=self.model)).data
+            for i in (
+                await self.async_client.embeddings.create(input=texts, model=self.model_name)
+            ).data
         ]
 
     async def aembed_query(self, text: str) -> list[float]:

packages/uipath_langchain_client/src/uipath_langchain_client/clients/normalized/chat_models.py

@@ -24,7 +24,7 @@
 """
 
 import json
-from collections.abc import AsyncIterator, Callable, Iterator, Sequence
+from collections.abc import AsyncGenerator, Callable, Generator, Sequence
 from typing import Any
 
 from langchain_core.callbacks import (
@@ -155,7 +155,6 @@ def _default_params(self) -> dict[str, Any]:
         }
 
         return {
-            "model": self.model_name,
             **{k: v for k, v in exclude_if_none.items() if v is not None},
             **self.model_kwargs,
         }
@@ -329,9 +328,7 @@ async def _uipath_agenerate(
         response = await self.uipath_arequest(request_body=request_body, raise_status_error=True)
         return self._postprocess_response(response.json())
 
-    def _generate_chunk(
-        self, original_message: str, json_data: dict[str, Any]
-    ) -> ChatGenerationChunk:
+    def _generate_chunk(self, json_data: dict[str, Any]) -> ChatGenerationChunk:
         generation_info = {
             "id": json_data.get("id"),
             "created": json_data.get("created", ""),
@@ -377,10 +374,10 @@ def _generate_chunk(
                 )
 
         return ChatGenerationChunk(
-            text=original_message,
+            text=content or "",
             generation_info=generation_info,
             message=AIMessageChunk(
-                content=content,
+                content=content or "",
                 usage_metadata=usage_metadata,
                 tool_call_chunks=tool_call_chunks,
             ),
@@ -392,40 +389,42 @@ def _uipath_stream(
         stop: list[str] | None = None,
         run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
-    ) -> Iterator[ChatGenerationChunk]:
+    ) -> Generator[ChatGenerationChunk, None, None]:
         request_body = self._preprocess_request(messages, stop=stop, **kwargs)
+        request_body["stream"] = True
         for chunk in self.uipath_stream(
             request_body=request_body, stream_type="lines", raise_status_error=True
         ):
             chunk = str(chunk)
             if chunk.startswith("data:"):
-                chunk = chunk.split("data:")[1].strip()
+                chunk = chunk[len("data:") :].strip()
             try:
                 json_data = json.loads(chunk)
             except json.JSONDecodeError:
                 continue
             if "id" in json_data and not json_data["id"]:
                 continue
-            yield self._generate_chunk(chunk, json_data)
+            yield self._generate_chunk(json_data)
 
     async def _uipath_astream(
         self,
         messages: list[BaseMessage],
         stop: list[str] | None = None,
         run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
-    ) -> AsyncIterator[ChatGenerationChunk]:
+    ) -> AsyncGenerator[ChatGenerationChunk, None]:
         request_body = self._preprocess_request(messages, stop=stop, **kwargs)
+        request_body["stream"] = True
         async for chunk in self.uipath_astream(
             request_body=request_body, stream_type="lines", raise_status_error=True
         ):
             chunk = str(chunk)
             if chunk.startswith("data:"):
-                chunk = chunk.split("data:")[1].strip()
+                chunk = chunk[len("data:") :].strip()
             try:
                 json_data = json.loads(chunk)
             except json.JSONDecodeError:
                 continue
             if "id" in json_data and not json_data["id"]:
                 continue
-            yield self._generate_chunk(chunk, json_data)
+            yield self._generate_chunk(json_data)

packages/uipath_langchain_client/src/uipath_langchain_client/clients/normalized/embeddings.py

@@ -18,14 +18,20 @@ class UiPathEmbeddings(UiPathBaseEmbeddings, Embeddings):
     )
 
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
-        response = self.uipath_request(request_body={"input": texts})
+        response = self.uipath_request(
+            request_body={"input": texts},
+            raise_status_error=True,
+        )
         return [r["embedding"] for r in response.json()["data"]]
 
     def embed_query(self, text: str) -> list[float]:
         return self.embed_documents([text])[0]
 
     async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
-        response = await self.uipath_arequest(request_body={"input": texts})
+        response = await self.uipath_arequest(
+            request_body={"input": texts},
+            raise_status_error=True,
+        )
         return [r["embedding"] for r in response.json()["data"]]
 
     async def aembed_query(self, text: str) -> list[float]:

packages/uipath_langchain_client/src/uipath_langchain_client/clients/vertexai/chat_models.py

@@ -41,7 +41,7 @@ def setup_uipath_client(self) -> Self:
             project_id="PLACEHOLDER",
             access_token="PLACEHOLDER",
             base_url=str(self.uipath_sync_client.base_url),
-            default_headers=self.uipath_sync_client.headers,
+            default_headers=dict(self.uipath_sync_client.headers),
             max_retries=0,  # handled by the UiPath client
             http_client=self.uipath_sync_client,
         )
@@ -50,7 +50,7 @@ def setup_uipath_client(self) -> Self:
             project_id="PLACEHOLDER",
             access_token="PLACEHOLDER",
             base_url=str(self.uipath_async_client.base_url),
-            default_headers=self.uipath_async_client.headers,
+            default_headers=dict(self.uipath_async_client.headers),
             max_retries=0,  # handled by the UiPath client
             http_client=self.uipath_async_client,
         )