chore: refactor agents runtime and fix mypy (#25)

Author: cotovanu-cristian

pyproject.toml

@@ -23,7 +23,7 @@ maintainers = [
 register = "uipath_agents.middlewares:register_middleware"
 
 [project.entry-points."uipath.runtime.factories"]
-langgraph = "uipath_agents.runtime:register_runtime_factory"
+agents = "uipath_agents.runtime:register_runtime_factory"
 
 [project.urls]
 Homepage = "https://uipath.com"

src/uipath_agents/__init__.py

@@ -1,5 +1,4 @@
 from .middlewares import register_middleware
+from .runtime import register_runtime_factory
 
-__all__ = [
-    "register_middleware",
-]
+__all__ = ["register_middleware", "register_runtime_factory"]

src/uipath_agents/_cli/cli_debug.py

@@ -3,8 +3,7 @@
 from typing import Optional
 
 from dotenv import load_dotenv
-from uipath._cli._debug._bridge import ConsoleDebugBridge, get_debug_bridge
-from uipath._cli._utils._debug import setup_debugging
+from uipath._cli._debug._bridge import get_debug_bridge
 from uipath._cli._utils._studio_project import StudioClient
 from uipath._cli.middlewares import MiddlewareResult
 from uipath._config import UiPathConfig
@@ -18,10 +17,7 @@
     UiPathRuntimeFactoryProtocol,
     UiPathRuntimeFactoryRegistry,
     UiPathRuntimeProtocol,
-    UiPathRuntimeResult,
-    UiPathStreamOptions,
 )
-from uipath.runtime.events import UiPathRuntimeStateEvent
 from uipath.tracing import LlmOpsHttpExporter
 from uipath_langchain._cli._runtime._exception import LangGraphRuntimeError
 
@@ -33,64 +29,15 @@
 logger = logging.getLogger(__name__)
 
 
-async def execute_runtime(ctx: UiPathRuntimeContext) -> UiPathRuntimeResult:
-    with ctx:
-        runtime: UiPathRuntimeProtocol | None = None
-        factory: UiPathRuntimeFactoryProtocol | None = None
-        try:
-            factory = UiPathRuntimeFactoryRegistry.get(context=ctx)
-            runtime = await factory.new_runtime(ctx.entrypoint, ctx.job_id or "default")
-            options = UiPathExecuteOptions(resume=ctx.resume)
-            ctx.result = await runtime.execute(input=ctx.get_input(), options=options)
-            return ctx.result
-        finally:
-            if runtime:
-                await runtime.dispose()
-            if factory:
-                await factory.dispose()
-
-
-async def debug_runtime(
-    ctx: UiPathRuntimeContext,
-) -> UiPathRuntimeResult | None:
-    with ctx:
-        runtime: UiPathRuntimeProtocol | None = None
-        factory: UiPathRuntimeFactoryProtocol | None = None
-        try:
-            factory = UiPathRuntimeFactoryRegistry.get(context=ctx)
-            runtime = await factory.new_runtime(ctx.entrypoint, "default")
-            debug_bridge: UiPathDebugBridgeProtocol = ConsoleDebugBridge()
-            await debug_bridge.emit_execution_started()
-            options = UiPathStreamOptions(resume=ctx.resume)
-            async for event in runtime.stream(ctx.get_input(), options=options):
-                if isinstance(event, UiPathRuntimeResult):
-                    await debug_bridge.emit_execution_completed(event)
-                    ctx.result = event
-                elif isinstance(event, UiPathRuntimeStateEvent):
-                    await debug_bridge.emit_state_update(event)
-            return ctx.result
-        finally:
-            if runtime:
-                await runtime.dispose()
-            if factory:
-                await factory.dispose()
-
-
 def agents_debug_middleware(
     entrypoint: Optional[str],
     input: Optional[str],
     resume: bool,
     input_file: Optional[str],
     output_file: Optional[str],
-    debug: bool,
-    debug_port: int,
     **kwargs,
 ) -> MiddlewareResult:
-    """Middleware to handle LangGraph execution"""
-
-    if not setup_debugging(debug, debug_port):
-        logger.error(f"Failed to start debug server on port {debug_port}")
-
+    """Middleware to handle Agents LangGraph execution"""
     _prepare_agent_run_files()
 
     try:
@@ -99,6 +46,7 @@ async def execute_debug_runtime():
             trace_manager = UiPathTraceManager()
 
             with UiPathRuntimeContext.with_defaults(
+                entrypoint=entrypoint,
                 input=input,
                 input_file=input_file,
                 output_file=output_file,
@@ -117,7 +65,7 @@ async def execute_debug_runtime():
                     factory = UiPathRuntimeFactoryRegistry.get(context=ctx)
 
                     runtime = await factory.new_runtime(
-                        entrypoint, ctx.job_id or "default"
+                        ctx.entrypoint, ctx.job_id or "agents"
                     )
 
                     debug_bridge: UiPathDebugBridgeProtocol = get_debug_bridge(ctx)

src/uipath_agents/_cli/cli_run.py

@@ -5,7 +5,6 @@
 from dotenv import load_dotenv
 from uipath._cli._debug._bridge import ConsoleDebugBridge
 from uipath._cli._utils._common import read_resource_overwrites_from_file
-from uipath._cli._utils._debug import setup_debugging
 from uipath._cli.middlewares import MiddlewareResult
 from uipath._utils._bindings import ResourceOverwritesContext
 from uipath.core import UiPathTraceManager
@@ -36,6 +35,9 @@ async def execute_runtime(ctx: UiPathRuntimeContext) -> UiPathRuntimeResult:
         runtime: UiPathRuntimeProtocol | None = None
         factory: UiPathRuntimeFactoryProtocol | None = None
         try:
+            if ctx.entrypoint is None:
+                raise ValueError("Entrypoint is required for runtime execution")
+
             factory = UiPathRuntimeFactoryRegistry.get(context=ctx)
             runtime = await factory.new_runtime(ctx.entrypoint, ctx.job_id or "default")
             options = UiPathExecuteOptions(resume=ctx.resume)
@@ -55,8 +57,10 @@ async def debug_runtime(
         runtime: UiPathRuntimeProtocol | None = None
         factory: UiPathRuntimeFactoryProtocol | None = None
         try:
+            if ctx.entrypoint is None:
+                raise ValueError("Entrypoint is required for runtime debugging")
             factory = UiPathRuntimeFactoryRegistry.get(context=ctx)
-            runtime = await factory.new_runtime(ctx.entrypoint, "default")
+            runtime = await factory.new_runtime(ctx.entrypoint, "agents")
             debug_bridge: UiPathDebugBridgeProtocol = ConsoleDebugBridge()
             await debug_bridge.emit_execution_started()
             options = UiPathStreamOptions(resume=ctx.resume)
@@ -81,15 +85,9 @@ def agents_run_middleware(
     input_file: Optional[str],
     output_file: Optional[str],
     trace_file: Optional[str],
-    debug: bool,
-    debug_port: int,
     **kwargs,
 ) -> MiddlewareResult:
-    """Middleware to handle LangGraph execution"""
-
-    if not setup_debugging(debug, debug_port):
-        logger.error(f"Failed to start debug server on port {debug_port}")
-
+    """Middleware to handle Agents LangGraph execution"""
     _prepare_agent_run_files()
 
     try:

src/uipath_agents/_cli/constants.py

@@ -1,2 +1,2 @@
-AGENT_FILENAME = "agent.json"
+AGENT_ENTRYPOINT = "agent.json"
 AGENT_BUILDER_FILENAME = ".agent-builder"

src/uipath_agents/_cli/runtime.py

@@ -0,0 +1,79 @@
+import os
+
+from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
+from opentelemetry.instrumentation.asyncio import AsyncioInstrumentor
+from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+from opentelemetry.instrumentation.sqlite3 import SQLite3Instrumentor
+from uipath.runtime import UiPathRuntimeContext, UiPathRuntimeProtocol
+from uipath_langchain._cli._runtime._factory import LangGraphRuntimeFactory
+
+from ..constants import AGENT_ENTRYPOINT
+from .runtime import create_agent_langgraph_runtime
+
+_telemetry_initialized = False
+
+
+def _configure_agents_telemetry() -> None:
+    """Configure telemetry for agents. Idempotent - only runs once."""
+    global _telemetry_initialized
+    if _telemetry_initialized:
+        return
+
+    os.environ.setdefault("OTEL_SERVICE_NAME", "uipath-agents")
+    AsyncioInstrumentor().instrument()
+    HTTPXClientInstrumentor().instrument()
+    AioHttpClientInstrumentor().instrument()
+    SQLite3Instrumentor().instrument()
+    _telemetry_initialized = True
+
+
+class AgentRuntimeFactory(LangGraphRuntimeFactory):
+    """Factory for creating Agent runtimes from agent.json configuration."""
+
+    def __init__(self, context: UiPathRuntimeContext) -> None:
+        """Initialize the factory.
+
+        Args:
+            context: UiPathRuntimeContext to use for runtime creation
+        """
+        super().__init__(context)
+        _configure_agents_telemetry()
+
+    def discover_entrypoints(self) -> list[str]:
+        """Discover the Agent entrypoint agent.json"""
+        return [AGENT_ENTRYPOINT]
+
+    async def discover_runtimes(self) -> list[UiPathRuntimeProtocol]:
+        """Discover runtime instances for all entrypoints.
+
+        Returns:
+            List of AgentLangGraphRuntime instances
+        """
+        entrypoints = self.discover_entrypoints()
+        memory = await self._get_memory()
+
+        runtimes: list[UiPathRuntimeProtocol] = []
+        for entrypoint in entrypoints:
+            runtime = create_agent_langgraph_runtime(
+                entrypoint, entrypoint, self.context, memory
+            )
+            runtimes.append(runtime)
+
+        return runtimes
+
+    async def new_runtime(
+        self, entrypoint: str, runtime_id: str
+    ) -> UiPathRuntimeProtocol:
+        """Create a new Agent runtime instance.
+
+        Args:
+            entrypoint: Agent entrypoint (agent.json)
+            runtime_id: Unique identifier for the runtime instance
+
+        Returns:
+            Configured AgentLangGraphRuntime instance
+        """
+        memory = await self._get_memory()
+        return create_agent_langgraph_runtime(
+            runtime_id, entrypoint, self.context, memory
+        )

src/uipath_agents/_cli/runtime/factory.py

@@ -0,0 +1,80 @@
+import logging
+from pathlib import Path
+from typing import Any
+
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+from uipath.runtime import (
+    UiPathRuntimeContext,
+)
+from uipath.runtime.schema import UiPathRuntimeSchema
+from uipath_langchain._cli._runtime._runtime import (
+    LangGraphRuntime,
+)
+
+from uipath_agents.agent_graph_builder import build_agent_graph
+
+from ..utils import load_agent_configuration
+from .utils import validate_json_against_json_schema
+
+logger = logging.getLogger(__name__)
+
+
+class AgentLangGraphRuntime(LangGraphRuntime):
+    """Agents runtime extending LangGraph base runtime."""
+
+    def __init__(
+        self,
+        runtime_id: str,
+        graph_resolver: Any,
+        memory: AsyncSqliteSaver,
+        entrypoint: str,
+    ) -> None:
+        super().__init__(runtime_id, graph_resolver, memory)
+        self.entrypoint = entrypoint
+
+    async def get_schema(self) -> UiPathRuntimeSchema:
+        """Return the runtime schema for this agent."""
+        agent_json_path = Path.cwd() / self.entrypoint
+        agent_definition = load_agent_configuration(agent_json_path)
+
+        return UiPathRuntimeSchema(
+            filePath=self.entrypoint,
+            uniqueId=self.runtime_id,
+            type="agent",
+            input=agent_definition.input_schema or {},
+            output=agent_definition.output_schema or {},
+        )
+
+
+def create_agent_langgraph_runtime(
+    runtime_id: str,
+    entrypoint: str,
+    ctx: UiPathRuntimeContext,
+    memory: AsyncSqliteSaver,
+) -> AgentLangGraphRuntime:
+    """Create runtime for agents with input validation.
+
+    Args:
+        runtime_id: Unique identifier for the runtime instance
+        entrypoint: Agent file path containing the Agent definition
+        ctx: Runtime context containing input, resume flag, and other metadata
+        memory: AsyncSqliteSaver instance for checkpoint/state management
+
+    Returns:
+        AgentLangGraphRuntime instance configured for the agent
+    """
+
+    async def graph_resolver():
+        """Load agent config, validate input on new runs, and build graph."""
+        agent_json_path = Path.cwd() / entrypoint
+        agent_definition = load_agent_configuration(agent_json_path)
+
+        agent_input: dict[str, Any] = {}
+        if not ctx.resume:
+            agent_input = validate_json_against_json_schema(
+                agent_definition.input_schema, ctx.input
+            )
+
+        return await build_agent_graph(agent_definition, input_data=agent_input)
+
+    return AgentLangGraphRuntime(runtime_id, graph_resolver, memory, entrypoint)

src/uipath_agents/_cli/runtime/runtime.py

@@ -3,7 +3,7 @@
 from pydantic import ValidationError
 from uipath.utils.dynamic_schema import jsonschema_to_pydantic
 
-from .exceptions import InputValidationError
+from ..exceptions import InputValidationError
 
 
 def validate_json_against_json_schema(