feat: add file support for more LLM API Flavors

Author: ionut-mihalache-uipath

src/uipath_langchain/agent/react/file_type_handler.py

@@ -1,11 +1,13 @@
 import base64
 import re
-from enum import StrEnum
 from typing import Any
 
 import httpx
+from langchain_core.language_models import BaseChatModel
 from uipath._utils._ssl_context import get_httpx_client_kwargs
 
+from uipath_langchain.chat.types import APIFlavor, LLMProvider
+
 IMAGE_MIME_TYPES: set[str] = {
     "image/png",
     "image/jpeg",
@@ -14,13 +16,6 @@
 }
 
 
-class LlmProvider(StrEnum):
-    OPENAI = "openai"
-    BEDROCK = "bedrock"
-    VERTEX = "vertex"
-    UNKNOWN = "unknown"
-
-
 def is_pdf(mime_type: str) -> bool:
     """Check if the MIME type represents a PDF document."""
     return mime_type.lower() == "application/pdf"
@@ -31,23 +26,28 @@ def is_image(mime_type: str) -> bool:
     return mime_type.lower() in IMAGE_MIME_TYPES
 
 
-def detect_provider(model_name: str) -> LlmProvider:
-    """Detect the LLM provider (Bedrock, OpenAI, or Vertex) based on the model name."""
-    if not model_name:
-        raise ValueError(f"Unsupported model: {model_name}")
+def get_llm_provider(model: BaseChatModel) -> LLMProvider:
+    """Get the LLM provider from a model instance.
 
-    model_lower = model_name.lower()
+    Checks for the llm_provider attribute on UiPath LLM classes first,
+    then falls back to model name string matching for other models.
+    """
+    if hasattr(model, "llm_provider") and isinstance(model.llm_provider, LLMProvider):
+        return model.llm_provider
 
-    if "anthropic" in model_lower or "claude" in model_lower:
-        return LlmProvider.BEDROCK
+    raise ValueError(
+        f"Can't determine llm_provider in file_type_handler for model={model}"
+    )
 
-    if "gpt" in model_lower:
-        return LlmProvider.OPENAI
 
-    if "gemini" in model_lower:
-        return LlmProvider.VERTEX
+def get_api_flavor(model: BaseChatModel) -> APIFlavor:
+    """Get the API flavor from a model instance."""
+    if hasattr(model, "api_flavor") and isinstance(model.api_flavor, APIFlavor):
+        return model.api_flavor
 
-    raise ValueError(f"Unsupported model: {model_name}")
+    raise ValueError(
+        f"Can't determine api_flavor in file_type_handler for model={model}"
+    )
 
 
 def sanitize_filename_for_anthropic(filename: str) -> str:
@@ -85,24 +85,30 @@ async def build_message_content_part_from_data(
     url: str,
     filename: str,
     mime_type: str,
-    model: str,
+    model: BaseChatModel,
 ) -> dict[str, Any]:
     """Download a file and build a provider-specific message content part.
 
     The format varies based on the detected provider (Bedrock, OpenAI, or Vertex).
+    Uses model.llm_provider and model.api_flavor attributes.
     """
-    provider = detect_provider(model)
+    provider = get_llm_provider(model)
+    api_flavor = get_api_flavor(model)
 
-    if provider == LlmProvider.BEDROCK:
-        return await _build_bedrock_content_part_from_data(url, mime_type, filename)
+    if provider == LLMProvider.BEDROCK:
+        return await _build_bedrock_content_part_from_data(
+            url, mime_type, filename, api_flavor
+        )
 
-    if provider == LlmProvider.OPENAI:
+    if provider == LLMProvider.OPENAI:
         return await _build_openai_content_part_from_data(
-            url, mime_type, filename, False
+            url, mime_type, filename, True, api_flavor
         )
 
-    if provider == LlmProvider.VERTEX:
-        return await _build_vertex_content_part_from_data(url, mime_type, False)
+    if provider == LLMProvider.VERTEX:
+        return await _build_vertex_content_part_from_data(
+            url, mime_type, True, api_flavor
+        )
 
     raise ValueError(f"Unsupported provider: {provider}")
 
