uipath-langchain 0.0.133__py3-none-any.whl → 0.1.24__py3-none-any.whl

uipath_langchain/chat/mapper.py

@@ -0,0 +1,309 @@
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any, cast
+from uuid import uuid4
+
+from langchain_core.messages import (
+    AIMessageChunk,
+    BaseMessage,
+    HumanMessage,
+    TextContentBlock,
+    ToolCallChunk,
+    ToolMessage,
+)
+from pydantic import ValidationError
+from uipath.core.chat import (
+    UiPathConversationContentPartChunkEvent,
+    UiPathConversationContentPartEndEvent,
+    UiPathConversationContentPartEvent,
+    UiPathConversationContentPartStartEvent,
+    UiPathConversationMessage,
+    UiPathConversationMessageEndEvent,
+    UiPathConversationMessageEvent,
+    UiPathConversationMessageStartEvent,
+    UiPathConversationToolCallEndEvent,
+    UiPathConversationToolCallEvent,
+    UiPathConversationToolCallStartEvent,
+    UiPathInlineValue,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class UiPathChatMessagesMapper:
+    """Stateful mapper that converts LangChain messages to UiPath message events.
+
+    Maintains state across multiple message conversions to properly track:
+    - The AI message ID associated with each tool call for proper correlation with ToolMessage
+    """
+
+    def __init__(self):
+        """Initialize the mapper with empty state."""
+        self.tool_call_to_ai_message: dict[str, str] = {}
+        self.seen_message_ids: set[str] = set()
+
+    def _extract_text(self, content: Any) -> str:
+        """Normalize LangGraph message.content to plain text."""
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            return "".join(
+                part.get("text", "")
+                for part in content
+                if isinstance(part, dict) and part.get("type") == "text"
+            )
+        return str(content or "")
+
+    def map_messages(self, messages: list[Any]) -> list[Any]:
+        """Normalize any 'messages' list into LangChain messages.
+
+        - If already BaseMessage instances: return as-is.
+        - If UiPathConversationMessage: convert to HumanMessage.
+        """
+        if not isinstance(messages, list):
+            raise TypeError("messages must be a list")
+
+        if not messages:
+            return []
+
+        first = messages[0]
+
+        # Case 1: already LangChain messages
+        if isinstance(first, BaseMessage):
+            return cast(list[BaseMessage], messages)
+
+        # Case 2: UiPath messages -> convert to HumanMessage
+        if isinstance(first, UiPathConversationMessage):
+            if not all(isinstance(m, UiPathConversationMessage) for m in messages):
+                raise TypeError("Mixed message types not supported")
+            return self._map_messages_internal(
+                cast(list[UiPathConversationMessage], messages)
+            )
+
+        # Case3: List[dict] -> parse to List[UiPathConversationMessage]
+        if isinstance(first, dict):
+            try:
+                parsed_messages = [
+                    UiPathConversationMessage.model_validate(message)
+                    for message in messages
+                ]
+                return self._map_messages_internal(parsed_messages)
+            except ValidationError:
+                pass
+
+        # Fallback: unknown type – just pass through
+        return messages
+
+    def _map_messages_internal(
+        self, messages: list[UiPathConversationMessage]
+    ) -> list[HumanMessage]:
+        """
+        Converts a UiPathConversationMessage into a list of HumanMessages for LangGraph.
+        Supports multimodal content parts (text, external content) and preserves metadata.
+        """
+        human_messages: list[HumanMessage] = []
+
+        for uipath_msg in messages:
+            # Loop over each content part
+            if uipath_msg.content_parts:
+                for part in uipath_msg.content_parts:
+                    data = part.data
+                    content = ""
+                    metadata: dict[str, Any] = {
+                        "message_id": uipath_msg.message_id,
+                        "content_part_id": part.content_part_id,
+                        "mime_type": part.mime_type,
+                        "created_at": uipath_msg.created_at,
+                        "updated_at": uipath_msg.updated_at,
+                    }
+
+                    if isinstance(data, UiPathInlineValue):
+                        content = str(data.inline)
+
+                    # Append a HumanMessage for this content part
+                    human_messages.append(
+                        HumanMessage(content=content, metadata=metadata)
+                    )
+
+            # Handle the case where there are no content parts
+            else:
+                metadata = {
+                    "message_id": uipath_msg.message_id,
+                    "role": uipath_msg.role,
+                    "created_at": uipath_msg.created_at,
+                    "updated_at": uipath_msg.updated_at,
+                }
+                human_messages.append(HumanMessage(content="", metadata=metadata))
+
+        return human_messages
+
+    def map_event(
+        self,
+        message: BaseMessage,
+    ) -> UiPathConversationMessageEvent | None:
+        """Convert LangGraph BaseMessage (chunk or full) into a UiPathConversationMessageEvent.
+
+        Args:
+            message: The LangChain message to convert
+
+        Returns:
+            A UiPathConversationMessageEvent if the message should be emitted, None otherwise.
+        """
+        # Format timestamp as ISO 8601 UTC with milliseconds: 2025-01-04T10:30:00.123Z
+        timestamp = (
+            datetime.now(timezone.utc)
+            .isoformat(timespec="milliseconds")
+            .replace("+00:00", "Z")
+        )
+
+        # --- Streaming AIMessageChunk ---
+        if isinstance(message, AIMessageChunk):
+            if message.id is None:
+                return None
+
+            msg_event = UiPathConversationMessageEvent(
+                message_id=message.id,
+            )
+
+            # Check if this is the last chunk by examining chunk_position
+            if message.chunk_position == "last":
+                msg_event.end = UiPathConversationMessageEndEvent(timestamp=timestamp)
+                msg_event.content_part = UiPathConversationContentPartEvent(
+                    content_part_id=f"chunk-{message.id}-0",
+                    end=UiPathConversationContentPartEndEvent(),
+                )
+                return msg_event
+
+            # For every new message_id, start a new message
+            if message.id not in self.seen_message_ids:
+                self.seen_message_ids.add(message.id)
+                msg_event.start = UiPathConversationMessageStartEvent(
+                    role="assistant", timestamp=timestamp
+                )
+                msg_event.content_part = UiPathConversationContentPartEvent(
+                    content_part_id=f"chunk-{message.id}-0",
+                    start=UiPathConversationContentPartStartEvent(
+                        mime_type="text/plain"
+                    ),
+                )
+
+            elif message.content_blocks:
+                for block in message.content_blocks:
+                    block_type = block.get("type")
+
+                    if block_type == "text":
+                        text_block = cast(TextContentBlock, block)
+                        text = text_block["text"]
+
+                        msg_event.content_part = UiPathConversationContentPartEvent(
+                            content_part_id=f"chunk-{message.id}-0",
+                            chunk=UiPathConversationContentPartChunkEvent(
+                                data=text,
+                                content_part_sequence=0,
+                            ),
+                        )
+
+                    elif block_type == "tool_call_chunk":
+                        tool_chunk_block = cast(ToolCallChunk, block)
+
+                        tool_call_id = tool_chunk_block.get("id")
+                        if tool_call_id:
+                            # Track tool_call_id -> ai_message_id mapping
+                            self.tool_call_to_ai_message[str(tool_call_id)] = message.id
+
+                        args = tool_chunk_block.get("args") or ""
+
+                        msg_event.content_part = UiPathConversationContentPartEvent(
+                            content_part_id=f"chunk-{message.id}-0",
+                            chunk=UiPathConversationContentPartChunkEvent(
+                                data=args,
+                                content_part_sequence=0,
+                            ),
+                        )
+                        # Continue so that multiple tool_call_chunks in the same block list
+                        # are handled correctly
+                        continue
+
+            # Fallback: raw string content on the chunk (rare when using content_blocks)
+            elif isinstance(message.content, str) and message.content:
+                msg_event.content_part = UiPathConversationContentPartEvent(
+                    content_part_id=f"content-{message.id}",
+                    chunk=UiPathConversationContentPartChunkEvent(
+                        data=message.content,
+                        content_part_sequence=0,
+                    ),
+                )
+
+            if (
+                msg_event.start
+                or msg_event.content_part
+                or msg_event.tool_call
+                or msg_event.end
+            ):
+                return msg_event
+
+            return None
+
+        # --- ToolMessage ---
+        if isinstance(message, ToolMessage):
+            # Look up the AI message ID using the tool_call_id
+            result_message_id = (
+                self.tool_call_to_ai_message.get(message.tool_call_id)
+                if message.tool_call_id
+                else None
+            )
+
+            # If no AI message ID was found, we cannot properly associate this tool result
+            if not result_message_id:
+                logger.warning(
+                    f"Tool message {message.tool_call_id} has no associated AI message ID. Skipping."
+                )
+
+            # Clean up the mapping after use
+            if (
+                message.tool_call_id
+                and message.tool_call_id in self.tool_call_to_ai_message
+            ):
+                del self.tool_call_to_ai_message[message.tool_call_id]
+
+            content_value: Any = message.content
+            if isinstance(content_value, str):
+                try:
+                    content_value = json.loads(content_value)
+                except (json.JSONDecodeError, TypeError):
+                    # Keep as string if not valid JSON
+                    pass
+
+            return UiPathConversationMessageEvent(
+                message_id=result_message_id or str(uuid4()),
+                tool_call=UiPathConversationToolCallEvent(
+                    tool_call_id=message.tool_call_id,
+                    start=UiPathConversationToolCallStartEvent(
+                        tool_name=message.name,
+                        arguments=None,
+                        timestamp=timestamp,
+                    ),
+                    end=UiPathConversationToolCallEndEvent(
+                        timestamp=timestamp,
+                        output=UiPathInlineValue(inline=content_value),
+                    ),
+                ),
+            )
+
+        # --- Fallback for other BaseMessage types ---
+        text_content = self._extract_text(message.content)
+        return UiPathConversationMessageEvent(
+            message_id=message.id,
+            start=UiPathConversationMessageStartEvent(
+                role="assistant", timestamp=timestamp
+            ),
+            content_part=UiPathConversationContentPartEvent(
+                content_part_id=f"cp-{message.id}",
+                chunk=UiPathConversationContentPartChunkEvent(data=text_content),
+            ),
+            end=UiPathConversationMessageEndEvent(),
+        )
+
+
+__all__ = ["UiPathChatMessagesMapper"]

uipath_langchain/chat/models.py

@@ -1,15 +1,19 @@
 import json
 import logging
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, AsyncIterator, Iterator, Literal, Union
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
 from langchain_core.language_models import LanguageModelInput
-from langchain_core.messages import AIMessage, BaseMessage
+from langchain_core.language_models.chat_models import (
+    agenerate_from_stream,
+    generate_from_stream,
+)
+from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
 from langchain_core.messages.ai import UsageMetadata
-from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.runnables import Runnable
 from langchain_openai.chat_models import AzureChatOpenAI
 from pydantic import BaseModel
@@ -25,37 +29,118 @@ class UiPathAzureChatOpenAI(UiPathRequestMixin, AzureChatOpenAI):
 
     def _generate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         if "tools" in kwargs and not kwargs["tools"]:
             del kwargs["tools"]
+
+        if self.streaming:
+            stream_iter = self._stream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return generate_from_stream(stream_iter)
+
         payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = self._call(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
     async def _agenerate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         if "tools" in kwargs and not kwargs["tools"]:
             del kwargs["tools"]
+
+        if self.streaming:
+            stream_iter = self._astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return await agenerate_from_stream(stream_iter)
+
         payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = await self._acall(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
+    def _stream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        if "tools" in kwargs and not kwargs["tools"]:
+            del kwargs["tools"]
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        default_chunk_class = AIMessageChunk
+
+        for chunk in self._stream_request(self.url, payload, self.auth_headers):
+            if self.logger:
+                self.logger.debug(f"[Stream] Got chunk from _stream_request: {chunk}")
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                if self.logger:
+                    self.logger.debug("[Stream] Skipping None generation_chunk")
+                continue
+
+            if self.logger:
+                self.logger.debug(
+                    f"[Stream] Yielding generation_chunk: {generation_chunk}"
+                )
+
+            if run_manager:
+                run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
+    async def _astream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        if "tools" in kwargs and not kwargs["tools"]:
+            del kwargs["tools"]
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        default_chunk_class = AIMessageChunk
+
+        async for chunk in self._astream_request(self.url, payload, self.auth_headers):
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                continue
+
+            if run_manager:
+                await run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
     def with_structured_output(
         self,
-        schema: Optional[Any] = None,
+        schema: Any = None,
         *,
         method: Literal["function_calling", "json_mode", "json_schema"] = "json_schema",
         include_raw: bool = False,
-        strict: Optional[bool] = None,
+        strict: bool | None = None,
         **kwargs: Any,
     ) -> Runnable[LanguageModelInput, Any]:
         """Model wrapper that returns outputs formatted to match the given schema."""
@@ -86,8 +171,8 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
 
     def _create_chat_result(
         self,
-        response: Union[Dict[str, Any], BaseModel],
-        generation_info: Optional[Dict[Any, Any]] = None,
+        response: Union[dict[str, Any], BaseModel],
+        generation_info: dict[Any, Any] | None = None,
     ) -> ChatResult:
         if not isinstance(response, dict):
             response = response.model_dump()
@@ -128,9 +213,9 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
         self,
         input_: LanguageModelInput,
         *,
-        stop: Optional[List[str]] = None,
+        stop: list[str] | None = None,
         **kwargs: Any,
-    ) -> Dict[Any, Any]:
+    ) -> dict[Any, Any]:
         payload = super()._get_request_payload(input_, stop=stop, **kwargs)
         # hacks to make the request work with uipath normalized
         for message in payload["messages"]:
@@ -149,11 +234,40 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
                 }
         return payload
 
