feat: ad conversational agent contracts

Author: cristipufu

src/uipath/_cli/_runtime/_contracts.py

@@ -22,6 +22,7 @@
 from opentelemetry.trace import Tracer
 from pydantic import BaseModel, Field
 
+from uipath.agent.conversation import UiPathConversationEvent
 from uipath.tracing import TracingManager
 
 from ._logging import LogsInterceptor
@@ -134,6 +135,17 @@ def to_dict(self) -> Dict[str, Any]:
         return result
 
 
+class UiPathConversationHandler(ABC):
+    """Base delegate for handling UiPath conversation events."""
+
+    use_streaming: bool = True
+
+    @abstractmethod
+    def on_event(self, event: UiPathConversationEvent) -> None:
+        """Handle a conversation event for a given execution run."""
+        pass
+
+
 class UiPathTraceContext(BaseModel):
     """Trace context information for tracing and debugging."""
 
@@ -173,6 +185,8 @@ class UiPathRuntimeContext(BaseModel):
     input_file: Optional[str] = None
     is_eval_run: bool = False
     log_handler: Optional[logging.Handler] = None
+    chat_handler: Optional[UiPathConversationHandler] = None
+
     model_config = {"arbitrary_types_allowed": True}
 
     @classmethod

src/uipath/agent/conversation/__init__.py

@@ -0,0 +1,131 @@
+"""UiPath Conversation Models.
+
+This module provides Pydantic models that represent the JSON event schema for conversations between a client (UI) and an LLM/agent.
+
+The event objects define a hierarchal conversation structure:
+
+* Conversation
+    * Exchange
+        * Message
+            * Content Parts
+                * Citations
+            * Tool Calls
+                * Tool Results
+
+ A conversation may contain multiple exchanges, and an exchange may contain multiple messages. A message may contain
+ multiple content parts, each of which can be text or binary, including media input and output streams; and each
+ content part can include multiple citations. A message may also contain multiple tool calls, which may contain a tool
+ result.
+
+ The protocol also supports a top level, "async", input media streams (audio and video), which can span multiple
+ exchanges. These are used for Gemini's automatic turn detection mode, where the LLM determines when the user has
+ stopped talking and starts producing output. The output forms one or more messages in an exchange with no explicit
+ input message. However, the LLM may produce an input transcript which can be used to construct the implicit input
+ message that started the exchange.
+
+ In addition, the protocol also supports "async" tool calls that span multiple exchanges. This can be used with
+ Gemini's asynchronous function calling protocol, which allows function calls to produce results that interrupt the
+ conversation when ready, even after multiple exchanges. They also support generating multiple results from a single
+ tool call. By contrast most tool calls are scoped to a single message, which contains both the call and the single
+ result produced by that call.
+
+ Not all features supported by the protocol will be supported by all clients and LLMs. The optional top level
+ `capabilities` property can be used to communicate information about supported features. This property should be set
+ on the first event written to a new websocket connection. This initial event may or may not contain additional
+ sub-events.
+"""
+
+from .async_stream import (
+    UiPathConversationAsyncInputStreamEndEvent,
+    UiPathConversationAsyncInputStreamEvent,
+    UiPathConversationAsyncInputStreamStartEvent,
+    UiPathConversationInputStreamChunkEvent,
+)
+from .citation import (
+    UiPathConversationCitationEndEvent,
+    UiPathConversationCitationEvent,
+    UiPathConversationCitationSource,
+    UiPathConversationCitationSourceMedia,
+    UiPathConversationCitationSourceUrl,
+    UiPathConversationCitationStartEvent,
+)
+from .content import (
+    UiPathConversationContentPart,
+    UiPathConversationContentPartChunkEvent,
+    UiPathConversationContentPartEndEvent,
+    UiPathConversationContentPartEvent,
+    UiPathConversationContentPartStartEvent,
+)
+from .conversation import (
+    UiPathConversationCapabilities,
+    UiPathConversationEndEvent,
+    UiPathConversationStartedEvent,
+    UiPathConversationStartEvent,
+)
+from .event import UiPathConversationEvent
+from .exchange import (
+    UiPathConversationExchange,
+    UiPathConversationExchangeEndEvent,
+    UiPathConversationExchangeEvent,
+    UiPathConversationExchangeStartEvent,
+)
+from .message import (
+    UiPathConversationMessage,
+    UiPathConversationMessageEndEvent,
+    UiPathConversationMessageEvent,
+    UiPathConversationMessageStartEvent,
+)
+from .meta import UiPathConversationMetaEvent
+from .tool import (
+    UiPathConversationToolCall,
+    UiPathConversationToolCallEndEvent,
+    UiPathConversationToolCallEvent,
+    UiPathConversationToolCallResult,
+    UiPathConversationToolCallStartEvent,
+)
+
+__all__ = [
+    # Root
+    "UiPathConversationEvent",
+    # Conversation
+    "UiPathConversationCapabilities",
+    "UiPathConversationStartEvent",
+    "UiPathConversationStartedEvent",
+    "UiPathConversationEndEvent",
+    # Exchange
+    "UiPathConversationExchangeStartEvent",
+    "UiPathConversationExchangeEndEvent",
+    "UiPathConversationExchangeEvent",
+    "UiPathConversationExchange",
+    # Message
+    "UiPathConversationMessageStartEvent",
+    "UiPathConversationMessageEndEvent",
+    "UiPathConversationMessageEvent",
+    "UiPathConversationMessage",
+    # Content
+    "UiPathConversationContentPartChunkEvent",
+    "UiPathConversationContentPartStartEvent",
+    "UiPathConversationContentPartEndEvent",
+    "UiPathConversationContentPartEvent",
+    "UiPathConversationContentPart",
+    # Citation
+    "UiPathConversationCitationStartEvent",
+    "UiPathConversationCitationEndEvent",
+    "UiPathConversationCitationEvent",
+    "UiPathConversationCitationSource",
+    "UiPathConversationCitationSourceUrl",
+    "UiPathConversationCitationSourceMedia",
+    # Tool
+    "UiPathConversationToolCallStartEvent",
+    "UiPathConversationToolCallEndEvent",
+    "UiPathConversationToolCallEvent",
+    "UiPathConversationToolCallResult",
+    "UiPathConversationToolCall",
+    # Async Stream
+    "UiPathConversationInputStreamChunkEvent",
+    "UiPathConversationAsyncInputStreamStartEvent",
+    "UiPathConversationAsyncInputStreamEndEvent",
+    "UiPathConversationAsyncInputStreamEvent",
+    # Meta
+    "UiPathConversationMetaEvent",
+]