@@ -111,8 +117,26 @@ async def _build_bedrock_content_part_from_data(
     url: str,
     mime_type: str,
     filename: str,
+    api_flavor: APIFlavor,
 ) -> dict[str, Any]:
-    """Build a content part for AWS Bedrock (Anthropic Claude models)."""
+    """Build a content part for AWS Bedrock (Anthropic Claude models).
+
+    Converse API uses raw bytes, Invoke API uses base64-encoded content.
+    """
+    if api_flavor == APIFlavor.AWS_BEDROCK_CONVERSE:
+        return await _build_bedrock_converse_content_part(url, mime_type, filename)
+    elif api_flavor == APIFlavor.AWS_BEDROCK_INVOKE:
+        return await _build_bedrock_invoke_content_part(url, mime_type, filename)
+    else:
+        raise ValueError(f"Unsupported Bedrock api_flavor: {api_flavor}")
+
+
+async def _build_bedrock_converse_content_part(
+    url: str,
+    mime_type: str,
+    filename: str,
+) -> dict[str, Any]:
+    """Build content part for Bedrock Converse API (uses raw bytes)."""
     if is_pdf(mime_type):
         file_bytes = await _download_file_bytes(url)
         name = filename.rsplit(".", 1)[0] if "." in filename else filename
@@ -143,39 +167,109 @@ async def _build_bedrock_content_part_from_data(
     raise ValueError(f"Unsupported mime_type: {mime_type}")
 
 
+async def _build_bedrock_invoke_content_part(
+    url: str,
+    mime_type: str,
+    filename: str,
+) -> dict[str, Any]:
+    """Build content part for Bedrock Invoke API (uses base64-encoded content)."""
+    base64_content = await _download_file(url)
+
+    if is_pdf(mime_type):
+        return {
+            "type": "document",
+            "source": {
+                "type": "base64",
+                "media_type": mime_type,
+                "data": base64_content,
+            },
+        }
+
+    if is_image(mime_type):
+        return {
+            "type": "image",
+            "source": {
+                "type": "base64",
+                "media_type": mime_type,
+                "data": base64_content,
+            },
+        }
+
+    raise ValueError(f"Unsupported mime_type: {mime_type}")
+
+
 async def _build_openai_content_part_from_data(
     url: str,
     mime_type: str,
     filename: str,
-    download_image: bool,
+    download_file: bool,
+    api_flavor: APIFlavor,
 ) -> dict[str, Any]:
-    """Build a content part for OpenAI models (base64-encoded or URL reference)."""
-    if download_image:
-        base64_content = await _download_file(url)
-        if is_image(mime_type):
-            data_url = f"data:{mime_type};base64,{base64_content}"
-            return {
-                "type": "input_image",
-                "image_url": data_url,
-            }
+    """Build a content part for OpenAI models"""
+    if api_flavor == APIFlavor.OPENAI_RESPONSES:
+        return await _build_openai_responses_content_part(
+            url, mime_type, filename, download_file
+        )
+    elif api_flavor == APIFlavor.OPENAI_COMPLETIONS:
+        return await _build_openai_completions_content_part(
+            url, mime_type, filename, download_file
+        )
+    else:
+        raise ValueError(f"Unsupported OpenAI api_flavor: {api_flavor}")
+
+
+async def _build_openai_responses_content_part(
+    url: str,
+    mime_type: str,
+    filename: str,
+    download_file: bool,
+) -> dict[str, Any]:
+    """Build content part for OpenAI Responses API."""
+    if download_file:
+        return await _build_openai_responses_downloaded(url, mime_type, filename)
+    return _build_openai_responses_from_url(url, mime_type)
 
-        if is_pdf(mime_type):
-            data = f"data:application/pdf;base64,{base64_content}"
-            return {
-                "type": "file",
-                "file": {
-                    "filename": filename,
-                    "file_data": data,
-                },
-            }
 
-    elif is_image(mime_type):
+async def _build_openai_responses_downloaded(
+    url: str,
+    mime_type: str,
+    filename: str,
+) -> dict[str, Any]:
+    """Build content part for OpenAI Responses API with downloaded file."""
+    base64_content = await _download_file(url)
+
+    if is_image(mime_type):
+        data_url = f"data:{mime_type};base64,{base64_content}"
+        return {
+            "type": "input_image",
+            "image_url": data_url,
+        }
+
+    if is_pdf(mime_type):
+        data = f"data:application/pdf;base64,{base64_content}"
+        return {
+            "type": "file",
+            "file": {
+                "filename": filename,
+                "file_data": data,
+            },
+        }
+
+    raise ValueError(f"Unsupported mime_type: {mime_type}")
+
+
+def _build_openai_responses_from_url(
+    url: str,
+    mime_type: str,
+) -> dict[str, Any]:
+    """Build content part for OpenAI Responses API with URL reference."""
+    if is_image(mime_type):
         return {
             "type": "input_image",
             "image_url": url,
         }
 
-    elif is_pdf(mime_type):
+    if is_pdf(mime_type):
         return {
             "type": "input_file",
             "file_url": url,
@@ -184,12 +278,68 @@ async def _build_openai_content_part_from_data(
     raise ValueError(f"Unsupported mime_type: {mime_type}")
 
 
+async def _build_openai_completions_content_part(
+    url: str,
+    mime_type: str,
+    filename: str,
+    download_file: bool,
+) -> dict[str, Any]:
+    """Build content part for OpenAI Completions API."""
+    if download_file:
+        return await _build_openai_completions_downloaded(url, mime_type, filename)
+    return await _build_openai_completions_from_url(url, mime_type, filename)
+
+
+async def _build_openai_completions_downloaded(
+    url: str,
+    mime_type: str,
+    filename: str,
+) -> dict[str, Any]:
+    """Build content part for OpenAI Completions API with downloaded file."""
+    base64_content = await _download_file(url)
+
+    if is_image(mime_type):
+        data_url = f"data:{mime_type};base64,{base64_content}"
+        return {
+            "type": "image_url",
+            "image_url": {"url": data_url},
+        }
+
+    if is_pdf(mime_type):
+        raise ValueError("PDFs are not supported when using the OpenAi completions API")
+
+    raise ValueError(f"Unsupported mime_type: {mime_type}")
+
+
+async def _build_openai_completions_from_url(
+    url: str,
+    mime_type: str,
+    filename: str,
+) -> dict[str, Any]:
+    """Build content part for OpenAI Completions API with URL reference."""
+    if is_image(mime_type):
+        return {
+            "type": "image_url",
+            "image_url": {"url": url},
+        }
+
+    if is_pdf(mime_type):
+        raise ValueError("PDFs are not supported when using the OpenAi completions API")
+
+    raise ValueError(f"Unsupported mime_type: {mime_type}")
+
+
 async def _build_vertex_content_part_from_data(
     url: str,
     mime_type: str,
     download_file: bool,
+    api_flavor: APIFlavor,
 ) -> dict[str, Any]:
     """Build a content part for Google Vertex AI / Gemini models."""
+
+    if api_flavor != APIFlavor.VERTEX_GEMINI_GENERATE_CONTENT:
+        raise ValueError(f"Unsupported api_flavor={api_flavor} for building file content parts")
+
     if download_file:
         base64_content = await _download_file(url)
         if is_image(mime_type) or is_pdf(mime_type):

src/uipath_langchain/agent/react/llm_node.py

@@ -1,11 +1,13 @@
 """LLM node for ReAct Agent graph."""
 
-from typing import Literal, Sequence
+from typing import Any, Sequence
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, AnyMessage
 from langchain_core.tools import BaseTool
 
+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
@@ -21,15 +23,22 @@
 
 def _get_required_tool_choice_by_model(
     model: BaseChatModel,
-) -> Literal["required", "any"]:
+) -> str | dict[str, Any]:
     """Get the appropriate tool_choice value to enforce tool usage based on model type.
 
-    "required" - OpenAI compatible required tool_choice value
-    "any" - Vertex and Bedrock parameter for required tool_choice value
+    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"