+    def _normalize_tool_choice(self, kwargs: dict[str, Any]) -> None:
+        """Normalize tool_choice for UiPath Gateway compatibility.
+
+        Converts LangChain tool_choice formats to UiPath Gateway format:
+        - String "required" -> {"type": "required"}
+        - String "auto" -> {"type": "auto"}
+        - Dict with function -> {"type": "tool", "name": "function_name"}
+        """
+        if "tool_choice" in kwargs:
+            tool_choice = kwargs["tool_choice"]
+
+            if isinstance(tool_choice, str):
+                if tool_choice in ("required", "auto", "none"):
+                    logger.debug(
+                        f"Converting tool_choice from '{tool_choice}' to {{'type': '{tool_choice}'}}"
+                    )
+                    kwargs["tool_choice"] = {"type": tool_choice}
+            elif (
+                isinstance(tool_choice, dict) and tool_choice.get("type") == "function"
+            ):
+                function_name = tool_choice["function"]["name"]
+                logger.debug(
+                    f"Converting tool_choice from function '{function_name}' to tool format"
+                )
+                kwargs["tool_choice"] = {
+                    "type": "tool",
+                    "name": function_name,
+                }
+
     def _generate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         """Override the _generate method to implement the chat model logic.
@@ -173,21 +287,23 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
         """
         if kwargs.get("tools"):
             kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
