Skip to content

feat: hydrate managed workspaces in cli #1845

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
radu-mocanu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: hydrate managed workspaces in cli #1845

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/uipath/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "uipath"
version = "2.14.1"
version = "2.15.0"
description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
readme = { file = "README.md", content-type = "text/markdown" }
requires-python = ">=3.11"
Expand Down
66 changes: 66 additions & 0 deletions packages/uipath/src/uipath/_cli/_managed_workspace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
from contextlib import AsyncExitStack

from uipath.platform import UiPath
from uipath.runtime import (
ConversationalWorkspaceRuntime,
HydrationRuntime,
UiPathRuntimeContext,
UiPathRuntimeFactoryProtocol,
UiPathRuntimeProtocol,
Workspace,
WorkspaceHydrator,
WorkspaceRegistryStore,
)


async def wrap_with_managed_workspace(
delegate: UiPathRuntimeProtocol,
*,
context: UiPathRuntimeContext,
factory: UiPathRuntimeFactoryProtocol,
enabled: bool,
cleanup: AsyncExitStack,
) -> UiPathRuntimeProtocol:
if context.job_id is None or not enabled:
return delegate

storage = await factory.get_storage()
if storage is None:
raise RuntimeError(
"Runtime factory advertises managed workspace support but provides no storage"
)

client = UiPath()
workspace = Workspace.create()
try:
workspace.path = workspace.path.resolve()
hydrator = WorkspaceHydrator(
workspace_path=workspace.path,
attachments=client.attachments,
jobs=client.jobs,
current_job_key=context.job_id,
folder_key=context.folder_key,
)
registry_store = WorkspaceRegistryStore(storage, context.job_id)
hydration_runtime = HydrationRuntime(
delegate,
workspace=workspace,
hydrator=hydrator,
registry_store=registry_store,
)

if context.conversation_id is None or context.exchange_id is None:
cleanup.push_async_callback(hydration_runtime.dispose)
return hydration_runtime

conversational_runtime = ConversationalWorkspaceRuntime(
hydration_runtime,
hydrator=hydrator,
registry_store=registry_store,
)
cleanup.push_async_callback(hydration_runtime.dispose)
cleanup.push_async_callback(conversational_runtime.dispose)
return conversational_runtime
except BaseException:
await workspace.dispose()
raise
146 changes: 87 additions & 59 deletions packages/uipath/src/uipath/_cli/cli_debug.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import asyncio
import logging
from contextlib import AsyncExitStack
from typing import Any, cast, get_args

import click
Expand Down Expand Up @@ -29,6 +30,7 @@
from uipath.tracing import LiveTrackingSpanProcessor, LlmOpsHttpExporter

from ._governance_bootstrap import GovernanceBootstrap, resolve_governance
from ._managed_workspace import wrap_with_managed_workspace
from ._telemetry import track_command
from ._utils._console import ConsoleLogger
from .middlewares import Middlewares
Expand Down Expand Up @@ -196,82 +198,108 @@ async def execute_debug_runtime():

