feat: add citation and reasoning mappings

Author: GabrielVasilescu04

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.18"
+version = "0.1.19"
 description = "UiPath Langchain"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
@@ -113,4 +113,3 @@ name = "testpypi"
 url = "https://test.pypi.org/simple/"
 publish-url = "https://test.pypi.org/legacy/"
 explicit = true
-

src/uipath_langchain/chat/mapper.py

@@ -7,12 +7,18 @@
 from langchain_core.messages import (
     AIMessageChunk,
     BaseMessage,
+    Citation,
     HumanMessage,
+    ReasoningContentBlock,
     TextContentBlock,
     ToolCallChunk,
     ToolMessage,
 )
 from uipath.core.chat import (
+    UiPathConversationCitationEndEvent,
+    UiPathConversationCitationEvent,
+    UiPathConversationCitationSource,
+    UiPathConversationCitationStartEvent,
     UiPathConversationContentPartChunkEvent,
     UiPathConversationContentPartEndEvent,
     UiPathConversationContentPartEvent,
@@ -184,11 +190,60 @@ def map_event(
                         text_block = cast(TextContentBlock, block)
                         text = text_block["text"]
 
+                        # Map citations if present
+                        annotations = text_block.get("annotations", [])
+                        citation_annotations = [
+                            cast(Citation, annotation)
+                            for annotation in annotations
+                            if annotation.get("type") == "citation"
+                        ]
+
+                        if citation_annotations:
+                            for citation_annotation in citation_annotations:
+                                # Build citation source, only include url if present
+                                block_index = text_block.get("index", 0)
+                                block_index = block_index // 2 + 1
+                                source_args = {
+                                    "title": citation_annotation.get("title"),
+                                    "number": block_index,
+                                }
+                                if citation_annotation.get("url") is not None:
+                                    source_args["url"] = citation_annotation.get("url")
+
+                                citation_source = UiPathConversationCitationSource(
+                                    **source_args
+                                )
+
+                                msg_event.content_part = UiPathConversationContentPartEvent(
+                                    content_part_id=f"chunk-{message.id}-0",
+                                    chunk=UiPathConversationContentPartChunkEvent(
+                                        data=citation_annotation["cited_text"],
+                                        citation=UiPathConversationCitationEvent(
+                                            citation_id=str(uuid4()),
+                                            start=UiPathConversationCitationStartEvent(),
+                                            end=UiPathConversationCitationEndEvent(
+                                                sources=[citation_source]
+                                            ),
+                                        ),
+                                    ),
+                                )
+                        else:
+                            msg_event.content_part = UiPathConversationContentPartEvent(
+                                content_part_id=f"chunk-{message.id}-0",
+                                chunk=UiPathConversationContentPartChunkEvent(
+                                    data=text
+                                ),
+                            )
+
+                    # Mapping to data temporarily until proper support is added
+                    elif block_type == "reasoning":
+                        reasoning_block = cast(ReasoningContentBlock, block)
+                        reasoning = reasoning_block["reasoning"]
+
                         msg_event.content_part = UiPathConversationContentPartEvent(
                             content_part_id=f"chunk-{message.id}-0",
                             chunk=UiPathConversationContentPartChunkEvent(
-                                data=text,
-                                content_part_sequence=0,
+                                data=reasoning
                             ),
                         )
 
@@ -204,10 +259,7 @@ def map_event(
 
                         msg_event.content_part = UiPathConversationContentPartEvent(
                             content_part_id=f"chunk-{message.id}-0",
-                            chunk=UiPathConversationContentPartChunkEvent(
-                                data=args,
-                                content_part_sequence=0,
-                            ),
+                            chunk=UiPathConversationContentPartChunkEvent(data=args),
                         )
                         # Continue so that multiple tool_call_chunks in the same block list
                         # are handled correctly
@@ -217,10 +269,7 @@ def map_event(
             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,
-                    ),
+                    chunk=UiPathConversationContentPartChunkEvent(data=message.content),
                 )
 
             if (
@@ -269,12 +318,12 @@ def map_event(
                     tool_call_id=message.tool_call_id,
                     start=UiPathConversationToolCallStartEvent(
                         tool_name=message.name,
-                        arguments=None,
+                        input=None,
                         timestamp=timestamp,
                     ),
                     end=UiPathConversationToolCallEndEvent(
                         timestamp=timestamp,
-                        result=UiPathInlineValue(inline=content_value),
+                        output=UiPathInlineValue(inline=content_value),
                     ),
                 ),
             )

src/uipath_langchain/runtime/runtime.py

@@ -7,6 +7,7 @@
 from langgraph.errors import EmptyInputError, GraphRecursionError, InvalidUpdateError
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.types import Command, Interrupt, StateSnapshot
+from pydantic import ValidationError
 from uipath.runtime import (
     UiPathBreakpointResult,
     UiPathExecuteOptions,
@@ -135,10 +136,17 @@ async def stream(
                 if chunk_type == "messages":
                     if isinstance(data, tuple):
                         message, _ = data
-                        event = UiPathRuntimeMessageEvent(
-                            payload=self.chat.map_event(message),
-                        )
-                        yield event
+
+                        try:
+                            mapped_event = self.chat.map_event(message)
+                            event = UiPathRuntimeMessageEvent(
+                                payload=mapped_event,
+                            )
+                            yield event
+                        except ValidationError as e:
+                            logger.warning(
+                                f"Failed to map event due to validation error, skipping: {e}"
+                            )
 
                 # Emit UiPathRuntimeStateEvent for state updates
                 elif chunk_type == "updates":