feat: address PR comments

Author: ionut-mihalache-uipath

src/uipath_langchain/agent/react/agent.py

@@ -8,6 +8,8 @@
 from pydantic import BaseModel
 from uipath.platform.guardrails import BaseGuardrail
 
+from uipath_langchain.chat.types import UiPathPassthroughChatModel
+
 from ..guardrails.actions import GuardrailAction
 from .guardrails.guardrails_subgraph import (
     create_agent_init_guardrails_subgraph,
@@ -33,6 +35,7 @@
     AgentGraphConfig,
     AgentGraphNode,
     AgentGraphState,
+    AgentSettings,
 )
 from .utils import create_state_with_input
 
@@ -62,6 +65,17 @@ def create_agent(
     """
     from ..tools import create_tool_node
 
+    if not isinstance(model, UiPathPassthroughChatModel):
+        raise TypeError(
+            f"Model {type(model).__name__} does not implement UiPathPassthroughChatModel. "
+            "The model must have llm_provider and api_flavor properties."
+        )
+
+    agent_settings = AgentSettings(
+        llm_provider=model.llm_provider,
+        api_flavor=model.api_flavor,
+    )
+
     if config is None:
         config = AgentGraphConfig()
 
@@ -71,7 +85,9 @@ def create_agent(
     )
     llm_tools: list[BaseTool] = [*agent_tools, *flow_control_tools]
 
-    init_node = create_init_node(messages, input_schema, config.is_conversational)
+    init_node = create_init_node(
+        messages, input_schema, config.is_conversational, agent_settings
+    )
 
     tool_nodes = create_tool_node(agent_tools)
     tool_nodes_with_guardrails = create_tools_guardrails_subgraph(

src/uipath_langchain/agent/react/init_node.py

@@ -9,13 +9,15 @@
 from .job_attachments import (
     get_job_attachments,
 )
+from .types import AgentSettings
 
 
 def create_init_node(
     messages: Sequence[SystemMessage | HumanMessage]
     | Callable[[Any], Sequence[SystemMessage | HumanMessage]],
     input_schema: type[BaseModel] | None,
     is_conversational: bool = False,
+    agent_settings: AgentSettings | None = None,
 ):
     def graph_state_init(state: Any) -> Any:
         resolved_messages: Sequence[SystemMessage | HumanMessage] | Overwrite
@@ -46,6 +48,7 @@ def graph_state_init(state: Any) -> Any:
             "messages": resolved_messages,
             "inner_state": {
                 "job_attachments": job_attachments_dict,
+                "agent_settings": agent_settings,
             },
         }
 

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

src/uipath_langchain/agent/react/llm_with_files.py

@@ -0,0 +1,16 @@
+"""Multimodal LLM input handling (images, PDFs, etc.)."""
+
+from .invoke import build_file_content_block, llm_call_with_files
+from .types import IMAGE_MIME_TYPES, FileInfo
+from .utils import download_file_base64, is_image, is_pdf, sanitize_filename
+
+__all__ = [
+    "FileInfo",
+    "IMAGE_MIME_TYPES",
+    "build_file_content_block",
+    "download_file_base64",
+    "is_image",
+    "is_pdf",
+    "llm_call_with_files",
+    "sanitize_filename",
+]

src/uipath_langchain/agent/react/multimodal/__init__.py

@@ -0,0 +1,86 @@
+"""LLM invocation with multimodal file attachments."""
+
+from typing import Any
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import (
+    AIMessage,
+    AnyMessage,
+    DataContentBlock,
+    HumanMessage,
+)
+from langchain_core.messages.content import create_file_block, create_image_block
+
+from .types import FileInfo
+from .utils import download_file_base64, is_image, is_pdf, sanitize_filename
+
+
+async def build_file_content_block(
+    file_info: FileInfo,
+) -> DataContentBlock:
+    """Build a LangChain content block for a file attachment.
+
+    Args:
+        file_info: File URL, name, and MIME type.
+
+    Returns:
+        A DataContentBlock for the file (image or PDF).
+
+    Raises:
+        ValueError: If the MIME type is not supported.
+    """
+    base64_file = await download_file_base64(file_info.url)
+
+    if is_image(file_info.mime_type):
+        return create_image_block(base64=base64_file, mime_type=file_info.mime_type)
+    if is_pdf(file_info.mime_type):
+        return create_file_block(
+            base64=base64_file,
+            mime_type=file_info.mime_type,
+            filename=sanitize_filename(file_info.name),
+        )
+
+    raise ValueError(f"Unsupported mime_type={file_info.mime_type}")
+
+
+async def llm_call_with_files(
+    messages: list[AnyMessage],
+    files: list[FileInfo],
+    model: BaseChatModel,
+) -> AIMessage:
+    """Invoke an LLM with file attachments.
+
+    Downloads files, creates content blocks, and appends them as a HumanMessage.
+    If no files are provided, equivalent to model.ainvoke().
+
+    Args:
+        messages: The conversation messages to send to the LLM.
+        files: List of file attachments to include.
+        model: The LLM model to invoke.
+
+    Returns:
+        The AIMessage response from the LLM.
+
+    Raises:
+        TypeError: If the LLM returns something other than AIMessage.
+    """
+    if not files:
+        response = await model.ainvoke(messages)
+        if not isinstance(response, AIMessage):
+            raise TypeError(
+                f"LLM returned {type(response).__name__} instead of AIMessage"
+            )
+        return response
+
+    content_blocks: list[Any] = []
+    for file_info in files:
+        content_block = await build_file_content_block(file_info)
+        content_blocks.append(content_block)
+
+    file_message = HumanMessage(content_blocks=content_blocks)
+    all_messages = list(messages) + [file_message]
+
+    response = await model.ainvoke(all_messages)
+    if not isinstance(response, AIMessage):
+        raise TypeError(f"LLM returned {type(response).__name__} instead of AIMessage")
+    return response

src/uipath_langchain/agent/react/multimodal/invoke.py

@@ -0,0 +1,20 @@
+"""Types and constants for multimodal LLM input handling."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class FileInfo:
+    """File information for LLM file attachments."""
+
+    url: str
+    name: str
+    mime_type: str
+
+
+IMAGE_MIME_TYPES: set[str] = {
+    "image/png",
+    "image/jpeg",
+    "image/gif",
+    "image/webp",
+}

src/uipath_langchain/agent/react/multimodal/types.py

@@ -0,0 +1,45 @@
+"""Utility functions for multimodal file handling."""
+
+import base64
+import re
+
+import httpx
+from uipath._utils._ssl_context import get_httpx_client_kwargs
+
+from .types import IMAGE_MIME_TYPES
+
+
+def sanitize_filename(filename: str) -> str:
+    """Sanitize a filename to conform to provider document naming requirements.
+
+    Bedrock only allows: alphanumeric characters, whitespace, hyphens,
+    parentheses, and square brackets. No consecutive whitespace allowed.
+    """
+    if not filename or filename.isspace():
+        return "document"
+
+    sanitized = re.sub(r"[^a-zA-Z0-9\s\-\(\)\[\]]", "-", filename)
+    sanitized = re.sub(r"\s+", " ", sanitized)
+    sanitized = re.sub(r"-+", "-", sanitized)
+    sanitized = sanitized.strip(" -")
+
+    return sanitized if sanitized else "document"
+
+
+def is_pdf(mime_type: str) -> bool:
+    """Check if the MIME type represents a PDF document."""
+    return mime_type.lower() == "application/pdf"
+
+
+def is_image(mime_type: str) -> bool:
+    """Check if the MIME type represents a supported image format."""
+    return mime_type.lower() in IMAGE_MIME_TYPES
+
+
+async def download_file_base64(url: str) -> str:
+    """Download a file from a URL and return its content as a base64 string."""
+    async with httpx.AsyncClient(**get_httpx_client_kwargs()) as client:
+        response = await client.get(url)
+        response.raise_for_status()
+        file_content = response.content
+    return base64.b64encode(file_content).decode("utf-8")