src/uipath/agent/conversation/async_stream.py

@@ -0,0 +1,54 @@
+"""Async input stream events."""
+
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class UiPathConversationInputStreamChunkEvent(BaseModel):
+    """Represents a single chunk of input stream data."""
+
+    input_stream_sequence: Optional[int] = Field(None, alias="inputStreamSequence")
+    data: str
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationAsyncInputStreamStartEvent(BaseModel):
+    """Signals the start of an asynchronous input stream."""
+
+    mime_type: str = Field(..., alias="mimeType")
+    start_of_speech_sensitivity: Optional[str] = Field(
+        None, alias="startOfSpeechSensitivity"
+    )
+    end_of_speech_sensitivity: Optional[str] = Field(
+        None, alias="endOfSpeechSensitivity"
+    )
+    prefix_padding_ms: Optional[int] = Field(None, alias="prefixPaddingMs")
+    silence_duration_ms: Optional[int] = Field(None, alias="silenceDurationMs")
+    meta_data: Optional[Dict[str, Any]] = Field(None, alias="metaData")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationAsyncInputStreamEndEvent(BaseModel):
+    """Signals the end of an asynchronous input stream."""
+
+    meta_data: Optional[Dict[str, Any]] = Field(None, alias="metaData")
+    last_chunk_content_part_sequence: Optional[int] = Field(
+        None, alias="lastChunkContentPartSequence"
+    )
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationAsyncInputStreamEvent(BaseModel):
+    """Encapsulates sub-events related to an asynchronous input stream."""
+
+    stream_id: str = Field(..., alias="streamId")
+    start: Optional[UiPathConversationAsyncInputStreamStartEvent] = None
+    end: Optional[UiPathConversationAsyncInputStreamEndEvent] = None
+    chunk: Optional[UiPathConversationInputStreamChunkEvent] = None
+    meta_event: Optional[Dict[str, Any]] = Field(None, alias="metaEvent")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)

src/uipath/agent/conversation/citation.py