async def execute_debug_runtime():
chat_runtime: UiPathRuntimeProtocol | None = None
debug_bridge: UiPathDebugProtocol = get_debug_bridge(
ctx, attach=attach_mode
)
new_runtime_kwargs: dict[str, Any] = {}
if governance_bootstrap is not None:
new_runtime_kwargs["evaluator"] = (
governance_bootstrap.evaluator
managed_workspace_created = False
managed_workspace_cleanup = AsyncExitStack()
debug_runtime: UiPathRuntimeProtocol | None = None
execution_runtime: UiPathRuntimeProtocol | None = None
runtime: UiPathRuntimeProtocol | None = None
try:
debug_bridge: UiPathDebugProtocol = get_debug_bridge(
ctx, attach=attach_mode
)
new_runtime_kwargs: dict[str, Any] = {}
if governance_bootstrap is not None:
new_runtime_kwargs["evaluator"] = (
governance_bootstrap.evaluator
)
runtime = await factory.new_runtime(
entrypoint,
governance_runtime_id,
**new_runtime_kwargs,
)
runtime = await factory.new_runtime(
entrypoint,
governance_runtime_id,
**new_runtime_kwargs,
)

if governance_bootstrap is not None:
runtime = governance_bootstrap.wrap_runtime(
if governance_bootstrap is not None:
runtime = governance_bootstrap.wrap_runtime(
runtime,
agent_name=entrypoint,
runtime_id=governance_runtime_id,
)

delegate = runtime
delegate = await wrap_with_managed_workspace(
runtime,
agent_name=entrypoint,
runtime_id=governance_runtime_id,
context=ctx,
factory=factory,
enabled=bool(
factory_settings
and factory_settings.managed_workspace
),
cleanup=managed_workspace_cleanup,
)
managed_workspace_created = delegate is not runtime

delegate = runtime
if ctx.conversation_id and ctx.exchange_id:
chat_bridge: UiPathChatProtocol = get_chat_bridge(
context=ctx
)
chat_runtime = UiPathChatRuntime(
delegate=delegate, chat_bridge=chat_bridge
)
delegate = chat_runtime
if ctx.conversation_id and ctx.exchange_id:
chat_bridge: UiPathChatProtocol = get_chat_bridge(
context=ctx
)
chat_runtime = UiPathChatRuntime(
delegate=delegate, chat_bridge=chat_bridge
)
delegate = chat_runtime

debug_runtime = UiPathDebugRuntime(
delegate=delegate,
debug_bridge=debug_bridge,
trigger_poll_interval=trigger_poll_interval,
)
debug_runtime = UiPathDebugRuntime(
delegate=delegate,
debug_bridge=debug_bridge,
trigger_poll_interval=trigger_poll_interval,
)

# Build mocking context with agent model for simulations
schema = await runtime.get_schema()
agent_model = None
if schema.metadata and "settings" in schema.metadata:
agent_model = schema.metadata["settings"].get("model")
schema = await runtime.get_schema()
agent_model = None
if schema.metadata and "settings" in schema.metadata:
agent_model = schema.metadata["settings"].get(
"model"
)

delegate_runtime: UiPathDebugRuntime | UiPathMockRuntime = (
debug_runtime
)
if simulation_config:
mocking_context = build_mocking_context(
simulation_config, agent_model
)
if mocking_context:
delegate_runtime = UiPathMockRuntime(
if simulation_config:
mocking_context = build_mocking_context(
simulation_config, agent_model
)
if mocking_context:
execution_runtime = UiPathMockRuntime(
delegate=debug_runtime,
mocking_context=mocking_context,
)
else:
mocking_context = load_simulation_config(
agent_model=agent_model
)
execution_runtime = UiPathMockRuntime(
delegate=debug_runtime,
mocking_context=mocking_context,
)
else:
mocking_context = load_simulation_config(
agent_model=agent_model
)
delegate_runtime = UiPathMockRuntime(
delegate=debug_runtime,
mocking_context=mocking_context,
)

try:
ctx.result = await delegate_runtime.execute(
awaitable_runtime = execution_runtime or debug_runtime
ctx.result = await awaitable_runtime.execute(
ctx.get_input(),
options=UiPathExecuteOptions(resume=resume),
)
finally:
if delegate_runtime is not debug_runtime:
await delegate_runtime.dispose()
await debug_runtime.dispose()
cleanup = AsyncExitStack()
if not managed_workspace_created:
if runtime is not None:
cleanup.push_async_callback(runtime.dispose)
cleanup.push_async_callback(
managed_workspace_cleanup.aclose
)
if chat_runtime:
await chat_runtime.dispose()
await runtime.dispose()
cleanup.push_async_callback(chat_runtime.dispose)
if debug_runtime is not None:
cleanup.push_async_callback(debug_runtime.dispose)
if execution_runtime is not None:
cleanup.push_async_callback(
execution_runtime.dispose
)
await cleanup.aclose()

if project_id := UiPathConfig.project_id:
studio_client = StudioClient(project_id)
Expand Down
51 changes: 39 additions & 12 deletions packages/uipath/src/uipath/_cli/cli_run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import asyncio
from contextlib import AsyncExitStack
from typing import Any

import click
Expand Down Expand Up @@ -36,6 +37,7 @@

from ._errors import EntrypointDiscoveryException
from ._governance_bootstrap import GovernanceBootstrap, resolve_governance
from ._managed_workspace import wrap_with_managed_workspace
from ._telemetry import track_command
from ._utils._console import ConsoleLogger
from .middlewares import Middlewares
Expand Down Expand Up @@ -218,7 +220,10 @@ async def execute() -> None:
with ExecutionSourceContext(ctx.execution_source), ctx:
base_runtime: UiPathRuntimeProtocol | None = None
runtime: UiPathRuntimeProtocol | None = None
workspace_delegate: UiPathRuntimeProtocol | None = None
chat_runtime: UiPathRuntimeProtocol | None = None
managed_workspace_created = False
managed_workspace_cleanup = AsyncExitStack()
factory: UiPathRuntimeFactoryProtocol | None = None
governance_bootstrap: GovernanceBootstrap | None = None
try:
Expand Down Expand Up @@ -293,6 +298,21 @@ async def execute() -> None:
mocking_context=mocking_context,
)

workspace_delegate = runtime
runtime = await wrap_with_managed_workspace(
workspace_delegate,
context=ctx,
factory=factory,
enabled=bool(
factory_settings
and factory_settings.managed_workspace
),
cleanup=managed_workspace_cleanup,
)
managed_workspace_created = (
runtime is not workspace_delegate
)

if ctx.job_id:
if UiPathConfig.is_tracing_enabled:
trace_manager.add_span_processor(
Expand All @@ -316,19 +336,26 @@ async def execute() -> None:
else:
ctx.result = await debug_runtime(ctx, runtime)
finally:
try:
if chat_runtime:
await chat_runtime.dispose()
cleanup = AsyncExitStack()
cleanup.callback(trace_manager.shutdown)
if factory:
cleanup.push_async_callback(factory.dispose)
if governance_bootstrap is not None:
cleanup.callback(governance_bootstrap.dispose)
if base_runtime is not None and (
not managed_workspace_created
or workspace_delegate is not base_runtime
):
cleanup.push_async_callback(base_runtime.dispose)
if not managed_workspace_created:
if runtime is not None and runtime is not base_runtime:
await runtime.dispose()
if base_runtime is not None:
await base_runtime.dispose()
if governance_bootstrap is not None:
governance_bootstrap.dispose()
if factory:
await factory.dispose()
finally:
trace_manager.shutdown()
cleanup.push_async_callback(runtime.dispose)
cleanup.push_async_callback(
managed_workspace_cleanup.aclose
)
if chat_runtime:
cleanup.push_async_callback(chat_runtime.dispose)
await cleanup.aclose()

asyncio.run(execute())

Expand Down
Loading
Loading