|
| 1 | +"""Mock runtime that replays events from a JSON file.""" |
| 2 | + |
| 3 | +import asyncio |
| 4 | +import json |
| 5 | +import logging |
| 6 | +from pathlib import Path |
| 7 | +from typing import Any, AsyncGenerator, cast |
| 8 | + |
| 9 | +from opentelemetry import trace |
| 10 | +from uipath.core.chat import UiPathConversationMessageEvent |
| 11 | +from uipath.runtime import ( |
| 12 | + UiPathExecuteOptions, |
| 13 | + UiPathRuntimeEvent, |
| 14 | + UiPathRuntimeResult, |
| 15 | + UiPathRuntimeStatus, |
| 16 | + UiPathStreamOptions, |
| 17 | +) |
| 18 | +from uipath.runtime.events import ( |
| 19 | + UiPathRuntimeMessageEvent, |
| 20 | + UiPathRuntimeStateEvent, |
| 21 | +) |
| 22 | +from uipath.runtime.schema import ( |
| 23 | + UiPathRuntimeEdge, |
| 24 | + UiPathRuntimeGraph, |
| 25 | + UiPathRuntimeNode, |
| 26 | + UiPathRuntimeSchema, |
| 27 | +) |
| 28 | + |
| 29 | +logger = logging.getLogger(__name__) |
| 30 | + |
| 31 | + |
| 32 | +class MockTemplateRuntime: |
| 33 | + """Mock runtime that replays events from a JSON file.""" |
| 34 | + |
| 35 | + def __init__( |
| 36 | + self, |
| 37 | + events_file: str | Path, |
| 38 | + schema_file: str | Path, |
| 39 | + ) -> None: |
| 40 | + self.events_file = Path(events_file) |
| 41 | + self.schema_file = Path(schema_file) |
| 42 | + self.tracer = trace.get_tracer("uipath.dev.mock.template") |
| 43 | + self._events: list[dict[str, Any]] = [] |
| 44 | + self._schema_data: dict[str, Any] | None = None |
| 45 | + self._load_events() |
| 46 | + self._load_schema() |
| 47 | + |
| 48 | + def _load_events(self) -> None: |
| 49 | + """Load events from JSON file.""" |
| 50 | + with open(self.events_file, "r", encoding="utf-8") as f: |
| 51 | + data = json.load(f) |
| 52 | + |
| 53 | + if isinstance(data, list): |
| 54 | + self._events = data |
| 55 | + else: |
| 56 | + raise ValueError("Expected JSON array of events") |
| 57 | + |
| 58 | + def _load_schema(self) -> None: |
| 59 | + """Load schema from separate JSON file if provided.""" |
| 60 | + with open(self.schema_file, "r", encoding="utf-8") as f: |
| 61 | + self._schema_data = json.load(f) |
| 62 | + |
| 63 | + async def get_schema(self) -> UiPathRuntimeSchema: |
| 64 | + """Get runtime schema from the loaded data or build default.""" |
| 65 | + assert self._schema_data is not None |
| 66 | + entry_points = self._schema_data.get("entryPoints", []) |
| 67 | + assert entry_points is not None and len(entry_points) > 0 |
| 68 | + ep = entry_points[0] |
| 69 | + |
| 70 | + # Build graph if present |
| 71 | + graph = None |
| 72 | + if "graph" in ep and ep["graph"]: |
| 73 | + graph_data = ep["graph"] |
| 74 | + nodes = [UiPathRuntimeNode(**node) for node in graph_data.get("nodes", [])] |
| 75 | + edges = [UiPathRuntimeEdge(**edge) for edge in graph_data.get("edges", [])] |
| 76 | + graph = UiPathRuntimeGraph(nodes=nodes, edges=edges) |
| 77 | + |
| 78 | + return UiPathRuntimeSchema( |
| 79 | + filePath=ep["filePath"], |
| 80 | + uniqueId=ep["uniqueId"], |
| 81 | + type=ep["type"], |
| 82 | + input=ep["input"], |
| 83 | + output=ep["output"], |
| 84 | + graph=graph, |
| 85 | + ) |
| 86 | + |
| 87 | + async def execute( |
| 88 | + self, |
| 89 | + input: dict[str, Any] | None = None, |
| 90 | + options: UiPathExecuteOptions | None = None, |
| 91 | + ) -> UiPathRuntimeResult: |
| 92 | + """Execute by replaying all events and returning final result.""" |
| 93 | + logger.info("MockTemplateRuntime: starting execution") |
| 94 | + |
| 95 | + # Stream all events and capture the final result |
| 96 | + final_result = None |
| 97 | + async for event in self.stream( |
| 98 | + input=input, options=cast(UiPathStreamOptions, options) |
| 99 | + ): |
| 100 | + if isinstance(event, UiPathRuntimeResult): |
| 101 | + final_result = event |
| 102 | + |
| 103 | + return final_result or UiPathRuntimeResult( |
| 104 | + output={}, |
| 105 | + status=UiPathRuntimeStatus.SUCCESSFUL, |
| 106 | + ) |
| 107 | + |
| 108 | + async def stream( |
| 109 | + self, |
| 110 | + input: dict[str, Any] | None = None, |
| 111 | + options: UiPathStreamOptions | None = None, |
| 112 | + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: |
| 113 | + """Stream events from the JSON file.""" |
| 114 | + logger.info(f"MockTemplateRuntime: streaming {len(self._events)} events") |
| 115 | + |
| 116 | + with self.tracer.start_as_current_span( |
| 117 | + "template.stream", |
| 118 | + attributes={ |
| 119 | + "uipath.runtime.name": "MockTemplateRuntime", |
| 120 | + "uipath.event.count": len(self._events), |
| 121 | + }, |
| 122 | + ): |
| 123 | + for i, event_data in enumerate(self._events): |
| 124 | + event_type = event_data.get("event_type") |
| 125 | + |
| 126 | + # Add small delay between events for realistic streaming |
| 127 | + if i > 0: |
| 128 | + await asyncio.sleep(0.01) |
| 129 | + |
| 130 | + try: |
| 131 | + if event_type == "runtime_message": |
| 132 | + payload_data = event_data.get("payload", {}) |
| 133 | + |
| 134 | + conversation_event = ( |
| 135 | + UiPathConversationMessageEvent.model_validate(payload_data) |
| 136 | + ) |
| 137 | + |
| 138 | + yield UiPathRuntimeMessageEvent( |
| 139 | + payload=conversation_event, |
| 140 | + execution_id=event_data.get("execution_id"), |
| 141 | + metadata=event_data.get("metadata"), |
| 142 | + ) |
| 143 | + |
| 144 | + elif event_type == "runtime_state": |
| 145 | + yield UiPathRuntimeStateEvent( |
| 146 | + payload=event_data.get("payload", {}), |
| 147 | + node_name=event_data.get("node_name"), |
| 148 | + execution_id=event_data.get("execution_id"), |
| 149 | + metadata=event_data.get("metadata"), |
| 150 | + ) |
| 151 | + |
| 152 | + elif event_type == "runtime_result": |
| 153 | + yield UiPathRuntimeResult( |
| 154 | + output=event_data.get("output"), |
| 155 | + status=UiPathRuntimeStatus( |
| 156 | + event_data.get("status", "successful") |
| 157 | + ), |
| 158 | + execution_id=event_data.get("execution_id"), |
| 159 | + metadata=event_data.get("metadata"), |
| 160 | + ) |
| 161 | + |
| 162 | + else: |
| 163 | + logger.warning(f"Unknown event type: {event_type}") |
| 164 | + |
| 165 | + except Exception as e: |
| 166 | + logger.error(f"Error processing event {i}: {e}", exc_info=True) |
| 167 | + continue |
| 168 | + |
| 169 | + logger.info("MockReplayRuntime: streaming completed") |
| 170 | + |
| 171 | + async def dispose(self) -> None: |
| 172 | + """Cleanup resources.""" |
| 173 | + logger.info("MockReplayRuntime: dispose() invoked") |
| 174 | + |
| 175 | + |
| 176 | +def create_template_runtime( |
| 177 | + events_json: str | Path, |
| 178 | + schema_json: str | Path, |
| 179 | +) -> MockTemplateRuntime: |
| 180 | + """Create a template runtime from JSON files. |
| 181 | +
|
| 182 | + Args: |
| 183 | + events_json: Path to JSON file containing events array |
| 184 | + schema_json: Optional path to entry-points.json schema file |
| 185 | +
|
| 186 | + Returns: |
| 187 | + MockTemplateRuntime instance |
| 188 | + """ |
| 189 | + return MockTemplateRuntime( |
| 190 | + events_file=events_json, |
| 191 | + schema_file=schema_json, |
| 192 | + ) |
0 commit comments