@@ -0,0 +1,70 @@
+"""Citation events for message content."""
+
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class UiPathConversationCitationStartEvent(BaseModel):
+    """Indicates the start of a citation target in a content part."""
+
+    pass
+
+
+class UiPathConversationCitationEndEvent(BaseModel):
+    """Indicates the end of a citation target in a content part."""
+
+    sources: List[Dict[str, Any]]
+
+
+class UiPathConversationCitationEvent(BaseModel):
+    """Encapsulates sub-events related to citations."""
+
+    citation_id: str = Field(..., alias="citationId")
+    start: Optional[UiPathConversationCitationStartEvent] = None
+    end: Optional[UiPathConversationCitationEndEvent] = None
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationCitationSourceUrl(BaseModel):
+    """Represents a citation source that can be rendered as a link (URL)."""
+
+    url: str
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationCitationSourceMedia(BaseModel):
+    """Represents a citation source that references media, such as a PDF document."""
+
+    mime_type: str = Field(..., alias="mimeType")
+    download_url: Optional[str] = Field(None, alias="downloadUrl")
+    page_number: Optional[str] = Field(None, alias="pageNumber")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationCitationSource(BaseModel):
+    """Represents a citation source, either a URL or media reference."""
+
+    title: Optional[str] = None
+
+    # Union of Url or Media
+    url: Optional[str] = None
+    mime_type: Optional[str] = Field(None, alias="mimeType")
+    download_url: Optional[str] = Field(None, alias="downloadUrl")
+    page_number: Optional[str] = Field(None, alias="pageNumber")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationCitation(BaseModel):
+    """Represents a citation or reference inside a content part."""
+
+    citation_id: str = Field(..., alias="citationId")
+    offset: int
+    length: int
+    sources: List[UiPathConversationCitationSource]
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)

src/uipath/agent/conversation/content.py

@@ -0,0 +1,81 @@
+"""Message content part events."""
+
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .citation import UiPathConversationCitation, UiPathConversationCitationEvent
+
+
+class UiPathConversationContentPartChunkEvent(BaseModel):
+    """Contains a chunk of a message content part."""
+
+    content_part_sequence: Optional[int] = Field(None, alias="contentPartSequence")
+    data: Optional[str] = None
+    citation: Optional[UiPathConversationCitationEvent] = None
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationContentPartStartEvent(BaseModel):
+    """Signals the start of a message content part."""
+
+    mime_type: str = Field(..., alias="mimeType")
+    meta_data: Optional[Dict[str, Any]] = Field(None, alias="metaData")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationContentPartEndEvent(BaseModel):
+    """Signals the end of a message content part."""
+
+    last_chunk_content_part_sequence: Optional[int] = Field(
+        None, alias="lastChunkContentPartSequence"
+    )
+    interrupted: Optional[Dict[str, Any]] = None
+    meta_data: Optional[Dict[str, Any]] = Field(None, alias="metaData")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathConversationContentPartEvent(BaseModel):
+    """Encapsulates events related to message content parts."""
+
+    content_part_id: str = Field(..., alias="contentPartId")
+    start: Optional[UiPathConversationContentPartStartEvent] = None
+    end: Optional[UiPathConversationContentPartEndEvent] = None
+    chunk: Optional[UiPathConversationContentPartChunkEvent] = None
+    meta_event: Optional[Dict[str, Any]] = Field(None, alias="metaEvent")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+class UiPathInlineValue(BaseModel):
+    """Used when a value is small enough to be returned inline."""
+
+    inline: Any
+
+
+class UiPathExternalValue(BaseModel):
+    """Used when a value is too large to be returned inline."""
+
+    url: str
+    byte_count: Optional[int] = Field(None, alias="byteCount")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+
+InlineOrExternal = Union[UiPathInlineValue, UiPathExternalValue]
+
+
+class UiPathConversationContentPart(BaseModel):
+    """Represents a single part of message content."""
+
+    content_part_id: str = Field(..., alias="contentPartId")
+    mime_type: str = Field(..., alias="mimeType")
+    data: InlineOrExternal
+    citations: Optional[List[UiPathConversationCitation]] = None
+    is_transcript: Optional[bool] = Field(None, alias="isTranscript")
+    is_incomplete: Optional[bool] = Field(None, alias="isIncomplete")
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)