-        if "tool_choice" in kwargs and kwargs["tool_choice"]["type"] == "function":
-            kwargs["tool_choice"] = {
-                "type": "tool",
-                "name": kwargs["tool_choice"]["function"]["name"],
-            }
-        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+        self._normalize_tool_choice(kwargs)
+
+        if self.streaming:
+            stream_iter = self._stream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return generate_from_stream(stream_iter)
 
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = self._call(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
     async def _agenerate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         """Override the _generate method to implement the chat model logic.
@@ -207,25 +323,122 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
         """
         if kwargs.get("tools"):
             kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
-        if "tool_choice" in kwargs and kwargs["tool_choice"]["type"] == "function":
-            kwargs["tool_choice"] = {
-                "type": "tool",
-                "name": kwargs["tool_choice"]["function"]["name"],
-            }
-        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+        self._normalize_tool_choice(kwargs)
+
+        if self.streaming:
+            stream_iter = self._astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return await agenerate_from_stream(stream_iter)
 
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = await self._acall(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
+    def _stream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        """Stream the LLM on a given prompt.
+
+        Args:
+            messages: the prompt composed of a list of messages.
+            stop: a list of strings on which the model should stop generating.
+            run_manager: A run manager with callbacks for the LLM.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            An iterator of ChatGenerationChunk objects.
+        """
+        if kwargs.get("tools"):
+            kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
+        self._normalize_tool_choice(kwargs)
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        default_chunk_class = AIMessageChunk
+
+        for chunk in self._stream_request(self.url, payload, self.auth_headers):
+            if self.logger:
+                self.logger.debug(f"[Stream] Got chunk from _stream_request: {chunk}")
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                if self.logger:
+                    self.logger.debug("[Stream] Skipping None generation_chunk")
+                continue
+
+            if self.logger:
+                self.logger.debug(
+                    f"[Stream] Yielding generation_chunk: {generation_chunk}"
+                )
+
+            if run_manager:
+                run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
+    async def _astream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        """Async stream the LLM on a given prompt.
+
+        Args:
+            messages: the prompt composed of a list of messages.
+            stop: a list of strings on which the model should stop generating.
+            run_manager: A run manager with callbacks for the LLM.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            An async iterator of ChatGenerationChunk objects.
+        """
+        if kwargs.get("tools"):
+            kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
+        self._normalize_tool_choice(kwargs)
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        # Update headers to enable streaming
+        headers = {**self.auth_headers}
+        headers["X-UiPath-Streaming-Enabled"] = "true"
+
+        default_chunk_class = AIMessageChunk
+
+        async for chunk in self._astream_request(self.url, payload, headers):
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                continue
+
+            if run_manager:
+                await run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
     def with_structured_output(
         self,
-        schema: Optional[Any] = None,
+        schema: Any = None,
         *,
         method: Literal[
             "function_calling", "json_mode", "json_schema"
         ] = "function_calling",
         include_raw: bool = False,
-        strict: Optional[bool] = None,
+        strict: bool | None = None,
         **kwargs: Any,
     ) -> Runnable[LanguageModelInput, Any]:
         """Model wrapper that returns outputs formatted to match the given schema."""