diff --git a/src/uipath_langchain/agent/advanced/__init__.py b/src/uipath_langchain/agent/advanced/__init__.py index ad1ecd9a6..dc1946b5c 100644 --- a/src/uipath_langchain/agent/advanced/__init__.py +++ b/src/uipath_langchain/agent/advanced/__init__.py @@ -14,6 +14,7 @@ build_code_interpreter_middleware, ptc_tool_names, subagent_dispatch_is_replay_safe, + warm_code_interpreter, ) from .types import AdvancedAgentGraphState, ConversationalAdvancedAgentGraphState from .utils import ( @@ -42,4 +43,5 @@ "create_state_with_input", "ptc_tool_names", "subagent_dispatch_is_replay_safe", + "warm_code_interpreter", ] diff --git a/src/uipath_langchain/agent/advanced/code_interpreter.py b/src/uipath_langchain/agent/advanced/code_interpreter.py index ab8319a57..08108250a 100644 --- a/src/uipath_langchain/agent/advanced/code_interpreter.py +++ b/src/uipath_langchain/agent/advanced/code_interpreter.py @@ -17,6 +17,7 @@ """ import atexit +import functools import logging from collections.abc import Iterable, Sequence from typing import Any, Literal, get_args @@ -214,6 +215,7 @@ def build_code_interpreter_middleware( ImportError: If the ``code-interpreter`` extra is not installed. """ middleware_cls = _code_interpreter_middleware_cls() + warm_code_interpreter() exposed = ptc_tool_names(tools) dispatch = subagent_dispatch_is_replay_safe(subagents, tools) logger.info( @@ -234,6 +236,32 @@ def build_code_interpreter_middleware( return [middleware] +@functools.cache +def warm_code_interpreter() -> None: + """Compile the interpreter's WebAssembly modules before any eval deadline is armed. + + quickjs_rs compiles its source-transform module lazily, inside the first eval of + the process, and that eval runs under the same per-call deadline as user code. + On a CPU-starved instance the compile alone can outlast the deadline, which + fails the run at its first model call, before the model ever asks for ``eval``. + The compiled modules are cached for the life of the process, so one call here + serves every REPL built afterwards. A host that preloads modules at process + start can call it then. Without the ``code-interpreter`` extra it does nothing. + """ + try: + from quickjs_rs import Runtime, SourceTransform, transform_source + except ImportError: + return + transform_source( + "warmup.js", "const x = 1;", flags=SourceTransform.TOP_LEVEL_CONST_TO_VAR + ) + runtime = Runtime() + try: + runtime.new_context().close() + finally: + runtime.close() + + def _close_at_exit(middleware: Any) -> None: """Close the REPL registry from ``atexit`` instead of leaving it to ``__del__``. diff --git a/tests/agent/advanced/test_code_interpreter.py b/tests/agent/advanced/test_code_interpreter.py index 31a5ffb64..ae97d87ee 100644 --- a/tests/agent/advanced/test_code_interpreter.py +++ b/tests/agent/advanced/test_code_interpreter.py @@ -29,6 +29,7 @@ create_advanced_agent, ptc_tool_names, subagent_dispatch_is_replay_safe, + warm_code_interpreter, ) from uipath_langchain.agent.advanced.code_interpreter import ( EVAL_TOOL_NAME, @@ -293,6 +294,58 @@ def test_factory_closes_the_repl_registry_at_exit( assert registered == [cast(Any, middleware)._registry.close] +def test_factory_warms_the_interpreter_before_building_the_repl( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The WebAssembly compile happens at graph build, outside any eval deadline. + + Upstream compiles its source-transform module inside the first eval, under + the per-call deadline. On a slow instance that compile alone exceeds the + deadline and the first model call fails, so the factory must pay it first. + """ + calls: list[str] = [] + monkeypatch.setattr( + "uipath_langchain.agent.advanced.code_interpreter.warm_code_interpreter", + lambda: calls.append("warm"), + ) + + build_code_interpreter_middleware([_tool("read_invoice")]) + + assert calls == ["warm"] + + +def test_warm_up_compiles_the_transform_module_upstream_uses_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import quickjs_rs + + seen: list[Any] = [] + + def fake_transform_source(name: str, source: str, *, flags: Any = None) -> str: + seen.append(flags) + return source + + monkeypatch.setattr(quickjs_rs, "transform_source", fake_transform_source) + warm_code_interpreter.cache_clear() + + warm_code_interpreter() + warm_code_interpreter() + + assert seen == [quickjs_rs.SourceTransform.TOP_LEVEL_CONST_TO_VAR] + warm_code_interpreter.cache_clear() + + +def test_warm_up_without_the_extra_is_a_no_op( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem(sys.modules, "quickjs_rs", None) + warm_code_interpreter.cache_clear() + + warm_code_interpreter() + + warm_code_interpreter.cache_clear() + + def test_factory_without_the_extra_raises_install_guidance( monkeypatch: pytest.MonkeyPatch, ) -> None: