feat: address PR comments

ionut-mihalache-uipath · ionut-mihalache-uipath · commit f4bc8e1c85b8 · 2026-01-20T12:31:17.000+02:00
diff --git a/src/uipath_langchain/agent/react/llm_node.py b/src/uipath_langchain/agent/react/llm_node.py
@@ -1,53 +1,22 @@
 """LLM node for ReAct Agent graph."""
 
-from typing import Any, Sequence
+from typing import Sequence
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, AnyMessage, ToolCall
 from langchain_core.tools import BaseTool
 from uipath.runtime.errors import UiPathErrorCategory, UiPathErrorCode
 
+from uipath_langchain.llm import get_payload_handler
+
 from ..exceptions import AgentTerminationException
 from .constants import (
     DEFAULT_MAX_CONSECUTIVE_THINKING_MESSAGES,
     DEFAULT_MAX_LLM_MESSAGES,
 )
 from .types import FLOW_CONTROL_TOOLS, AgentGraphState
-from uipath_langchain.chat.types import APIFlavor
-
-from .constants import MAX_CONSECUTIVE_THINKING_MESSAGES
-from .types import AgentGraphState
 from .utils import count_consecutive_thinking_messages
 
-OPENAI_COMPATIBLE_CHAT_MODELS = (
-    "UiPathChatOpenAI",
-    "AzureChatOpenAI",
-    "ChatOpenAI",
-    "UiPathChat",
-    "UiPathAzureChatOpenAI",
-)
-
-
-def _get_required_tool_choice_by_model(
-    model: BaseChatModel,
-) -> str | dict[str, Any]:
-    """Get the appropriate tool_choice value to enforce tool usage based on model type.
-
-    Returns:
-        - "required" for OpenAI compatible models
-        - "any" for Bedrock Converse and Vertex models (string format)
-        - {"type": "any"} for Bedrock Invoke API (dict format required)
-    """
-    model_class_name = model.__class__.__name__
-    if model_class_name in OPENAI_COMPATIBLE_CHAT_MODELS:
-        return "required"
-
-    api_flavor = getattr(model, "api_flavor", None)
-    if api_flavor == APIFlavor.AWS_BEDROCK_INVOKE:
-        return {"type": "any"}
-
-    return "any"
-
 
 def _filter_control_flow_tool_calls(
     tool_calls: list[ToolCall],
@@ -81,7 +50,8 @@ def create_llm_node(
     """
     bindable_tools = list(tools) if tools else []
     base_llm = model.bind_tools(bindable_tools) if bindable_tools else model
-    tool_choice_required_value = _get_required_tool_choice_by_model(model)
+    payload_handler = get_payload_handler(model)
+    tool_choice_required_value = payload_handler.get_required_tool_choice()
 
     async def llm_node(state: AgentGraphState):
         messages: list[AnyMessage] = state.messages
diff --git a/src/uipath_langchain/chat/__init__.py b/src/uipath_langchain/chat/__init__.py
@@ -29,7 +29,7 @@ def __getattr__(name):
         from . import supported_models
 
         return getattr(supported_models, name)
-    if name in ("LLMProvider", "APIFlavor"):
+    if name in ("LLMProvider", "APIFlavor", "UiPathChatModel"):
         from . import types
 
         return getattr(types, name)
@@ -40,6 +40,7 @@ def __getattr__(name):
     "UiPathChat",
     "UiPathAzureChatOpenAI",
     "UiPathChatOpenAI",
+    "UiPathChatModel",
     "OpenAIModels",
     "BedrockModels",
     "GeminiModels",
diff --git a/src/uipath_langchain/chat/models.py b/src/uipath_langchain/chat/models.py
@@ -29,9 +29,6 @@
 class UiPathAzureChatOpenAI(UiPathRequestMixin, AzureChatOpenAI):
     """Custom LLM connector for LangChain integration with UiPath."""
 
-    llm_provider: LLMProvider = LLMProvider.OPENAI
-    api_flavor: APIFlavor = APIFlavor.OPENAI_COMPLETIONS
-
     def _generate(
         self,
         messages: list[BaseMessage],
@@ -174,9 +171,6 @@ def endpoint(self) -> str:
 class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
     """Custom LLM connector for LangChain integration with UiPath Normalized."""
 
-    llm_provider: LLMProvider = LLMProvider.OPENAI
-    api_flavor: APIFlavor = APIFlavor.OPENAI_COMPLETIONS
-
     def _create_chat_result(
         self,
         response: Union[dict[str, Any], BaseModel],
diff --git a/src/uipath_langchain/chat/types.py b/src/uipath_langchain/chat/types.py
@@ -1,4 +1,5 @@
 from enum import StrEnum
+from typing import Protocol, runtime_checkable
 
 
 class LLMProvider(StrEnum):
@@ -18,3 +19,23 @@ class APIFlavor(StrEnum):
     AWS_BEDROCK_INVOKE = "AwsBedrockInvoke"
     VERTEX_GEMINI_GENERATE_CONTENT = "GeminiGenerateContent"
     VERTEX_ANTHROPIC_CLAUDE = "AnthropicClaude"
+
+
+@runtime_checkable
+class UiPathChatModel(Protocol):
+    """Protocol for UiPath chat models with provider and flavor information.
+
+    All UiPath chat model classes (UiPathChatOpenAI, UiPathChatBedrock,
+    UiPathChatBedrockConverse, UiPathChatVertex, UiPathChat, UiPathAzureChatOpenAI)
+    implement this protocol.
+    """
+
+    @property
+    def llm_provider(self) -> LLMProvider:
+        """The LLM provider for this model."""
+        ...
+
+    @property
+    def api_flavor(self) -> APIFlavor:
+        """The API flavor for this model."""
+        ...
diff --git a/src/uipath_langchain/llm/__init__.py b/src/uipath_langchain/llm/__init__.py
@@ -1,12 +1,19 @@
 """Common LLM provider utilities."""
 
+from uipath_langchain.chat.types import UiPathChatModel
+
 from .builders import MessageContentBuilder
 from .content_builder import get_content_builder
-from .provider import get_api_flavor, get_llm_provider
+from .handlers import ModelPayloadHandler
+from .payload_handler import get_payload_handler
+from .utils import get_api_flavor, get_llm_provider
 
 __all__ = [
     "MessageContentBuilder",
+    "ModelPayloadHandler",
+    "UiPathChatModel",
     "get_api_flavor",
     "get_content_builder",
     "get_llm_provider",
+    "get_payload_handler",
 ]
diff --git a/src/uipath_langchain/llm/builders/base.py b/src/uipath_langchain/llm/builders/base.py
@@ -1,13 +1,9 @@
 """Abstract base class for message content builders."""
 
-import base64
 import re
 from abc import ABC, abstractmethod
 from typing import Any
 
-import httpx
-from uipath._utils._ssl_context import get_httpx_client_kwargs
-
 IMAGE_MIME_TYPES: set[str] = {
     "image/png",
     "image/jpeg",
@@ -38,20 +34,6 @@ def sanitize_filename_for_anthropic(filename: str) -> str:
     return sanitized if sanitized else "document"
 
 
-async def download_file_bytes(url: str) -> bytes:
-    """Download a file from a URL and return its content bytes."""
-    async with httpx.AsyncClient(**get_httpx_client_kwargs()) as client:
-        response = await client.get(url)
-        response.raise_for_status()
-        return response.content
-
-
-async def download_file_base64(url: str) -> str:
-    """Download a file from a URL and return its content as a base64 string."""
-    file_content = await download_file_bytes(url)
-    return base64.b64encode(file_content).decode("utf-8")
-
-
 class MessageContentBuilder(ABC):
     """Abstract base class for building provider-specific message content parts."""
 
diff --git a/src/uipath_langchain/llm/builders/bedrock_converse.py b/src/uipath_langchain/llm/builders/bedrock_converse.py
@@ -4,12 +4,11 @@
 
 from .base import (
     MessageContentBuilder,
-    download_file_base64,
-    download_file_bytes,
     is_image,
     is_pdf,
     sanitize_filename_for_anthropic,
 )
+from .download import download_file_base64, download_file_bytes
 
 
 class BedrockConverseBuilder(MessageContentBuilder):
diff --git a/src/uipath_langchain/llm/builders/bedrock_invoke.py b/src/uipath_langchain/llm/builders/bedrock_invoke.py
@@ -2,12 +2,8 @@
 
 from typing import Any
 
-from .base import (
-    MessageContentBuilder,
-    download_file_base64,
-    is_image,
-    is_pdf,
-)
+from .base import MessageContentBuilder, is_image, is_pdf
+from .download import download_file_base64
 
 
 class BedrockInvokeBuilder(MessageContentBuilder):
diff --git a/src/uipath_langchain/llm/builders/download.py b/src/uipath_langchain/llm/builders/download.py
@@ -0,0 +1,20 @@
+"""Utilities for downloading files from URLs."""
+
+import base64
+
+import httpx
+from uipath._utils._ssl_context import get_httpx_client_kwargs
+
+
+async def download_file_bytes(url: str) -> bytes:
+    """Download a file from a URL and return its content bytes."""
+    async with httpx.AsyncClient(**get_httpx_client_kwargs()) as client:
+        response = await client.get(url)
+        response.raise_for_status()
+        return response.content
+
+
+async def download_file_base64(url: str) -> str:
+    """Download a file from a URL and return its content as a base64 string."""
+    file_content = await download_file_bytes(url)
+    return base64.b64encode(file_content).decode("utf-8")
diff --git a/src/uipath_langchain/llm/builders/openai_completions.py b/src/uipath_langchain/llm/builders/openai_completions.py
@@ -2,12 +2,8 @@
 
 from typing import Any
 
-from .base import (
-    MessageContentBuilder,
-    download_file_base64,
-    is_image,
-    is_pdf,
-)
+from .base import MessageContentBuilder, is_image, is_pdf
+from .download import download_file_base64
 
 
 class OpenAICompletionsBuilder(MessageContentBuilder):
diff --git a/src/uipath_langchain/llm/builders/openai_responses.py b/src/uipath_langchain/llm/builders/openai_responses.py
@@ -2,12 +2,8 @@
 
 from typing import Any
 
-from .base import (
-    MessageContentBuilder,
-    download_file_base64,
-    is_image,
-    is_pdf,
-)
+from .base import MessageContentBuilder, is_image, is_pdf
+from .download import download_file_base64
 
 
 class OpenAIResponsesBuilder(MessageContentBuilder):
diff --git a/src/uipath_langchain/llm/builders/vertex_gemini.py b/src/uipath_langchain/llm/builders/vertex_gemini.py
@@ -2,12 +2,8 @@
 
 from typing import Any
 
-from .base import (
-    MessageContentBuilder,
-    download_file_base64,
-    is_image,
-    is_pdf,
-)
+from .base import MessageContentBuilder, is_image, is_pdf
+from .download import download_file_base64
 
 
 class VertexGeminiBuilder(MessageContentBuilder):
diff --git a/src/uipath_langchain/llm/content_builder.py b/src/uipath_langchain/llm/content_builder.py
@@ -1,8 +1,6 @@
 """Factory for creating message content builders."""
 
-from langchain_core.language_models import BaseChatModel
-
-from uipath_langchain.chat.types import APIFlavor, LLMProvider
+from uipath_langchain.chat.types import APIFlavor, LLMProvider, UiPathChatModel
 
 from .builders import (
     BedrockConverseBuilder,
@@ -12,7 +10,6 @@
     OpenAIResponsesBuilder,
     VertexGeminiBuilder,
 )
-from .provider import get_api_flavor, get_llm_provider
 
 _BUILDER_REGISTRY: dict[tuple[LLMProvider, APIFlavor], type[MessageContentBuilder]] = {
     (LLMProvider.BEDROCK, APIFlavor.AWS_BEDROCK_CONVERSE): BedrockConverseBuilder,
@@ -23,20 +20,20 @@
 }
 
 
-def get_content_builder(model: BaseChatModel) -> MessageContentBuilder:
+def get_content_builder(model: UiPathChatModel) -> MessageContentBuilder:
     """Get the appropriate content builder for a model.
 
     Args:
-        model: The LLM model instance.
+        model: A UiPath chat model instance with llm_provider and api_flavor.
 
     Returns:
         A MessageContentBuilder instance for the model's provider and API flavor.
 
     Raises:
         ValueError: If no builder is registered for the model's provider/API flavor.
     """
-    provider = get_llm_provider(model)
-    api_flavor = get_api_flavor(model)
+    provider = model.llm_provider
+    api_flavor = model.api_flavor
 
     key = (provider, api_flavor)
     builder_class = _BUILDER_REGISTRY.get(key)
diff --git a/src/uipath_langchain/llm/handlers/__init__.py b/src/uipath_langchain/llm/handlers/__init__.py
@@ -0,0 +1,17 @@
+"""Model payload handlers for different LLM providers and API flavors."""
+
+from .base import ModelPayloadHandler
+from .bedrock_converse import BedrockConversePayloadHandler
+from .bedrock_invoke import BedrockInvokePayloadHandler
+from .openai_completions import OpenAICompletionsPayloadHandler
+from .openai_responses import OpenAIResponsesPayloadHandler
+from .vertex_gemini import VertexGeminiPayloadHandler
+
+__all__ = [
+    "ModelPayloadHandler",
+    "BedrockConversePayloadHandler",
+    "BedrockInvokePayloadHandler",
+    "OpenAICompletionsPayloadHandler",
+    "OpenAIResponsesPayloadHandler",
+    "VertexGeminiPayloadHandler",
+]
diff --git a/src/uipath_langchain/llm/handlers/base.py b/src/uipath_langchain/llm/handlers/base.py
@@ -0,0 +1,22 @@
+"""Abstract base class for model payload handlers."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class ModelPayloadHandler(ABC):
+    """Abstract base class for handling model-specific LLM parameters.
+
+    Each handler provides provider-specific parameter values for LLM operations.
+    """
+
+    @abstractmethod
+    def get_required_tool_choice(self) -> str | dict[str, Any]:
+        """Get the tool_choice value that enforces tool usage.
+
+        Returns:
+            Provider-specific value to force tool usage:
+            - "required" for OpenAI-compatible models
+            - "any" for Bedrock Converse and Vertex models (string format)
+            - {"type": "any"} for Bedrock Invoke API (dict format required)
+        """
diff --git a/src/uipath_langchain/llm/handlers/bedrock_converse.py b/src/uipath_langchain/llm/handlers/bedrock_converse.py
@@ -0,0 +1,13 @@
+"""Bedrock Converse payload handler."""
+
+from typing import Any
+
+from .base import ModelPayloadHandler
+
+
+class BedrockConversePayloadHandler(ModelPayloadHandler):
+    """Payload handler for AWS Bedrock Converse API."""
+
+    def get_required_tool_choice(self) -> str | dict[str, Any]:
+        """Get tool_choice value for Bedrock Converse API."""
+        return "any"
diff --git a/src/uipath_langchain/llm/handlers/bedrock_invoke.py b/src/uipath_langchain/llm/handlers/bedrock_invoke.py
@@ -0,0 +1,13 @@
+"""Bedrock Invoke payload handler."""
+
+from typing import Any
+
+from .base import ModelPayloadHandler
+
+
+class BedrockInvokePayloadHandler(ModelPayloadHandler):
+    """Payload handler for AWS Bedrock Invoke API."""
+
+    def get_required_tool_choice(self) -> str | dict[str, Any]:
+        """Get tool_choice value for Bedrock Invoke API."""
+        return {"type": "any"}
diff --git a/src/uipath_langchain/llm/handlers/openai_completions.py b/src/uipath_langchain/llm/handlers/openai_completions.py
diff --git a/src/uipath_langchain/llm/handlers/openai_responses.py b/src/uipath_langchain/llm/handlers/openai_responses.py
diff --git a/src/uipath_langchain/llm/handlers/vertex_gemini.py b/src/uipath_langchain/llm/handlers/vertex_gemini.py
diff --git a/src/uipath_langchain/llm/payload_handler.py b/src/uipath_langchain/llm/payload_handler.py
diff --git a/src/uipath_langchain/llm/utils.py b/src/uipath_langchain/llm/utils.py
