From d84fa9170665c6149c0297354ba7089b511a5a0b Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Tue, 1 Sep 2026 09:53:33 -0400 Subject: [PATCH 1/4] Add standalone agent harness SDK --- README.md | 9 +- agents/mcp/src/hyperforge_mcp/agent.py | 1 + .../src/hyperforge_perplexity_search/agent.py | 1 + .../src/hyperforge_restricted/agent.py | 71 +- .../src/hyperforge_restricted/config.py | 10 - .../src/hyperforge_restricted/sandbox.py | 378 ---- .../src/hyperforge_restricted/worker.py | 170 -- agents/restricted/tests/test_sandbox.py | 70 +- docs/harness-sdk.md | 485 ++++++ hyperforge/pyproject.toml | 4 +- .../src/hyperforge/codemode/__init__.py | 39 + .../src/hyperforge/codemode}/model.py | 77 +- hyperforge/src/hyperforge/codemode/sandbox.py | 426 +++++ hyperforge/src/hyperforge/codemode/worker.py | 241 +++ hyperforge/src/hyperforge/definition.py | 3 + hyperforge/src/hyperforge/engine.py | 21 +- hyperforge/src/hyperforge/harness/__init__.py | 4 + hyperforge/src/hyperforge/harness/agent.py | 117 ++ hyperforge/src/hyperforge/harness/config.py | 50 + .../src/hyperforge/harness_sdk/__init__.py | 92 + .../src/hyperforge/harness_sdk/agents.py | 297 ++++ .../src/hyperforge/harness_sdk/clients.py | 458 +++++ .../src/hyperforge/harness_sdk/context.py | 74 + .../src/hyperforge/harness_sdk/harness.py | 1100 ++++++++++++ .../src/hyperforge/harness_sdk/models.py | 107 ++ .../src/hyperforge/harness_sdk/schema.py | 45 + .../src/hyperforge/harness_sdk/storage.py | 95 + .../hyperforge/harness_sdk/tools/__init__.py | 8 + .../hyperforge/harness_sdk/tools/codemode.py | 110 ++ .../src/hyperforge/harness_sdk/tools/core.py | 230 +++ .../src/hyperforge/harness_sdk/tools/utils.py | 126 ++ .../src/hyperforge/harness_sdk/usage.py | 34 + hyperforge/src/hyperforge/llm.py | 75 +- hyperforge/src/hyperforge/server/sandbox.py | 22 +- hyperforge/tests/codemode/test_model.py | 29 + hyperforge/tests/codemode/test_protocol.py | 21 + .../tests/harness/test_harness_agent.py | 122 ++ hyperforge/tests/harness_sdk/test_clients.py | 241 +++ hyperforge/tests/harness_sdk/test_codemode.py | 182 ++ hyperforge/tests/harness_sdk/test_harness.py | 1525 +++++++++++++++++ uv.lock | 19 +- 41 files changed, 6520 insertions(+), 669 deletions(-) delete mode 100644 agents/restricted/src/hyperforge_restricted/sandbox.py delete mode 100644 agents/restricted/src/hyperforge_restricted/worker.py create mode 100644 docs/harness-sdk.md create mode 100644 hyperforge/src/hyperforge/codemode/__init__.py rename {agents/restricted/src/hyperforge_restricted => hyperforge/src/hyperforge/codemode}/model.py (67%) create mode 100644 hyperforge/src/hyperforge/codemode/sandbox.py create mode 100644 hyperforge/src/hyperforge/codemode/worker.py create mode 100644 hyperforge/src/hyperforge/harness/__init__.py create mode 100644 hyperforge/src/hyperforge/harness/agent.py create mode 100644 hyperforge/src/hyperforge/harness/config.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/__init__.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/agents.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/clients.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/context.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/harness.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/models.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/schema.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/storage.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/tools/__init__.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/tools/codemode.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/tools/core.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/tools/utils.py create mode 100644 hyperforge/src/hyperforge/harness_sdk/usage.py create mode 100644 hyperforge/tests/codemode/test_model.py create mode 100644 hyperforge/tests/codemode/test_protocol.py create mode 100644 hyperforge/tests/harness/test_harness_agent.py create mode 100644 hyperforge/tests/harness_sdk/test_clients.py create mode 100644 hyperforge/tests/harness_sdk/test_codemode.py create mode 100644 hyperforge/tests/harness_sdk/test_harness.py diff --git a/README.md b/README.md index 0a3adfc3..ed29d2a3 100644 --- a/README.md +++ b/README.md @@ -20,13 +20,18 @@ NucliaDB Agentic API service, and a set of reusable Hyperforge agent packages. ## Packages -The workspace is managed with `uv` and Python 3.10 or newer. +The workspace is managed with `uv` and Python 3.12 or newer. Core packages: - `hyperforge`: agentic framework runtime and API. - `nucliadb_agentic_api`: NucliaDB Agentic API service. +Framework guides: + +- [Harness SDK](docs/harness-sdk.md): build asynchronous tool-loop agents with + streaming events, persistence, limits, feedback, memory, and sub-agents. + Agent packages: - `hyperforge_conditional` @@ -120,4 +125,4 @@ Contributions are welcome. Please format the code with `make fmt`, run [Apache-2.0](LICENSE) -Disclaimer: _This is not an officially supported Progress Software Corporation product._ \ No newline at end of file +Disclaimer: _This is not an officially supported Progress Software Corporation product._ diff --git a/agents/mcp/src/hyperforge_mcp/agent.py b/agents/mcp/src/hyperforge_mcp/agent.py index 0e7ae3a3..9a998a46 100644 --- a/agents/mcp/src/hyperforge_mcp/agent.py +++ b/agents/mcp/src/hyperforge_mcp/agent.py @@ -901,6 +901,7 @@ async def preload(self, manager: Manager, memory: QuestionMemory) -> None: name=tool.name, description=tool.description or tool.name, parameters=_tool_parameters(tool.inputSchema), + input_schema=tool.inputSchema, ) setattr(self, tool.name, self._make_tool_caller(tool.name)) diff --git a/agents/perplexity_search/src/hyperforge_perplexity_search/agent.py b/agents/perplexity_search/src/hyperforge_perplexity_search/agent.py index 82387249..6f77c6ea 100644 --- a/agents/perplexity_search/src/hyperforge_perplexity_search/agent.py +++ b/agents/perplexity_search/src/hyperforge_perplexity_search/agent.py @@ -28,6 +28,7 @@ class PerplexitySearchAgent(ContextAgent, Agent[PerplexitySearchAgentConfig]): __published_functions__: ClassVar[Dict[str, FunctionDefinition]] = { "internet_search": FunctionDefinition( name="internet_search", + method="search", description="Performs an internet search using Perplexity and returns the results as context to answer questions. Does not generate an answer, only search and return the results as context.", parameters={ "question": { diff --git a/agents/restricted/src/hyperforge_restricted/agent.py b/agents/restricted/src/hyperforge_restricted/agent.py index af9545d0..8f0ece8f 100644 --- a/agents/restricted/src/hyperforge_restricted/agent.py +++ b/agents/restricted/src/hyperforge_restricted/agent.py @@ -2,7 +2,7 @@ import asyncio import functools -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from concurrent.futures import ProcessPoolExecutor from copy import deepcopy from dataclasses import dataclass, field from time import time @@ -10,6 +10,13 @@ from hyperforge import logger from hyperforge.agent import Agent +from hyperforge.codemode import ( + RestrictedPythonTask, + SandboxRunner, + WorkerExecutionRequest, + WorkerTypes, +) +from hyperforge.codemode import settings as sandbox_settings from hyperforge.configure import agent from hyperforge.context.agent import ContextAgent, build_context_agent from hyperforge.db.agents import SERVICE_NAME @@ -35,14 +42,6 @@ EXTRACT_AGENT_TEMPLATE, TRANSFORM_REPHRASE, ) -from hyperforge_restricted.model import ( - RestrictedPythonTask, - WorkerExecutionRequest, - WorkerTypes, -) - -from .sandbox import SandboxRunner -from .sandbox import settings as sandbox_settings def tracer(): @@ -81,6 +80,23 @@ class Output: final_answer: Optional[str] = None +CONTROLLER_FUNCTIONS = { + "add_step", + "choose", + "context_match_question", + "extract", + "rephrase", + "remi_validation", + "save", + "save_context", + "set_context", + "summary", + "transform", + "user_validation", + "validate", +} + + if sandbox_settings.sandbox_socket is None: import os @@ -330,19 +346,11 @@ def _runner(self, memory: QuestionMemory, manager: Manager): sandbox_settings.sandbox_socket, functools.partial(self.handle_queue_item, manager, memory), ) - elif self.config.debug: - # Debug mode, run in the same process (another thread) - return SandboxRunner.with_pool( - ThreadPoolExecutor(1), - functools.partial(self.handle_queue_item, manager, memory), - debug=True, - ) - else: - # Process pool mode, run in a separate process for isolation - return SandboxRunner.with_pool( - SANDBOX_POOL, - functools.partial(self.handle_queue_item, manager, memory), - ) + # Process pool mode, run in a separate process for isolation + return SandboxRunner.with_pool( + SANDBOX_POOL, + functools.partial(self.handle_queue_item, manager, memory), + ) async def _preload_child_agents( self, manager: Manager, memory: QuestionMemory @@ -419,11 +427,6 @@ async def _get_question_context( local_vars["needs_rephrase"] = ( True if depth == 0 and self.config.needs_rephrase else False ) - if self.config.debug: - logger.debug( - f"Executing restricted code, loop {depth}, question: {question}" - ) - tasks.append( asyncio.create_task( self._runner(memory, manager).run( @@ -507,6 +510,20 @@ async def handle_queue_item( ) -> WorkerTypes: contexts: List[Context] = [] + if not ( + item.agent == "_controller" and item.function == "save" + ) and item.function not in self.function_names.get(item.agent, {}): + raise ValueError( + f"Function {item.function!r} is not authorized for agent {item.agent!r}" + ) + if item.function in CONTROLLER_FUNCTIONS and item.agent not in { + "_controller", + "self", + }: + raise ValueError( + f"Function {item.function!r} is reserved for the restricted agent" + ) + with tracer().start_as_current_span(f"RestrictedAgent.{item.function}"): try: item.keyword_args.update( diff --git a/agents/restricted/src/hyperforge_restricted/config.py b/agents/restricted/src/hyperforge_restricted/config.py index c8be6419..8bed4552 100644 --- a/agents/restricted/src/hyperforge_restricted/config.py +++ b/agents/restricted/src/hyperforge_restricted/config.py @@ -40,16 +40,6 @@ class PythonAgentConfig(ContextAgentConfig): }, ) - debug: bool = Field( - False, - title="Debug mode", - description="Enable debug mode to get more information about the execution", - json_schema_extra={ - "show_in_node": False, - "widget": WidgetType.NOT_SHOWN, - }, - ) - decision_model: LLMField = Field( default=LLMConfig(model_id=llm_defaults.reasoning), title="Generative model", diff --git a/agents/restricted/src/hyperforge_restricted/sandbox.py b/agents/restricted/src/hyperforge_restricted/sandbox.py deleted file mode 100644 index 7b5366b1..00000000 --- a/agents/restricted/src/hyperforge_restricted/sandbox.py +++ /dev/null @@ -1,378 +0,0 @@ -import asyncio -import json -import logging -import os -from concurrent.futures import Executor, ThreadPoolExecutor -from multiprocessing import Pipe, Process -from multiprocessing.connection import Connection -from typing import Any, Callable, Coroutine - -import httpx -import nucliadb_telemetry.metrics -import prometheus_client -from pydantic_settings import BaseSettings - -from hyperforge_restricted.model import ( - RestrictedPythonTask, - SandboxMessage, - WorkerExecutionRequest, - WorkerTypes, -) -from hyperforge_restricted.worker import PythonAgentWorker - -logger = logging.getLogger("hyperforge_restricted_sandbox") - -# Maximum time in seconds that the worker can execute CPU bound work. This is reset every time the worker -# calls any RAO function (another agent, LLM questions, etc.). It should be kept really low to avoid -# denial of service. It can be as low as a few milliseconds, it probably should not exceed 1 second. -WORKER_CPU_LIMIT = 1 - - -class SandboxSettings(BaseSettings): - # Verify that the sandbox is secure on startup or raise an exception - sandbox_verify: bool = True - - # Socket for the sandbox server to listen to - sandbox_socket: str | None = None - - # Port for prometheus metrics - sandbox_metrics_port: int = 8091 - - -settings = SandboxSettings() - - -class InsecureSandbox(Exception): - pass - - -class SandboxRunner: - pool: Executor | None - socket: str | None - callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]] - debug: bool - - @classmethod - def with_pool( - cls, - pool: Executor, - callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], - debug: bool = False, - ) -> "SandboxRunner": - return cls(pool=pool, socket=None, callback=callback, debug=debug) - - @classmethod - def remote( - cls, - socket: str, - callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], - debug: bool = False, - ) -> "SandboxRunner": - return cls(pool=None, socket=socket, callback=callback, debug=debug) - - def __init__( - self, - pool: Executor | None, - socket: str | None, - callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], - debug: bool, - ): - self.pool = pool - self.socket = socket - self.callback = callback - self.debug = debug - - async def run(self, request: WorkerExecutionRequest): - if self.pool is not None: - return await self._run_in_pool(request) - elif self.socket is not None: - return await self._run_remotely(request) - else: - raise ValueError( - "SandboxRunner must be initialized with a pool or a socket" - ) - - async def _run_in_pool(self, request: WorkerExecutionRequest): - """ - Runs the two tasks needed to run a sandboxed Python agent: - - _process_question_context_sync runs the untrusted code in a process pool executor - under RestrictedPython and has limited I/O. - - agent_runner_task calls read_write_queue and starts a "server" listening for messages - from the sandbox. It then executes the desired function in behalf of the user provided - code and sends the results back over a Pipe. - """ - loop = asyncio.get_running_loop() - (pipe_worker, pipe_restricted) = Pipe() - worker = PythonAgentWorker(pipe_worker, debug=self.debug) - agent_runner_task = asyncio.create_task(self.background_task(pipe_restricted)) - await loop.run_in_executor( - self.pool, - worker._process_question_context_sync, - request.code, - request.question, - request.local_vars, - request.global_vars, - request.function_names, - ) - # _process_question_context_sync always sends a message to the pipe on exit (finally block) - # this causes the agent_runner_task to finish cleanly. Otherwise (e.g: if we call cancel on the task) - # we would be leaking threads from the executor used for agent_runner_tasks - await agent_runner_task - - async def _run_remotely(self, request: WorkerExecutionRequest): - """ - Runs the code in a remote sandbox server (another container in the same machine) - Communicates via Unix socket with the sandbox server and runs agent tasks for it - until the code finishes - """ - rx, tx = await asyncio.open_unix_connection(self.socket) - reader, writer = SandboxReader(rx), SandboxWriter(tx) - await writer.write_message(SandboxMessage.Run(run=request)) - while True: - try: - msg = await reader.read_message() - except asyncio.IncompleteReadError: - raise RuntimeError("Sandbox connection closed unexpectedly") - - if isinstance(msg, SandboxMessage.Done): - break - elif isinstance(msg, SandboxMessage.Error): - raise RuntimeError(f"Python agent error: {msg.error}") - - assert isinstance(msg, SandboxMessage.Request) - - try: - response = await self.callback(msg.task) - except Exception as e: - # On error, close the connection so the server can terminate the process - writer.close() - raise e - - await writer.write_message(SandboxMessage.Response(result=response)) - - writer.close() - - def run_in_process(self, request: WorkerExecutionRequest): - """ - Runs the two tasks needed to run a sandboxed Python agent: - - _process_question_context_sync runs the untrusted code in a process under RestrictedPython - - agent_runner_task calls read_write_queue and starts a "server" listening for messages - from the sandbox. It then executes the desired function in behalf of the user provided - code and sends the results back over a Pipe. - """ - (pipe_worker, pipe_restricted) = Pipe() - worker = PythonAgentWorker(pipe_worker, debug=self.debug) - agent_runner_task = asyncio.create_task(self.background_task(pipe_restricted)) - process = Process( - target=worker._process_question_context_sync, - args=( - request.code, - request.question, - request.local_vars, - request.global_vars, - request.function_names, - ), - ) - process.start() - - return agent_runner_task, process - - async def background_task(self, pipe: Connection): - loop = asyncio.get_running_loop() - await loop.run_in_executor( - ThreadPoolExecutor(max_workers=1), - self.read_write_queue, - pipe, - loop, - ) - - def read_write_queue( - self, - pipe: Connection, - loop: asyncio.AbstractEventLoop, - ): - """ - This is a sync function that runs on a thread to read from the pipe that - communicates with the RestrictedPython process. It then calls an async callback - which can be the actual function in the agent, or a function in SandboxSession - that sends a socket message to the actual agent. - """ - try: - while True: - item = pipe.recv() - if item.function == "_error": - raise Exception(item.args[0]) - elif item.function == "_close": - break - future = asyncio.run_coroutine_threadsafe(self.callback(item), loop) - result = future.result() - pipe.send(result) - except EOFError: - pass - - -def _url_reachable(url: str) -> bool: - with httpx.Client() as client: - try: - client.get(url) - return True - except httpx.ConnectError: - return False - - -def _verify_connectivity(): - kubernetes_host = os.environ.get("KUBERNETES_SERVICE_HOST") - if kubernetes_host and _url_reachable(f"https://{kubernetes_host}"): - raise InsecureSandbox() - if _url_reachable("https://nuclia.com"): - raise InsecureSandbox() - - -class SandboxWriter: - def __init__(self, writer: asyncio.StreamWriter): - self.writer = writer - - async def _write_packet(self, msg: dict[str, Any]): - msg_bytes = json.dumps(msg).encode("utf-8") - self.writer.write(len(msg_bytes).to_bytes(4, "little")) - self.writer.write(msg_bytes) - await self.writer.drain() - - async def write_message(self, message: SandboxMessage.AnyMessage): - await self._write_packet(SandboxMessage.serialize(message)) - - def close(self): - self.writer.close() - - -class SandboxReader: - def __init__(self, reader: asyncio.StreamReader): - self.reader = reader - - async def _read_packet(self) -> dict[str, Any]: - len_bytes = await self.reader.readexactly(4) - msg_len = int.from_bytes(len_bytes, "little") - msg_bytes = await self.reader.readexactly(msg_len) - msg = json.loads(msg_bytes.decode("utf-8")) - - return msg - - async def read_message(self) -> SandboxMessage.AnyMessage: - return SandboxMessage.parse(await self._read_packet()) - - -class SandboxSession: - reader: SandboxReader - writer: SandboxWriter - runner: SandboxRunner - task: asyncio.Task[None] - process: Process - waiting_for_agent: asyncio.Task[None] | None - - def __init__( - self, - reader: SandboxReader, - writer: SandboxWriter, - ): - self.reader = reader - self.writer = writer - self.runner = SandboxRunner( - pool=None, socket=None, callback=self._callback, debug=False - ) - - async def run(self, request: WorkerExecutionRequest): - self.task, self.process = self.runner.run_in_process(request) - try: - await self._start_timeout() - await self.task - await self._stop_timeout() - await self.writer.write_message(SandboxMessage.Done()) - self.process.join(timeout=1) - if self.process.exitcode is None: - logger.warning("Sandbox process did not exit cleanly, killing it") - self.process.kill() - except Exception as e: - logger.error(f"Error in sandbox session: {e}") - self.process.kill() - self.process.join() - self.task.cancel() - try: - await self.writer.write_message(SandboxMessage.Error(error=str(e))) - except Exception: - pass - - async def _callback(self, item: RestrictedPythonTask) -> WorkerTypes: - await self._stop_timeout() - await self.writer.write_message(SandboxMessage.Request(task=item)) - - try: - response = await self.reader.read_message() - except asyncio.IncompleteReadError: - logger.warning("Got disconnected from agent") - self.process.kill() - self.process.join() - self.task.cancel() - # Irrelevant since the process has just been killed - return None - - assert isinstance(response, SandboxMessage.Response) - await self._start_timeout() - return response.result - - async def _start_timeout(self): - self.waiting_for_agent = asyncio.create_task(self._timeout()) - - async def _stop_timeout(self): - if self.waiting_for_agent: - self.waiting_for_agent.cancel() - self.waiting_for_agent = None - - async def _timeout(self): - await asyncio.sleep(WORKER_CPU_LIMIT) - await self.writer.write_message( - SandboxMessage.Error(error="Python agent timeout") - ) - self.process.kill() - - -agents_running = prometheus_client.Gauge( - "arag_sandbox_agents_running", "Number of Python agents currently running" -) -execution_observer = nucliadb_telemetry.metrics.Observer("arag_sandbox_execution") - - -async def run_sandbox_server(): - assert settings.sandbox_socket is not None - if settings.sandbox_verify: - _verify_connectivity() - - async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): - reader = SandboxReader(rx) - writer = SandboxWriter(tx) - - # Get initial request - try: - async with asyncio.timeout(1): - msg = await reader.read_message() - assert isinstance(msg, SandboxMessage.Run) - except TimeoutError: - logger.warning("Timed out waiting for Run() message. Closing.") - return - - agents_running.inc() - observation = execution_observer() - observation.start() - try: - session = SandboxSession(reader, writer) - await session.run(msg.run) - except Exception as e: - observation.set_status("error") - await writer.write_message(SandboxMessage.Error(error=str(e))) - finally: - agents_running.dec() - observation.end() - writer.close() - - server = await asyncio.start_unix_server(handler, settings.sandbox_socket) - async with server: - await server.serve_forever() diff --git a/agents/restricted/src/hyperforge_restricted/worker.py b/agents/restricted/src/hyperforge_restricted/worker.py deleted file mode 100644 index 669454f0..00000000 --- a/agents/restricted/src/hyperforge_restricted/worker.py +++ /dev/null @@ -1,170 +0,0 @@ -from dataclasses import dataclass -from functools import partial -from multiprocessing.connection import Connection -from typing import Any, Dict, List, Optional - -from hyperforge.definition import FunctionDefinition -from hyperforge.memory import Chunk -from hyperforge.memory.memory import Context -from RestrictedPython import compile_restricted -from RestrictedPython.Guards import ( - guarded_iter_unpack_sequence, - guarded_unpack_sequence, - safe_builtins, -) - -from hyperforge_restricted.model import RestrictedPythonTask - -EXAMPLES = { - "discourse": """ -import httpx -resp = get("discourse_url", param={q=question}).json() -for message in resp: - chunks.append(Chunk(chunk_id=message.id, text=message.text, labels=message.labels)) -""" -} - - -# Define allowed builtins and imports -allowed_builtins = { - **safe_builtins, # Default safe builtins - "__import__": lambda name, *args: ( - __import__(name) if name in {"math"} else None - ), # Restrict imports -} - -# Define restricted globals -restricted_globals = { - "__builtins__": allowed_builtins, # Restrict built-ins -} - - -class PythonAgentWorker: - context: Optional[Context] = None - - def __init__(self, pipe: Connection, debug: bool = False): - self.pipe = pipe - self.functions_agent_id: Dict[str, List[str]] = {} - self.debug = debug - - def _process_question_context_sync( - self, - code: str, - question: str, - local_vars: Dict[str, Any], - global_vars: Dict[str, Any], - function_names: Dict[str, Dict[str, FunctionDefinition]], - ): - try: - for agent, functions in function_names.items(): - for function in functions.keys(): - if function in global_vars: - self.functions_agent_id[function].append(agent) - else: - global_vars[function] = partial( - self.valid_functions_call, function - ) - self.functions_agent_id[function] = [agent] - - byte_code = compile_restricted(code, filename="", mode="exec") - - global_vars.update( - { - "__builtins__": safe_builtins, - "_getitem_": lambda obj, index: obj[index], - "_getiter_": iter, - "dataclass": dataclass, - "Chunk": Chunk, - "Context": Context, - "List": List, - "Any": Any, - "Dict": Dict, - "_iter_unpack_sequence_": guarded_iter_unpack_sequence, - "_unpack_sequence_": guarded_unpack_sequence, - "save": self.save, - "question": question, - } - ) - if self.debug: - global_vars["pdb"] = __import__("pdb") - - exec( - byte_code, - global_vars, - local_vars, - ) - except Exception as e: - self.pipe.send( - RestrictedPythonTask( - function="_error", - agent="_controller", - args=(str(e),), - keyword_args={}, - ) - ) - - self.pipe.send( - RestrictedPythonTask( - function="_close", - agent="_controller", - args=(None,), - keyword_args={}, - ) - ) - - def save( - self, - contexts: List[Context] = [], - resolved: bool = True, - missing: list[str] = [], - question: str = "", - final_answer: Optional[str] = None, - ): - self.pipe.send( - RestrictedPythonTask( - function="save", - agent="_controller", - args=(None,), - keyword_args={ - "contexts": contexts, - "resolved": resolved, - "missing": missing, - "question": question, - "final_answer": final_answer, - }, - ) - ) - return self.pipe.recv() - - def valid_functions_call( - self, - function_name: str, - *args: Any, - **kwargs: Any, - ) -> Any: - if "agent_id" in kwargs: - agent_id = kwargs.pop("agent_id") - else: - agent_id = None - if agent_id is not None: - if agent_id not in self.functions_agent_id.get(function_name, []): - raise Exception( - f"Agent ID ({agent_id}) not authorized for this function {function_name}" - ) - else: - agents = self.functions_agent_id.get(function_name, []) - if len(agents) == 0: - raise ValueError(f"No agent found for function {function_name}") - elif len(agents) > 1: - raise ValueError( - f"Multiple agents found for function {function_name}: {agents}, please specify agent_id" - ) - - agent_id = agents[0] - - self.pipe.send( - RestrictedPythonTask( - function=function_name, agent=agent_id, args=args, keyword_args=kwargs - ) - ) - return self.pipe.recv() diff --git a/agents/restricted/tests/test_sandbox.py b/agents/restricted/tests/test_sandbox.py index 8825468c..f5c35327 100644 --- a/agents/restricted/tests/test_sandbox.py +++ b/agents/restricted/tests/test_sandbox.py @@ -1,16 +1,19 @@ import asyncio +import os import tempfile from unittest.mock import patch import pytest -from hyperforge.definition import FunctionDefinition - -from hyperforge_restricted import sandbox -from hyperforge_restricted.model import ( +from hyperforge.codemode import sandbox +from hyperforge.codemode.model import ( RestrictedPythonTask, SandboxMessage, WorkerExecutionRequest, ) +from hyperforge.definition import FunctionDefinition + +from hyperforge_restricted.agent import PythonAgent +from hyperforge_restricted.config import PythonAgentConfig @pytest.fixture @@ -19,21 +22,29 @@ async def sandbox_conn(): socket = f"{path}/sandbox.sock" with ( patch( - "hyperforge_restricted.sandbox.settings.sandbox_socket", + "hyperforge.codemode.sandbox.settings.sandbox_socket", socket, ), patch( - "hyperforge_restricted.sandbox.settings.sandbox_verify", + "hyperforge.codemode.sandbox.settings.sandbox_verify", False, ), + patch( + "hyperforge.codemode.sandbox.settings.sandbox_token", + "test-token", + ), ): task = asyncio.create_task(sandbox.run_sandbox_server()) await asyncio.sleep(0.1) # Wait for the server to start + assert oct(os.stat(socket).st_mode & 0o777) == "0o600" (rx, tx) = await asyncio.open_unix_connection(socket) yield (sandbox.SandboxReader(rx), sandbox.SandboxWriter(tx)) + tx.close() + await tx.wait_closed() task.cancel() + await asyncio.gather(task, return_exceptions=True) @pytest.fixture @@ -43,7 +54,7 @@ async def long_timeout(): timeout under coverage. Use a larger limit here to avoid flaky failures. """ with patch( - "hyperforge_restricted.sandbox.WORKER_CPU_LIMIT", + "hyperforge.codemode.sandbox.WORKER_CPU_LIMIT", 15, ): yield @@ -55,7 +66,7 @@ async def test_sandbox_no_code(sandbox_conn, long_timeout): ) (rx, tx) = sandbox_conn - await tx.write_message(SandboxMessage.Run(request)) + await tx.write_message(SandboxMessage.Run(request, token="test-token")) message = await rx.read_message() assert isinstance(message, SandboxMessage.Done) @@ -67,7 +78,7 @@ async def test_sandbox_bad_code(sandbox_conn, long_timeout): ) (rx, tx) = sandbox_conn - await tx.write_message(SandboxMessage.Run(request)) + await tx.write_message(SandboxMessage.Run(request, token="test-token")) message = await rx.read_message() assert isinstance(message, SandboxMessage.Error) assert "'return' outside function" in message.error @@ -80,12 +91,29 @@ async def test_sandbox_infinite_loop(sandbox_conn): ) (rx, tx) = sandbox_conn - await tx.write_message(SandboxMessage.Run(request)) + await tx.write_message(SandboxMessage.Run(request, token="test-token")) message = await rx.read_message() assert isinstance(message, SandboxMessage.Error) assert "timeout" in message.error +async def test_sandbox_runtime_limit(sandbox_conn, long_timeout): + request = WorkerExecutionRequest( + code="while True: pass", + question="Q?", + local_vars={}, + global_vars={}, + function_names={}, + max_runtime_seconds=0.1, + ) + + (rx, tx) = sandbox_conn + await tx.write_message(SandboxMessage.Run(request, token="test-token")) + message = await rx.read_message() + assert isinstance(message, SandboxMessage.Error) + assert "timed out" in message.error + + async def test_sandbox_function_call(sandbox_conn, long_timeout): request = WorkerExecutionRequest( code="show(msg=foo())", @@ -112,7 +140,7 @@ async def test_sandbox_function_call(sandbox_conn, long_timeout): ) (rx, tx) = sandbox_conn - await tx.write_message(SandboxMessage.Run(request)) + await tx.write_message(SandboxMessage.Run(request, token="test-token")) message = await rx.read_message() assert isinstance(message, SandboxMessage.Request) @@ -130,3 +158,23 @@ async def test_sandbox_function_call(sandbox_conn, long_timeout): message = await rx.read_message() assert isinstance(message, SandboxMessage.Done) + + +def test_restricted_agent_config_cannot_enable_debug_mode(): + config = PythonAgentConfig.model_validate({"code": "save()", "debug": True}) + + assert "debug" not in config.model_dump() + + +async def test_restricted_agent_rejects_undeclared_callback(): + agent = PythonAgent.__new__(PythonAgent) + agent.function_names = {"child": {}} + task = RestrictedPythonTask( + function="private_function", + agent="child", + args=(), + keyword_args={}, + ) + + with pytest.raises(ValueError, match="not authorized"): + await agent.handle_queue_item(None, None, task) diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md new file mode 100644 index 00000000..cb6ff05c --- /dev/null +++ b/docs/harness-sdk.md @@ -0,0 +1,485 @@ +# Harness SDK + +The Hyperforge Harness SDK provides an asynchronous agent loop for applications +that need streaming model responses, validated tools, conversation persistence, +live events, usage limits, feedback, memory, and sub-agent delegation. + +Import the framework from `hyperforge.harness_sdk`: + +```python +from hyperforge.harness_sdk import AgentHarness +``` + +## Agent Loop + +An `AgentHarness` run: + +1. Loads or creates its conversation. +2. Adds the system prompt and user prompt to message history. +3. Streams a response from the configured `ModelClient`. +4. Executes any requested tools and returns their validated results to the model. +5. Repeats until the model returns a final text response. +6. Persists messages and lifecycle events through the configured storage. + +Calls to `run()` on one harness are serialized. A harness instance can therefore +be reused for multiple turns in the same conversation. + +## Creating an Agent + +Hyperforge includes `NucliaModelClient`, which connects the harness to Nuclia's +chat-completions endpoint through Hyperforge's existing asynchronous NUA client: + +```python +from hyperforge.harness_sdk import AgentHarness, HarnessEventType, NucliaModelClient + + +async def build_agent() -> AgentHarness: + model_client = await NucliaModelClient.from_api_key("your-api-key") + return AgentHarness( + model="your-model", + model_client=model_client, + reasoning_effort="high", + system_prompt="You are a concise support assistant.", + ) + + +async def answer(agent: AgentHarness, prompt: str) -> str: + result = "" + async for event in agent.run(prompt): + if event.type == HarnessEventType.TURN_COMPLETED: + result = event.payload["text"] + return result +``` + +Close a `NucliaModelClient` created by the application during application +shutdown: + +```python +await model_client.aclose() +``` + +Applications can provide another model integration by implementing +`ModelClient.stream()`. It receives the model name, complete message history, +registered tools, and execution context, and yields `ModelDelta` objects: + +```python +from collections.abc import AsyncIterator + +from hyperforge.harness_sdk import ModelDelta + + +class ApplicationModelClient: + async def stream( + self, *, model, reasoning_effort, messages, tools, execution_context + ) -> AsyncIterator[ModelDelta]: + # Translate messages and tools to the provider's API here. + yield ModelDelta(text="Hello", input_tokens=10, output_tokens=1) +``` + +## Using Legacy Agents + +The legacy engine can use the harness as its root orchestrator. Configure an +`agents` list instead of the retrieval pipeline stages; each configured legacy +agent contributes its published functions as namespaced harness tools: + +```python +config = { + "model": "your-model", + "reasoning_effort": "high", + "agents": [ + {"id": "docs", "module": "basic_ask", "sources": ["nuclia-docs"]}, + ], + "drivers": [ + { + "name": "Nuclia docs", + "provider": "nucliadb", + "identifier": "nuclia-docs", + "config": {"url": "...", "kbid": "...", "key": "..."}, + } + ], + "rules": {"rules": ["Be concise"]}, + "memory": {}, + "workflow": { + "id": "default", + "name": "Default workflow", + "description": None, + "parameters": {}, + }, +} + +memory = await hyperforge.engine.main( + config=config, + question="How do I use the API?", + loaded_modules=["hyperforge_nucliadb"], + external_nua_api_key="...", +) +``` + +For example, an agent with ID `docs` that publishes `ask_agent` is exposed to +the model as `docs__ask_agent`. + +Token values in streaming deltas are cumulative for one model request. The +harness records the largest reported values and aggregates them across the run. + +## Defining Tools + +A tool uses Pydantic models for its arguments and return value. Field +descriptions are included in the JSON Schema sent to the model. + +```python +from pydantic import BaseModel, Field + +from hyperforge.harness_sdk import AgentHarness, tool + + +class WeatherInput(BaseModel): + city: str = Field(description="City for the weather lookup") + + +class WeatherOutput(BaseModel): + summary: str + + +@tool(description="Get the current weather for a city.") +async def weather( + harness: AgentHarness, + input_value: WeatherInput, +) -> WeatherOutput: + return WeatherOutput(summary=f"Sunny in {input_value.city}") + + +agent = AgentHarness( + model="your-model", + model_client=model_client, + tools=[weather], +) +``` + +The `@tool` decorator creates a `HarnessTool` and infers its input and output +models from the handler annotations. The decorated tool remains callable, so +application code can use `await weather(harness, WeatherInput(city="Boston"))` +as it would call the original function. Direct `HarnessTool(...)` construction +is available when a decorator is not appropriate. +Nested JSON Schema references are flattened before schemas are sent to the model; +the emitted schema does not contain `$ref`, `$defs`, or `definitions`. +The harness validates model-provided arguments before calling the handler and +validates the result against its annotated return model. Missing or non-Pydantic +input and return annotations fail when the tool is constructed. Tool failures are +returned to the model as failed tool results so it can recover or answer differently. +Multiple tool calls from one model response run concurrently; their result +messages retain the model's original call order. + +Tools marked `lazy_load=True` are registered for execution but omitted from the +model's initial tool list. The always-available `search_tools` and +`activate_tools` tools let the model discover matching lazy tools and expose +their schemas on subsequent model calls. Lazy activation remains active for the +life of the harness instance. + +```python +@tool(description="Get the current weather for a city.", lazy_load=True) +async def weather(harness: AgentHarness, input_value: WeatherInput) -> WeatherOutput: + return WeatherOutput(summary=f"Sunny in {input_value.city}") +``` + +`lazy_load` can also be passed directly to `HarnessTool`. Legacy published agent +functions support the same option on `FunctionDefinition`. + +For a tool without arguments, define and use an empty Pydantic input model. Tool +handlers always receive a validated model instance; they never receive `None`. + +Tools can attach typed context to their result by setting `context_type` and, if +needed, registering a schema and formatter with `register_context()`. + +## Usage Limits + +All limits are disabled by default. Configure only the limits needed by the +application: + +```python +from hyperforge.harness_sdk import AgentHarness, UsageLimits + +agent = AgentHarness( + model="your-model", + model_client=model_client, + usage_limits=UsageLimits( + max_tool_calls=20, + max_turns=10, + max_input_tokens=100_000, + max_output_tokens=20_000, + max_time=120, + ), +) +``` + +`max_turns` counts model requests in one `run()` call. `max_time` is in seconds +and covers model and tool work. A parent agent and its sub-agents share usage and +the same limits. Exceeding a limit raises `UsageLimitExceeded` and records a +failed turn event. + +After a run, inspect `agent.usage` for `turns`, `tool_calls`, `input_tokens`, and +`output_tokens`. + +## Conversations and Storage + +Pass a `conversation_id` to continue a known conversation. Agent IDs are +internal implementation details and are never supplied by application code. + +```python +from hyperforge.harness_sdk import AgentHarness, InMemoryHarnessStorage + +storage = InMemoryHarnessStorage() + +first = AgentHarness( + model="your-model", + model_client=model_client, + conversation_id="support-thread-42", + storage=storage, +) +async for _ in first.run("My order is late"): + pass + +resumed = AgentHarness( + model="your-model", + model_client=model_client, + conversation_id="support-thread-42", + storage=storage, +) +await resumed.load(create=False) +async for _ in resumed.run("What should I do next?"): + pass +``` + +`InMemoryHarnessStorage` is process-local and intended for tests and ephemeral +agents. Production applications should implement `HarnessStorageProtocol` with +durable conversation, event, and memory storage. `create_agent()` is a convenience +function that constructs and loads an agent in one call. + +Conversation metadata can be supplied under `conversation_metadata` in the +execution context: + +```python +agent = AgentHarness( + model="your-model", + model_client=model_client, + execution_context={ + "user_id": "user-42", + "conversation_metadata": {"tenant": "acme"}, + }, +) +``` + +The complete `execution_context` is also passed to the model client. + +## Streaming Events + +Iterate `run()` to receive text and reasoning deltas, +model and tool lifecycle events, inbox activity, feedback activity, and the final +turn event: + +```python +import asyncio + +from hyperforge.harness_sdk import HarnessEventType + + +async def run_with_events(agent: AgentHarness, prompt: str) -> str: + result = "interrupted" + async with agent: + async for event in agent.run(prompt): + if event.type == HarnessEventType.TEXT_DELTA: + send_to_client(event.payload["text"]) + elif event.type == HarnessEventType.TURN_COMPLETED: + result = event.payload["text"] + return result +``` + +`TEXT_DELTA` and `REASONING_DELTA` events are live-only. Other events are +persisted. Use `history()` to iterate over persisted conversation events. + +`run()` owns one turn-scoped event queue and supports one consumer. Event +publication never waits for the consumer. If the consumer falls behind by +`event_queue_size` events, additional `TEXT_DELTA` and `REASONING_DELTA` events +are dropped; persisted lifecycle events remain lossless. Events from child agents +may interleave, but `turn_id`, `agent_id`, and `parent_agent_id` identify their +origin. Use `async with agent` when a caller may stop consuming early so the +active turn and descendants are cleaned up. + +At most four agents run concurrently in a conversation by default, including +the root agent. Configure `max_concurrent_agents` to change this limit. When the +limit is reached, `spawn_agent` returns `concurrency_limit_reached` and identifies +children the caller can pass to `wait_agent`; it does not start another agent. + +## Steering and Interrupting + +`steer()` adds a message while the run is active. The message is delivered before +the next model request: + +```python +await agent.steer("Prefer a concise answer") +``` + +`interrupt()` stops active model work, tool calls, and descendant agents: + +```python +async for event in agent.run("Perform a long task"): + agent.interrupt() + if event.type == HarnessEventType.TURN_INTERRUPTED: + break +``` + +External task cancellation remains regular `asyncio` cancellation and is +re-raised rather than converted to an interrupted result. + +## Feedback + +Enable the built-in `feedback` tool when the model should be able to ask the +application for input: + +```python +agent = AgentHarness( + model="your-model", + model_client=model_client, + feedback_enabled=True, +) + +async for event in agent.run("Prepare the report"): + if event.type == HarnessEventType.FEEDBACK_REQUESTED: + await agent.respond_feedback(event.payload["request_id"], "Use markdown") +``` + +Feedback requests time out after five minutes when invoked through the built-in +tool. `request_feedback()` is also available for application-defined tools that +need a custom schema or timeout. + +## Memory and Compaction + +The following core tools are available to the model by default: + +- `remember`, `recall`, and `forget` manage records through the configured storage. +- `compact` replaces current history with the system prompt and a model-authored summary. +- `spawn_agent`, `send_message`, and `wait_agent` coordinate sub-agents. + +Disable core tools that should not be available to an agent: + +```python +agent = AgentHarness( + model="your-model", + model_client=model_client, + disabled_core_tools={"remember", "forget", "spawn_agent"}, +) +``` + +Memory scope is an application-defined string. A production storage +implementation is responsible for enforcing tenant and user isolation. + +## Sub-Agents + +The model can call `spawn_agent` with a self-contained prompt, then call +`wait_agent` with the returned internal ID. It can use `send_message` to steer a +running child. Children inherit the model client, model, external tools, +execution context, storage, and usage limits. They start with isolated history +unless `include_history=true` is requested. Included history contains the current +system and user messages but never an incomplete spawn tool exchange. + +Delegation is limited to depth one by default. Set `max_spawn_depth` on the +root harness to change the limit; an agent at the limit does not advertise the +`spawn_agent` tool. Children are turn-scoped and are stopped when the parent turn +completes, fails, times out, or is interrupted. + +Internal child IDs are only handles for these core tools. They are not part of +the application-facing agent or storage API. + +## Adapting Existing Hyperforge Agents + +`AgentHarness.to_tools()` converts an existing Hyperforge agent's +`__published_functions__` into harness tools. Each published function becomes one +tool with its configured name, description, and complete parameter schema. Use a +namespace when combining agents that publish the same function name. + +```python +from hyperforge.harness_sdk import AgentHarness +from hyperforge_perplexity import PerplexityAgent + + +perplexity = PerplexityAgent( + config=perplexity_config, +) + +agent = AgentHarness( + model="your-model", + model_client=model_client, + tools=AgentHarness.to_tools(perplexity, namespace="perplexity") + [application_tool], + system_prompt=( + "Use `perplexity__internet_search` when the answer requires current information." + ), + execution_context={ + "drivers": { + perplexity_config.source: perplexity_driver, + } + }, +) + +answer = "" +async for event in agent.run("What changed in Python this week?"): + if event.type == HarnessEventType.TURN_COMPLETED: + answer = event.payload["text"] +``` + +The conversion uses each `FunctionDefinition` in `__published_functions__` and +the corresponding instance method. Model arguments are validated from the +published JSON schema and method annotations. Framework-only `memory` and +`manager` arguments are omitted from the model-facing tool schema. Schema +constraints such as enums, nested objects, defaults, bounds, and +`additionalProperties` are preserved and references are flattened. + +On every tool call, the compatibility layer: + +- Creates a fresh no-op `QuestionMemory` for the harness conversation. +- Creates a compatibility `Manager` that shares the NUA client when the harness + uses `NucliaModelClient`. +- Adds prebuilt drivers from `execution_context["drivers"]`. +- Calls the published method and returns its result as structured tool context. + +For a custom `ModelClient`, pass an existing `Manager` to `to_tools()` or provide +an asynchronous NUA client as `execution_context["nua_client"]`. The harness model +loop still uses chat completions; the manager exists only for unchanged legacy +published functions that call the established manager API. + +Agents that discover published functions asynchronously can use +`preload_published_agent_to_tools()` with an existing manager. + +Agents whose published functions only use `manager.nua` require no extra manager +configuration. Agents such as Perplexity that use `manager.drivers` need their +initialized drivers passed through `execution_context` as shown above. + +`to_tools()` raises `TypeError` when the agent has no published functions or a +published function does not map to a callable method. `HarnessAgentToTool` +remains available for exceptional agents that need custom memory construction or +result adaptation. + +## Testing Applications + +Use a small `ModelClient` test double that yields deterministic `ModelDelta` +objects. This exercises the real message, tool, event, and persistence loop +without making a model request: + +```python +from collections.abc import AsyncIterator + +from hyperforge.harness_sdk import AgentHarness, ModelDelta + + +class TestModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + yield ModelDelta(text="test answer", input_tokens=2, output_tokens=2) + + +async def test_agent_answer() -> None: + agent = AgentHarness(model="test", model_client=TestModel()) + events = [event async for event in agent.run("question")] + assert events[-1].payload["text"] == "test answer" +``` + +For tool-loop tests, return a `HarnessToolCall` from the first model request and a +final text delta after the harness appends the tool result. See +`hyperforge/tests/harness_sdk/` for complete happy-path examples. diff --git a/hyperforge/pyproject.toml b/hyperforge/pyproject.toml index d5ab7427..809bf712 100644 --- a/hyperforge/pyproject.toml +++ b/hyperforge/pyproject.toml @@ -31,11 +31,13 @@ dependencies = [ "httpx>=0.28.1,<0.29", "httpcore>=1.0.9,<1.1", "jinja2", + "jsonschema", + "RestrictedPython", "databases[asyncpg]", "sqlalchemy", "cryptography", "mmh3", - "psycopg2", + "psycopg2-binary", "tenacity", "websockets", "redis[hiredis]>=6.4.0", diff --git a/hyperforge/src/hyperforge/codemode/__init__.py b/hyperforge/src/hyperforge/codemode/__init__.py new file mode 100644 index 00000000..c2a775ee --- /dev/null +++ b/hyperforge/src/hyperforge/codemode/__init__.py @@ -0,0 +1,39 @@ +from .model import ( + RestrictedPythonTask, + SandboxMessage, + WorkerError, + WorkerExecutionRequest, + WorkerTypes, + deserialize, + serialize, +) +from .sandbox import ( + InsecureSandbox, + SandboxReader, + SandboxRunner, + SandboxSession, + SandboxSettings, + SandboxWriter, + run_sandbox_server, + settings, +) +from .worker import PythonAgentWorker + +__all__ = [ + "InsecureSandbox", + "PythonAgentWorker", + "RestrictedPythonTask", + "SandboxMessage", + "SandboxReader", + "SandboxRunner", + "SandboxSession", + "SandboxSettings", + "SandboxWriter", + "WorkerExecutionRequest", + "WorkerError", + "WorkerTypes", + "deserialize", + "run_sandbox_server", + "serialize", + "settings", +] diff --git a/agents/restricted/src/hyperforge_restricted/model.py b/hyperforge/src/hyperforge/codemode/model.py similarity index 67% rename from agents/restricted/src/hyperforge_restricted/model.py rename to hyperforge/src/hyperforge/codemode/model.py index 5573dc5c..4e2f34f6 100644 --- a/agents/restricted/src/hyperforge_restricted/model.py +++ b/hyperforge/src/hyperforge/codemode/model.py @@ -1,18 +1,21 @@ from dataclasses import dataclass from typing import Any, Sequence, assert_never, cast -from hyperforge.definition import FunctionDefinition -from hyperforge.memory import Context from nuclia_models.predict.remi import RemiResponse from pydantic import BaseModel +from hyperforge.definition import FunctionDefinition +from hyperforge.memory import Context + class WorkerExecutionRequest(BaseModel): code: str - question: str + question: str = "" local_vars: dict[str, Any] global_vars: dict[str, Any] function_names: dict[str, dict[str, FunctionDefinition]] + max_runtime_seconds: float | None = None + max_memory_bytes: int | None = None class RestrictedPythonTask(BaseModel): @@ -22,21 +25,27 @@ class RestrictedPythonTask(BaseModel): keyword_args: dict[str, Any] -# The types that are visible to the worker in parameters and return types +class WorkerError(BaseModel): + error: str + + WorkerTypes = ( str | bool | int + | float | Context | Sequence["WorkerTypes"] | dict[str, "WorkerTypes"] | None | RemiResponse | RestrictedPythonTask + | WorkerError ) WorkerModels: dict[str, type[BaseModel]] = { "Context": Context, "RemiResponse": RemiResponse, + "WorkerError": WorkerError, } @@ -44,8 +53,6 @@ def deserialize(msg: Any) -> WorkerTypes: if isinstance(msg, dict) and "__model__" in msg: model_name = msg["__model__"] if model_name == "RestrictedPythonTask": - # Special case for RestrictedPythonTask because args are type Any and cannot be serialized - # to the original with the default pydantic serialization/deserialization return RestrictedPythonTask( function=msg["function"], agent=msg["agent"], @@ -54,62 +61,54 @@ def deserialize(msg: Any) -> WorkerTypes: k: deserialize(v) for k, v in msg["keyword_args"].items() }, ) - else: - return cast(WorkerTypes, WorkerModels[model_name].model_validate(msg)) - elif isinstance(msg, list): + model = WorkerModels.get(model_name) + if model is not None: + return cast(WorkerTypes, model.model_validate(msg)) + if isinstance(msg, list): return [deserialize(m) for m in msg] - else: - return msg + if isinstance(msg, dict): + return {k: deserialize(v) for k, v in msg.items()} + return msg -def serialize(message: WorkerTypes) -> Any: - converted: Any +def serialize(message: Any) -> Any: if isinstance(message, RestrictedPythonTask): - # Special case for RestrictedPythonTask because args are type Any and cannot be serialized - # to the original with the default pydantic serialization/deserialization - converted = { + return { "__model__": "RestrictedPythonTask", "function": message.function, "agent": message.agent, "args": [serialize(a) for a in message.args], "keyword_args": {k: serialize(v) for k, v in message.keyword_args.items()}, } - elif isinstance(message, BaseModel): - model = message.__class__.__name__ + if isinstance(message, BaseModel): converted = message.model_dump() - converted["__model__"] = model - elif isinstance(message, list): - converted = [serialize(m) for m in message] - elif isinstance(message, dict): - converted = {k: deserialize(v) for k, v in message.items()} - else: - converted = message - - return converted + converted["__model__"] = message.__class__.__name__ + return converted + if isinstance(message, (list, tuple)): + return [serialize(m) for m in message] + if isinstance(message, dict): + return {k: serialize(v) for k, v in message.items()} + return message class SandboxMessage: - # agent -> worker to initiate execution @dataclass class Run: run: WorkerExecutionRequest + token: str | None = None - # worker -> agent to run a function @dataclass class Request: task: RestrictedPythonTask - # agent -> worker with function results @dataclass class Response: result: WorkerTypes - # worker -> agent when it has finished @dataclass class Done: pass - # worker -> agent in case of error @dataclass class Error: error: str @@ -120,14 +119,16 @@ class Error: def parse(cls, data: dict[str, Any]) -> AnyMessage: match data["_"]: case "run": - return cls.Run(run=WorkerExecutionRequest.model_validate(data["run"])) + return cls.Run( + run=WorkerExecutionRequest.model_validate(data["run"]), + token=data.get("token"), + ) case "request": task = deserialize(data["task"]) assert isinstance(task, RestrictedPythonTask) return cls.Request(task=task) case "response": - result = deserialize(data["result"]) - return cls.Response(result=result) + return cls.Response(result=deserialize(data["result"])) case "done": return cls.Done() case "error": @@ -139,7 +140,11 @@ def parse(cls, data: dict[str, Any]) -> AnyMessage: def serialize(cls, message: AnyMessage) -> dict[str, Any]: match message: case cls.Run(): - return {"_": "run", "run": message.run.model_dump()} + return { + "_": "run", + "run": message.run.model_dump(), + "token": message.token, + } case cls.Request(): return {"_": "request", "task": serialize(message.task)} case cls.Response(): diff --git a/hyperforge/src/hyperforge/codemode/sandbox.py b/hyperforge/src/hyperforge/codemode/sandbox.py new file mode 100644 index 00000000..6129614b --- /dev/null +++ b/hyperforge/src/hyperforge/codemode/sandbox.py @@ -0,0 +1,426 @@ +import asyncio +import contextvars +import hmac +import json +import logging +import os +from concurrent.futures import Executor, ThreadPoolExecutor +from concurrent.futures import Future as ConcurrentFuture +from multiprocessing import Pipe, Process +from multiprocessing.connection import Connection +from typing import Any, Callable, Coroutine + +import httpx +import nucliadb_telemetry.metrics +import prometheus_client +from pydantic_settings import BaseSettings + +from .model import ( + RestrictedPythonTask, + SandboxMessage, + WorkerError, + WorkerExecutionRequest, + WorkerTypes, +) +from .worker import PythonAgentWorker + +logger = logging.getLogger("hyperforge_codemode_sandbox") +WORKER_CPU_LIMIT = 1 +WORKER_EXIT_TIMEOUT = 1 +MAX_PACKET_BYTES = 16 * 1024 * 1024 + + +class SandboxSettings(BaseSettings): + sandbox_verify: bool = True + sandbox_socket: str | None = None + sandbox_metrics_port: int = 8091 + sandbox_token: str | None = None + + +settings = SandboxSettings() + + +class InsecureSandbox(Exception): + pass + + +class SandboxRunner: + @classmethod + def with_pool( + cls, + pool: Executor, + callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], + debug: bool = False, + ) -> "SandboxRunner": + return cls(pool=pool, socket=None, callback=callback, debug=debug) + + @classmethod + def remote( + cls, + socket: str, + callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], + debug: bool = False, + ) -> "SandboxRunner": + return cls(pool=None, socket=socket, callback=callback, debug=debug) + + @classmethod + def isolated_process( + cls, + callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], + ) -> "SandboxRunner": + runner = cls(pool=None, socket=None, callback=callback, debug=False) + runner.isolated = True + return runner + + def __init__( + self, + pool: Executor | None, + socket: str | None, + callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], + debug: bool, + ): + self.pool = pool + self.socket = socket + self.callback = callback + self.debug = debug + self.isolated = False + + async def run(self, request: WorkerExecutionRequest): + if self.pool is not None: + return await self._run_in_pool(request) + if self.socket is not None: + return await self._run_remotely(request) + if self.isolated: + return await self._run_isolated(request) + raise ValueError("SandboxRunner must be initialized with a pool or a socket") + + async def _run_in_pool(self, request: WorkerExecutionRequest): + loop = asyncio.get_running_loop() + pipe_worker, pipe_restricted = Pipe() + worker = PythonAgentWorker(pipe_worker, debug=self.debug) + controller_task = asyncio.create_task(self.background_task(pipe_restricted)) + await loop.run_in_executor( + self.pool, + worker._process_question_context_sync, + request.code, + request.question, + request.local_vars, + request.global_vars, + request.function_names, + request.max_memory_bytes, + ) + await controller_task + + async def _run_isolated(self, request: WorkerExecutionRequest): + controller_task, process = self.run_in_process(request) + try: + if request.max_runtime_seconds is None: + clean_exit = await controller_task + else: + async with asyncio.timeout(request.max_runtime_seconds): + clean_exit = await controller_task + if not clean_exit: + raise RuntimeError("Codemode worker exited unexpectedly") + await asyncio.to_thread(process.join, WORKER_EXIT_TIMEOUT) + if process.exitcode != 0: + raise RuntimeError( + "Codemode worker exceeded its memory limit or exited unexpectedly" + ) + except TimeoutError as exc: + raise RuntimeError("Codemode execution timed out") from exc + finally: + if process.is_alive(): + process.kill() + await asyncio.to_thread(process.join) + if not controller_task.done(): + controller_task.cancel() + await asyncio.gather(controller_task, return_exceptions=True) + + async def _run_remotely(self, request: WorkerExecutionRequest): + if settings.sandbox_token is None: + raise RuntimeError("SANDBOX_TOKEN is required for remote codemode") + rx, tx = await asyncio.open_unix_connection(self.socket) + reader, writer = SandboxReader(rx), SandboxWriter(tx) + try: + await writer.write_message( + SandboxMessage.Run(run=request, token=settings.sandbox_token) + ) + while True: + try: + msg = await reader.read_message() + except asyncio.IncompleteReadError as exc: + raise RuntimeError( + "Sandbox connection closed unexpectedly" + ) from exc + if isinstance(msg, SandboxMessage.Done): + break + if isinstance(msg, SandboxMessage.Error): + raise RuntimeError(f"Python agent error: {msg.error}") + assert isinstance(msg, SandboxMessage.Request) + try: + response = await self.callback(msg.task) + except Exception as exc: + response = WorkerError(error=str(exc)) + await writer.write_message(SandboxMessage.Response(result=response)) + finally: + writer.close() + await writer.wait_closed() + + def run_in_process(self, request: WorkerExecutionRequest): + pipe_worker, pipe_restricted = Pipe() + worker = PythonAgentWorker(pipe_worker, debug=self.debug) + controller_task = asyncio.create_task(self.background_task(pipe_restricted)) + process = Process( + target=worker._process_question_context_sync, + args=( + request.code, + request.question, + request.local_vars, + request.global_vars, + request.function_names, + request.max_memory_bytes, + ), + ) + process.start() + pipe_worker.close() + return controller_task, process + + async def background_task(self, pipe: Connection) -> bool: + loop = asyncio.get_running_loop() + context = contextvars.copy_context() + try: + return await loop.run_in_executor( + ThreadPoolExecutor(max_workers=1), + self.read_write_queue, + pipe, + loop, + context, + ) + finally: + pipe.close() + + def read_write_queue( + self, + pipe: Connection, + loop: asyncio.AbstractEventLoop, + context: contextvars.Context, + ) -> bool: + try: + while True: + item = pipe.recv() + if item.function == "_error": + raise RuntimeError(item.args[0]) + if item.function == "_close": + return True + future: ConcurrentFuture[WorkerTypes] = ConcurrentFuture() + + def schedule_callback() -> None: + task = loop.create_task(self.callback(item), context=context.copy()) + + def copy_result(completed: asyncio.Task[WorkerTypes]) -> None: + try: + future.set_result(completed.result()) + except Exception as exc: + future.set_exception(exc) + + task.add_done_callback(copy_result) + + loop.call_soon_threadsafe(schedule_callback) + try: + result = future.result() + except Exception as exc: + result = WorkerError(error=str(exc)) + pipe.send(result) + except EOFError: + return False + + +def _url_reachable(url: str) -> bool: + with httpx.Client() as client: + try: + client.get(url) + return True + except httpx.ConnectError: + return False + + +def _verify_connectivity(): + kubernetes_host = os.environ.get("KUBERNETES_SERVICE_HOST") + if kubernetes_host and _url_reachable(f"https://{kubernetes_host}"): + raise InsecureSandbox() + if _url_reachable("https://nuclia.com"): + raise InsecureSandbox() + + +class SandboxWriter: + def __init__(self, writer: asyncio.StreamWriter): + self.writer = writer + + async def _write_packet(self, msg: dict[str, Any]): + msg_bytes = json.dumps(msg).encode("utf-8") + if len(msg_bytes) > MAX_PACKET_BYTES: + raise ValueError("Sandbox message exceeds maximum size") + self.writer.write(len(msg_bytes).to_bytes(4, "little")) + self.writer.write(msg_bytes) + await self.writer.drain() + + async def write_message(self, message: SandboxMessage.AnyMessage): + await self._write_packet(SandboxMessage.serialize(message)) + + def close(self): + self.writer.close() + + async def wait_closed(self): + await self.writer.wait_closed() + + +class SandboxReader: + def __init__(self, reader: asyncio.StreamReader): + self.reader = reader + + async def _read_packet(self) -> dict[str, Any]: + len_bytes = await self.reader.readexactly(4) + msg_len = int.from_bytes(len_bytes, "little") + if msg_len > MAX_PACKET_BYTES: + raise ValueError("Sandbox message exceeds maximum size") + msg_bytes = await self.reader.readexactly(msg_len) + return json.loads(msg_bytes.decode("utf-8")) + + async def read_message(self) -> SandboxMessage.AnyMessage: + return SandboxMessage.parse(await self._read_packet()) + + +class SandboxSession: + def __init__(self, reader: SandboxReader, writer: SandboxWriter): + self.reader = reader + self.writer = writer + self.runner = SandboxRunner( + pool=None, socket=None, callback=self._callback, debug=False + ) + self.waiting_for_agent: asyncio.Task[None] | None = None + self.timed_out = False + + async def run(self, request: WorkerExecutionRequest): + self.task, self.process = self.runner.run_in_process(request) + try: + await self._start_timeout(WORKER_CPU_LIMIT) + if request.max_runtime_seconds is None: + clean_exit = await self.task + else: + async with asyncio.timeout(request.max_runtime_seconds): + clean_exit = await self.task + await self._stop_timeout() + if self.timed_out: + raise RuntimeError("Python agent timeout") + if not clean_exit: + raise RuntimeError("Codemode worker exited unexpectedly") + await asyncio.to_thread(self.process.join, WORKER_EXIT_TIMEOUT) + if self.process.exitcode != 0: + raise RuntimeError( + "Codemode worker exceeded its memory limit or exited unexpectedly" + ) + await self.writer.write_message(SandboxMessage.Done()) + except TimeoutError: + await self._fail(RuntimeError("Codemode execution timed out")) + except Exception as exc: + await self._fail(exc) + finally: + await self._stop_timeout() + if self.process.is_alive(): + self.process.kill() + await asyncio.to_thread(self.process.join) + if not self.task.done(): + self.task.cancel() + await asyncio.gather(self.task, return_exceptions=True) + + async def _fail(self, exc: Exception): + logger.error("Error in sandbox session: %s", exc) + if self.process.is_alive(): + self.process.kill() + await asyncio.to_thread(self.process.join) + if not self.task.done(): + self.task.cancel() + await asyncio.gather(self.task, return_exceptions=True) + try: + await self.writer.write_message(SandboxMessage.Error(error=str(exc))) + except Exception: + pass + + async def _callback(self, item: RestrictedPythonTask) -> WorkerTypes: + await self._stop_timeout() + await self.writer.write_message(SandboxMessage.Request(task=item)) + try: + response = await self.reader.read_message() + except asyncio.IncompleteReadError: + logger.warning("Got disconnected from agent") + self.process.kill() + self.process.join() + self.task.cancel() + return None + assert isinstance(response, SandboxMessage.Response) + await self._start_timeout(WORKER_CPU_LIMIT) + return response.result + + async def _start_timeout(self, timeout: float): + self.waiting_for_agent = asyncio.create_task(self._timeout(timeout)) + + async def _stop_timeout(self): + if self.waiting_for_agent: + self.waiting_for_agent.cancel() + self.waiting_for_agent = None + + async def _timeout(self, timeout: float): + await asyncio.sleep(timeout) + self.timed_out = True + if self.process.is_alive(): + self.process.kill() + + +agents_running = prometheus_client.Gauge( + "arag_sandbox_agents_running", "Number of Python agents currently running" +) +execution_observer = nucliadb_telemetry.metrics.Observer("arag_sandbox_execution") + + +async def run_sandbox_server(): + assert settings.sandbox_socket is not None + sandbox_token = settings.sandbox_token + if sandbox_token is None: + raise RuntimeError("SANDBOX_TOKEN is required for the sandbox server") + if settings.sandbox_verify: + _verify_connectivity() + + async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): + reader = SandboxReader(rx) + writer = SandboxWriter(tx) + try: + async with asyncio.timeout(1): + msg = await reader.read_message() + assert isinstance(msg, SandboxMessage.Run) + if not hmac.compare_digest(msg.token or "", sandbox_token): + raise PermissionError("Invalid sandbox token") + except (AssertionError, PermissionError, TimeoutError, ValueError): + logger.warning("Rejected invalid sandbox connection") + writer.close() + await writer.wait_closed() + return + + agents_running.inc() + observation = execution_observer() + observation.start() + try: + session = SandboxSession(reader, writer) + await session.run(msg.run) + except Exception as exc: + observation.set_status("error") + await writer.write_message(SandboxMessage.Error(error=str(exc))) + finally: + agents_running.dec() + observation.end() + writer.close() + await writer.wait_closed() + + server = await asyncio.start_unix_server(handler, settings.sandbox_socket) + os.chmod(settings.sandbox_socket, 0o600) + async with server: + await server.serve_forever() diff --git a/hyperforge/src/hyperforge/codemode/worker.py b/hyperforge/src/hyperforge/codemode/worker.py new file mode 100644 index 00000000..89b134bc --- /dev/null +++ b/hyperforge/src/hyperforge/codemode/worker.py @@ -0,0 +1,241 @@ +import os +import sys +from dataclasses import dataclass +from functools import partial +from multiprocessing.connection import Connection +from operator import ( + add, + and_, + floordiv, + lshift, + matmul, + mod, + mul, + or_, + pow, + rshift, + sub, + truediv, + xor, +) +from typing import Any, Dict, List, Optional + +from RestrictedPython import compile_restricted +from RestrictedPython.Guards import ( + guarded_iter_unpack_sequence, + guarded_unpack_sequence, + safe_builtins, +) + +from hyperforge.definition import FunctionDefinition +from hyperforge.memory import Chunk +from hyperforge.memory.memory import Context + +from .model import RestrictedPythonTask, WorkerError + +BLOCKED_EXCEPTIONS = { + "BaseException", + "GeneratorExit", + "KeyboardInterrupt", + "SystemExit", +} + + +def _set_memory_limit(max_memory_bytes: int | None) -> None: + if max_memory_bytes is None or sys.platform == "darwin": + return + try: + import resource + + with open("/proc/self/statm", encoding="ascii") as statm: + current_virtual_bytes = int(statm.read().split()[0]) * os.sysconf( + "SC_PAGE_SIZE" + ) + address_space_limit = current_virtual_bytes + max_memory_bytes + resource.setrlimit( + resource.RLIMIT_AS, + (address_space_limit, address_space_limit), + ) + except (ImportError, OSError, ValueError, IndexError) as exc: + raise RuntimeError("Codemode memory limits are unavailable") from exc + + +def _harden_process(max_memory_bytes: int | None) -> None: + os.environ.clear() + os.umask(0o077) + try: + import resource + + resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) + except (ImportError, OSError, ValueError): + pass + + if sys.platform.startswith("linux"): + import ctypes + + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(38, 1, 0, 0, 0) != 0: # PR_SET_NO_NEW_PRIVS + raise RuntimeError("Unable to disable worker privilege escalation") + libc.prctl(4, 0, 0, 0, 0) # PR_SET_DUMPABLE + + _set_memory_limit(max_memory_bytes) + + +def guarded_inplace(operator: str, left: Any, right: Any) -> Any: + operations = { + "+=": add, + "-=": sub, + "*=": mul, + "/=": truediv, + "//=": floordiv, + "%=": mod, + "**=": pow, + "<<=": lshift, + ">>=": rshift, + "&=": and_, + "^=": xor, + "|=": or_, + "@=": matmul, + } + operation = operations.get(operator) + if operation is None: + raise ValueError(f"Unsupported inplace operator: {operator}") + return operation(left, right) + + +class PythonAgentWorker: + context: Optional[Context] = None + + def __init__(self, pipe: Connection, debug: bool = False): + self.pipe = pipe + self.functions_agent_id: Dict[str, List[str]] = {} + self.debug = debug + + def _process_question_context_sync( + self, + code: str, + question: str, + local_vars: Dict[str, Any], + global_vars: Dict[str, Any], + function_names: Dict[str, Dict[str, FunctionDefinition]], + max_memory_bytes: int | None = None, + ): + try: + _harden_process(max_memory_bytes) + for agent, functions in function_names.items(): + for function in functions: + if function in global_vars: + self.functions_agent_id[function].append(agent) + else: + global_vars[function] = partial( + self.valid_functions_call, function + ) + self.functions_agent_id[function] = [agent] + + byte_code = compile_restricted(code, filename="", mode="exec") + global_vars.update( + { + "__builtins__": { + **{ + name: value + for name, value in safe_builtins.items() + if name not in BLOCKED_EXCEPTIONS + }, + "sum": sum, + }, + "_getitem_": lambda obj, index: obj[index], + "_getiter_": iter, + "_inplacevar_": guarded_inplace, + "dataclass": dataclass, + "Chunk": Chunk, + "Context": Context, + "List": List, + "Any": Any, + "Dict": Dict, + "_iter_unpack_sequence_": guarded_iter_unpack_sequence, + "_unpack_sequence_": guarded_unpack_sequence, + "save": self.save, + "question": question, + } + ) + if self.debug: + global_vars["pdb"] = __import__("pdb") + exec(byte_code, global_vars, local_vars) + except BaseException as exc: + self.pipe.send( + RestrictedPythonTask( + function="_error", + agent="_controller", + args=(str(exc),), + keyword_args={}, + ) + ) + finally: + self.pipe.send( + RestrictedPythonTask( + function="_close", + agent="_controller", + args=(None,), + keyword_args={}, + ) + ) + + def save( + self, + contexts: List[Context] | None = None, + resolved: bool = True, + missing: list[str] | None = None, + question: str = "", + final_answer: Optional[str] = None, + ): + self.pipe.send( + RestrictedPythonTask( + function="save", + agent="_controller", + args=(None,), + keyword_args={ + "contexts": contexts or [], + "resolved": resolved, + "missing": missing or [], + "question": question, + "final_answer": final_answer, + }, + ) + ) + return self._receive() + + def valid_functions_call( + self, function_name: str, *args: Any, **kwargs: Any + ) -> Any: + agent_id = kwargs.pop("agent_id", None) + if agent_id is not None: + if agent_id not in self.functions_agent_id.get(function_name, []): + raise ValueError( + f"Agent ID ({agent_id}) not authorized for this function {function_name}" + ) + else: + agents = self.functions_agent_id.get(function_name, []) + if not agents: + raise ValueError(f"No agent found for function {function_name}") + if len(agents) > 1: + raise ValueError( + f"Multiple agents found for function {function_name}: {agents}, " + "please specify agent_id" + ) + agent_id = agents[0] + + self.pipe.send( + RestrictedPythonTask( + function=function_name, + agent=agent_id, + args=args, + keyword_args=kwargs, + ) + ) + return self._receive() + + def _receive(self) -> Any: + result = self.pipe.recv() + if isinstance(result, WorkerError): + raise RuntimeError(result.error) + return result diff --git a/hyperforge/src/hyperforge/definition.py b/hyperforge/src/hyperforge/definition.py index 6b62dcc9..b815afc2 100644 --- a/hyperforge/src/hyperforge/definition.py +++ b/hyperforge/src/hyperforge/definition.py @@ -11,3 +11,6 @@ class FunctionDefinition(BaseModel): name: str description: str parameters: Dict[str, Any] + method: str | None = None + input_schema: Dict[str, Any] | None = None + lazy_load: bool = False diff --git a/hyperforge/src/hyperforge/engine.py b/hyperforge/src/hyperforge/engine.py index 136d9840..2de7a33a 100644 --- a/hyperforge/src/hyperforge/engine.py +++ b/hyperforge/src/hyperforge/engine.py @@ -5,6 +5,7 @@ from nuclia.lib.nua import AsyncNuaClient from hyperforge.configure import GLOBAL_REGISTRY, load_all_configurations +from hyperforge.harness import HarnessAgent, HarnessAgentConfig from hyperforge.interaction import AragAnswer from hyperforge.llm import NoopNuaClient, NuaBaseModel, NUAConnection from hyperforge.manager import Manager @@ -18,7 +19,7 @@ @dataclass class State: - agent: RetrievalAgent + agent: RetrievalAgent | HarnessAgent manager: Manager @@ -31,7 +32,7 @@ async def init( local_openai_model: Optional[str] = None, external_nua_api_key: Optional[str] = None, loaded_modules: list[str] = [], - retrieval_config: Optional[RetrievalAgentConfig] = None, + retrieval_config: Optional[RetrievalAgentConfig | HarnessAgentConfig] = None, session_id: str = "default_session", memory_klass: type[BaseSessionMemory] = SessionMemory, ) -> Tuple[State, SessionMemory]: @@ -47,7 +48,10 @@ async def init( if retrieval_config is None: if config is None: raise ValueError("Either config or retrieval_config must be provided") - retrieval_config = RetrievalAgentConfig.model_validate(config) + config_class = ( + HarnessAgentConfig if "agents" in config else RetrievalAgentConfig + ) + retrieval_config = config_class.model_validate(config) state = await get_state( agent_id=agent_id, @@ -78,7 +82,7 @@ async def main( external_nua_api_key: Optional[str] = None, question: str = "", loaded_modules: list[str] = [], - retrieval_config: Optional[RetrievalAgentConfig] = None, + retrieval_config: Optional[RetrievalAgentConfig | HarnessAgentConfig] = None, callback: Optional[Callable[[AragAnswer], Awaitable[None]]] = None, session_id: str = "default_session", user_metadata: Optional[Dict[str, str]] = None, @@ -132,7 +136,7 @@ async def main( async def engine( manager: Manager, - agent: RetrievalAgent, + agent: RetrievalAgent | HarnessAgent, question_memory: QuestionMemory, user_metadata: Optional[Dict[str, str]] = None, ) -> None: @@ -147,7 +151,7 @@ async def engine( async def get_state( agent_id: str, - config: RetrievalAgentConfig, + config: RetrievalAgentConfig | HarnessAgentConfig, internal_nua_api: str = "http://predict.learning.svc.cluster.local:8080", internal_nua: bool = False, local_openai: Optional[str] = None, @@ -191,6 +195,9 @@ async def get_state( nua = NoopNuaClient() manager = await Manager.from_config(drivers=config.drivers, nua=nua) - agent = await RetrievalAgent.from_config_class(config) + if isinstance(config, HarnessAgentConfig): + agent = await HarnessAgent.from_config_class(config) + else: + agent = await RetrievalAgent.from_config_class(config) return State(manager=manager, agent=agent) diff --git a/hyperforge/src/hyperforge/harness/__init__.py b/hyperforge/src/hyperforge/harness/__init__.py new file mode 100644 index 00000000..93b78eb7 --- /dev/null +++ b/hyperforge/src/hyperforge/harness/__init__.py @@ -0,0 +1,4 @@ +from .agent import HarnessAgent +from .config import HarnessAgentConfig + +__all__ = ["HarnessAgent", "HarnessAgentConfig"] diff --git a/hyperforge/src/hyperforge/harness/agent.py b/hyperforge/src/hyperforge/harness/agent.py new file mode 100644 index 00000000..fe667513 --- /dev/null +++ b/hyperforge/src/hyperforge/harness/agent.py @@ -0,0 +1,117 @@ +from typing import Any, cast + +from hyperforge.agent import Agent +from hyperforge.configure import get_agent_klass +from hyperforge.harness.config import HarnessAgentConfig +from hyperforge.harness_sdk import ( + AgentHarness, + HarnessEventType, + HarnessMessage, + NucliaChatCompletionsClient, + NucliaModelClient, +) +from hyperforge.interaction import Feedback, StreamingChunk +from hyperforge.llm import AsyncNuaClient +from hyperforge.manager import Manager +from hyperforge.memory.memory import QuestionMemory + + +class HarnessAgent: + def __init__(self, config: HarnessAgentConfig, agents: list[Agent[Any]]) -> None: + self.config = config + self.agents = agents + + @classmethod + async def from_config_class(cls, config: HarnessAgentConfig) -> "HarnessAgent": + agents = [ + await get_agent_klass(agent_config.module).from_config(agent_config) + for agent_config in config.agents + ] + return cls(config, agents) + + async def __call__(self, memory: QuestionMemory, manager: Manager) -> None: + for agent in self.agents: + preload = getattr(agent, "preload", None) + if callable(preload): + await preload(manager, memory) + tools = [ + tool + for agent in self.agents + for tool in AgentHarness.to_tools( + agent, namespace=agent.agent_id, manager=manager + ) + ] + rules = [ + rule if isinstance(rule, str) else rule.prompt + for rule in memory.get_rules() + if isinstance(rule, str) or rule.prompt + ] + system_prompt = self.config.system_prompt + if rules: + system_prompt += "\n\nRules:\n" + "\n".join(f"- {rule}" for rule in rules) + harness = AgentHarness( + model=self.config.model, + model_client=NucliaModelClient( + NucliaChatCompletionsClient(cast(AsyncNuaClient, manager.nua)) + ), + reasoning_effort=self.config.reasoning_effort, + tools=tools, + system_prompt=system_prompt, + title=self.config.workflow.name, + conversation_id=memory.get_session_id(), + disabled_core_tools=self.config.disabled_core_tools, + execution_context={ + "user_id": memory.session.user_info.get("id", "system"), + "conversation_metadata": { + "agent_id": memory.get_agent_id(), + "workflow_id": memory.get_workflow_id(), + "question_id": memory.original_question_uuid, + }, + "memory": memory, + }, + feedback_enabled=self.config.feedback_enabled, + usage_limits=self.config.usage_limits, + max_spawn_depth=self.config.max_spawn_depth, + max_concurrent_agents=self.config.max_concurrent_agents, + ) + history = await memory.get_chat_history() + harness.add_messages( + HarnessMessage( + role="user" if message.author.value == "USER" else "assistant", + content=message.text, + ) + for message in history + ) + result = "" + async for event in harness.run(memory.original_question): + if event.agent_id != harness.agent_id: + continue + if event.type == HarnessEventType.TEXT_DELTA: + await memory.emit_streaming_chunk( + StreamingChunk(text=event.payload["text"]) + ) + elif event.type == HarnessEventType.REASONING_DELTA: + await memory.emit_streaming_chunk( + StreamingChunk(text=event.payload["text"]), reasoning=True + ) + elif event.type == HarnessEventType.FEEDBACK_REQUESTED: + response = await memory.send_feedback( + Feedback( + request_id=event.payload["request_id"], + question=event.payload["question"], + module="harness", + agent_id=memory.get_agent_id(), + data={}, + timeout_ms=event.payload["timeout_ms"], + response_schema=event.payload["response_schema"], + ) + ) + if response is not None: + await harness.respond_feedback( + event.payload["request_id"], response.response + ) + elif event.type == HarnessEventType.TURN_COMPLETED: + result = str(event.payload.get("text", "")) + await memory.emit_streaming_chunk(StreamingChunk(text="", last=True)) + await memory.add_answer(result, "harness", "/harness") + await memory.add_final_answer() diff --git a/hyperforge/src/hyperforge/harness/config.py b/hyperforge/src/hyperforge/harness/config.py new file mode 100644 index 00000000..c7e25be1 --- /dev/null +++ b/hyperforge/src/hyperforge/harness/config.py @@ -0,0 +1,50 @@ +from typing import Any + +from pydantic import BaseModel, Field, field_serializer, field_validator + +from hyperforge.agent import AgentConfig +from hyperforge.configure import get_agent_config_klass, get_driver_config_klass +from hyperforge.driver import DriverConfig +from hyperforge.harness_sdk import SYSTEM_PROMPT, ReasoningEffort, UsageLimits +from hyperforge.models import MemoryConfig, Rules +from hyperforge.workflows import WorkflowData + + +class HarnessAgentConfig(BaseModel): + model: str + agents: list[AgentConfig] = Field(default_factory=list) + drivers: list[DriverConfig] = Field(default_factory=list) + system_prompt: str = SYSTEM_PROMPT + reasoning_effort: ReasoningEffort | None = None + disabled_core_tools: list[str] = Field(default_factory=list) + feedback_enabled: bool = False + usage_limits: UsageLimits = Field(default_factory=UsageLimits) + max_spawn_depth: int = 1 + max_concurrent_agents: int = 4 + rules: Rules = Field(default_factory=Rules) + memory: MemoryConfig = Field(default_factory=MemoryConfig) + workflow: WorkflowData + + @field_validator("agents", mode="before") + @classmethod + def validate_agents(cls, value: list[Any]) -> list[AgentConfig]: + return [ + item + if isinstance(item, AgentConfig) + else get_agent_config_klass(item["module"]).model_validate(item) + for item in value + ] + + @field_validator("drivers", mode="before") + @classmethod + def validate_drivers(cls, value: list[Any]) -> list[DriverConfig]: + return [ + item + if isinstance(item, DriverConfig) + else get_driver_config_klass(item["provider"]).model_validate(item) + for item in value + ] + + @field_serializer("agents", "drivers") + def serialize_config_items(self, items: list[BaseModel]) -> list[dict[str, Any]]: + return [item.model_dump() for item in items] diff --git a/hyperforge/src/hyperforge/harness_sdk/__init__.py b/hyperforge/src/hyperforge/harness_sdk/__init__.py new file mode 100644 index 00000000..ece0d691 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/__init__.py @@ -0,0 +1,92 @@ +from .agents import ( + HarnessAgentToTool, + preload_published_agent_to_tools, + published_agent_to_tools, +) +from .clients import ( + ChatCompletionChoice, + ChatCompletionChunk, + ChatCompletionDelta, + ChatCompletionRequest, + ChatCompletionResponseFormat, + ChatCompletionToolCallDelta, + ChatCompletionToolCallFunctionDelta, + ChatCompletionUsage, + ModelClient, + ModelDelta, + NucliaChatCompletionsClient, + NucliaChatCompletionsError, + NucliaModelClient, + ReasoningEffort, +) +from .context import context_schema, format_context, make_context, register_context +from .harness import ( + SYSTEM_PROMPT, + AgentHarness, + AgentResult, + LLMCallError, + create_agent, +) +from .models import ( + HarnessContextReference, + HarnessContextType, + HarnessConversation, + HarnessEvent, + HarnessEventType, + HarnessInboxItem, + HarnessMemory, + HarnessMessage, + HarnessToolCall, +) +from .storage import HarnessStorageProtocol, InMemoryHarnessStorage +from .tools import HarnessTool, tool +from .tools.codemode import CodemodeInput, CodemodeOutput, codemode +from .usage import HarnessUsage, UsageLimitExceeded, UsageLimits + +__all__ = [ + "SYSTEM_PROMPT", + "AgentHarness", + "AgentResult", + "ChatCompletionChunk", + "ChatCompletionChoice", + "ChatCompletionDelta", + "ChatCompletionRequest", + "ChatCompletionResponseFormat", + "ChatCompletionToolCallDelta", + "ChatCompletionToolCallFunctionDelta", + "ChatCompletionUsage", + "CodemodeInput", + "CodemodeOutput", + "HarnessAgentToTool", + "HarnessContextReference", + "HarnessContextType", + "HarnessConversation", + "HarnessEvent", + "HarnessEventType", + "HarnessInboxItem", + "HarnessMemory", + "HarnessMessage", + "HarnessStorageProtocol", + "HarnessTool", + "HarnessToolCall", + "HarnessUsage", + "InMemoryHarnessStorage", + "LLMCallError", + "ModelClient", + "ModelDelta", + "NucliaModelClient", + "NucliaChatCompletionsClient", + "NucliaChatCompletionsError", + "ReasoningEffort", + "UsageLimitExceeded", + "UsageLimits", + "context_schema", + "codemode", + "create_agent", + "format_context", + "make_context", + "register_context", + "tool", + "published_agent_to_tools", + "preload_published_agent_to_tools", +] diff --git a/hyperforge/src/hyperforge/harness_sdk/agents.py b/hyperforge/src/hyperforge/harness_sdk/agents.py new file mode 100644 index 00000000..6a551c05 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/agents.py @@ -0,0 +1,297 @@ +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from types import GenericAlias, UnionType +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Union, + cast, + get_args, + get_origin, + get_type_hints, +) + +from pydantic import BaseModel, Field, create_model + +from hyperforge.definition import FunctionDefinition +from hyperforge.manager import Manager +from hyperforge.memory.memory import BaseSessionMemory +from hyperforge.models import Context, MemoryConfig, Rules + +from .models import HarnessContextReference, HarnessContextType +from .tools import HarnessTool + +if TYPE_CHECKING: + from hyperforge.agent import Agent as HyperforgeAgent + from hyperforge.manager import Manager + from hyperforge.memory.memory import QuestionMemory + + from .harness import AgentHarness + + +class AgentToolInput(BaseModel): + prompt: str + + +class AgentToolOutput(BaseModel): + value: Any + + +class PublishedFunctionOutput(BaseModel): + value: Any + + +type MemoryFactory = Callable[ + ["AgentHarness", str], "Awaitable[QuestionMemory] | QuestionMemory" +] +type ResultAdapter = Callable[[Any, "QuestionMemory"], "Awaitable[Any] | Any"] + + +@dataclass(frozen=True) +class HarnessAgentToTool: + """Adapts an unchanged legacy Hyperforge agent to a harness tool.""" + + name: str + agent: HyperforgeAgent + manager: Manager + memory_factory: MemoryFactory + result_adapter: ResultAdapter + description: str = "" + + def as_tool(self) -> HarnessTool[AgentToolInput, AgentToolOutput]: + async def execute( + harness: AgentHarness, input_value: AgentToolInput + ) -> AgentToolOutput: + memory = self.memory_factory(harness, input_value.prompt) + if inspect.isawaitable(memory): + memory = await memory + memory = cast("QuestionMemory", memory) + result = await self.agent(memory=memory, manager=self.manager) + adapted = self.result_adapter(result, memory) + if inspect.isawaitable(adapted): + adapted = await adapted + return AgentToolOutput(value=adapted) + + return HarnessTool( + name=self.name, + description=self.description or f"Run the {self.name} Hyperforge agent", + handler=execute, + context_type=HarnessContextType.STRUCTURED, + ) + + +def published_agent_to_tools( + agent: Any, + *, + namespace: str | None = None, + manager: Manager | None = None, +) -> list[HarnessTool[Any, PublishedFunctionOutput]]: + published: dict[str, FunctionDefinition] = getattr( + agent, "__published_functions__", {} + ) + if not published: + raise TypeError( + f"{type(agent).__name__} does not define any published functions" + ) + + return [ + _published_function_tool( + agent, function_id, definition, namespace=namespace, manager=manager + ) + for function_id, definition in published.items() + ] + + +async def preload_published_agent_to_tools( + agent: Any, + *, + namespace: str | None = None, + manager: Manager, +) -> list[HarnessTool[Any, PublishedFunctionOutput]]: + preload = getattr(agent, "preload", None) + if callable(preload): + session = BaseSessionMemory.from_config( + MemoryConfig(), agent_id="harness", workflow_id="harness", rules=Rules() + ) + session.init("harness") + memory = session.start_question("preload", headers={}) + result = preload(manager, memory) + if inspect.isawaitable(result): + await result + return published_agent_to_tools(agent, namespace=namespace, manager=manager) + + +def _published_function_tool( + agent: Any, + function_id: str, + definition: FunctionDefinition, + *, + namespace: str | None, + manager: Manager | None, +) -> HarnessTool[Any, PublishedFunctionOutput]: + method_name = definition.method or function_id + method = getattr(agent, method_name, None) + if method is None or not callable(method): + raise TypeError( + f"Published function {function_id!r} is not callable on {type(agent).__name__}" + ) + + input_model = _published_input_model(method, function_id, definition) + + async def execute( + harness: AgentHarness, input_value: BaseModel + ) -> PublishedFunctionOutput: + compat_manager = manager or _compat_manager(harness) + memory = harness.execution_context.get("memory") + if memory is None: + session = BaseSessionMemory.from_config( + MemoryConfig(), + agent_id="harness", + workflow_id="harness", + rules=Rules(), + ) + session.init(harness.conversation_id) + memory = session.start_question( + str(input_value.model_dump().get("question", function_id)), + headers={"conversation-id": harness.conversation_id}, + ) + arguments = input_value.model_dump(exclude_unset=True) + signature = inspect.signature(method) + accepts_kwargs = any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + for parameter in signature.parameters.values() + ) + if "memory" in signature.parameters or accepts_kwargs: + arguments["memory"] = memory + if "manager" in signature.parameters or accepts_kwargs: + arguments["manager"] = compat_manager + result = method(**arguments) + if inspect.isawaitable(result): + result = await result + return PublishedFunctionOutput(value=result) + + execute.__annotations__["input_value"] = input_model + return HarnessTool( + name=f"{namespace}__{definition.name}" if namespace else definition.name, + description=definition.description, + handler=execute, + context_type=HarnessContextType.STRUCTURED, + parameters_schema=_published_schema(method, definition), + context_factory=_published_context, + lazy_load=definition.lazy_load, + ) + + +def _compat_manager(harness: AgentHarness) -> Manager: + manager = Manager() + nua = getattr(harness.model_client, "nua", None) + if nua is None: + nua = harness.execution_context.get("nua_client") + if nua is not None: + manager.nua = nua + drivers = harness.execution_context.get("drivers") + if drivers is not None: + if not isinstance(drivers, dict): + raise TypeError("execution_context['drivers'] must be a dictionary") + manager.drivers.update(drivers) + return manager + + +def _published_schema( + method: Callable[..., Any], definition: FunctionDefinition +) -> dict[str, Any]: + if definition.input_schema is not None: + return definition.input_schema + signature = inspect.signature(method) + properties = {name: dict(schema) for name, schema in definition.parameters.items()} + required = [] + for name, schema in properties.items(): + parameter = signature.parameters.get(name) + if parameter is not None and parameter.kind != inspect.Parameter.VAR_KEYWORD: + if parameter.default is inspect.Parameter.empty: + required.append(name) + else: + schema.setdefault("default", parameter.default) + elif "default" not in schema: + schema.setdefault("default", None) + return { + "type": "object", + "properties": properties, + "required": required, + "additionalProperties": False, + } + + +def _published_context(output: PublishedFunctionOutput) -> HarnessContextReference: + value = output.value + if isinstance(value, Context): + content = value.model_dump(mode="json") + context_type = HarnessContextType.RETRIEVAL + elif isinstance(value, list) and all(isinstance(item, Context) for item in value): + content = {"items": [item.model_dump(mode="json") for item in value]} + context_type = HarnessContextType.RETRIEVAL + elif isinstance(value, BaseModel): + content = {"value": value.model_dump(mode="json")} + context_type = HarnessContextType.STRUCTURED + else: + content = {"value": value} + context_type = HarnessContextType.TOOL_RESULT + return HarnessContextReference(type=context_type, content=content) + + +def _published_input_model( + method: Callable[..., Any], + function_id: str, + definition: FunctionDefinition, +) -> type[BaseModel]: + signature = inspect.signature(method) + schema = _published_schema(method, definition) + properties = schema.get("properties", {}) + required = set(schema.get("required", [])) + try: + hints = get_type_hints(method) + except (NameError, TypeError): + hints = {} + fields: dict[str, Any] = {} + for name, property_schema in properties.items(): + parameter = signature.parameters.get(name) + annotation = hints.get(name) + if annotation is None: + annotation = _schema_type(property_schema) + default: Any = ... + if parameter is not None and parameter.default is not inspect.Parameter.empty: + default = parameter.default + elif name not in required: + default = property_schema.get("default", None) + fields[name] = ( + annotation, + Field(default=default, description=property_schema.get("description")), + ) + model_name = f"{function_id.title().replace('_', '')}Input" + return create_model(model_name, **fields) + + +def _schema_type(schema: dict[str, Any]) -> Any: + schema_type = schema.get("type") + if schema_type == "string": + return str + if schema_type == "integer": + return int + if schema_type == "number": + return float + if schema_type == "boolean": + return bool + if schema_type == "array": + return GenericAlias(list, _schema_type(schema.get("items", {}))) + if schema_type == "object": + return dict[str, Any] + return Any + + +def _allows_none(annotation: Any) -> bool: + origin = get_origin(annotation) + return origin in {Union, UnionType} and type(None) in get_args(annotation) diff --git a/hyperforge/src/hyperforge/harness_sdk/clients.py b/hyperforge/src/hyperforge/harness_sdk/clients.py new file mode 100644 index 00000000..6ab82fbe --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/clients.py @@ -0,0 +1,458 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import random +from collections.abc import AsyncIterator, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Literal, Protocol + +import httpx +from pydantic import BaseModel, Field + +from hyperforge.llm import AsyncNuaClient, NUAConnection + +from .models import HarnessMessage, HarnessToolCall +from .tools import HarnessTool + +logger = logging.getLogger(__name__) + +PUBLIC_CHAT_COMPLETIONS_PATH = "/api/v1/predict/compat/chat/completions" +INTERNAL_CHAT_COMPLETIONS_PATH = "/api/internal/predict/compat/chat/completions" + +type ReasoningEffort = Literal["minimal", "low", "medium", "high"] +type ToolChoice = Literal["auto", "none", "required"] | dict[str, Any] + + +class ChatCompletionResponseFormat(BaseModel): + type: str + json_schema: dict[str, Any] | None = None + + +class ChatCompletionRequest(BaseModel): + """Provider-neutral request for an OpenAI-compatible chat endpoint.""" + + messages: list[dict[str, Any]] = Field(min_length=1) + model: str | None = None + stream: bool = True + temperature: float | None = None + max_tokens: int = 50_000 + top_p: float | None = None + frequency_penalty: float | None = None + presence_penalty: float | None = None + stop: str | list[str] | None = None + response_format: ChatCompletionResponseFormat | None = None + json_schema: dict[str, Any] | None = None + tools: list[dict[str, Any]] = Field(default_factory=list) + tool_choice: ToolChoice | None = None + reasoning_effort: ReasoningEffort | None = None + user: str | None = None + stream_options: dict[str, bool] = Field( + default_factory=lambda: {"include_usage": True} + ) + + +class ChatCompletionToolCallFunctionDelta(BaseModel): + name: str | None = None + arguments: str | None = None + + +class ChatCompletionToolCallDelta(BaseModel): + index: int + id: str | None = None + type: Literal["function"] | None = None + function: ChatCompletionToolCallFunctionDelta | None = None + + +class ChatCompletionDelta(BaseModel): + role: str | None = None + content: str | None = None + reasoning_content: str | None = None + refusal: str | None = None + tool_calls: list[ChatCompletionToolCallDelta] = Field(default_factory=list) + + +class ChatCompletionChoice(BaseModel): + index: int = 0 + delta: ChatCompletionDelta = Field(default_factory=ChatCompletionDelta) + finish_reason: str | None = None + + +class ChatCompletionUsage(BaseModel): + prompt_tokens: float = 0 + completion_tokens: float = 0 + total_tokens: float = 0 + prompt_tokens_details: dict[str, Any] | None = None + completion_tokens_details: dict[str, Any] | None = None + + +class ChatCompletionChunk(BaseModel): + id: str | None = None + choices: list[ChatCompletionChoice] = Field(default_factory=list) + created: int | None = None + model: str | None = None + object: str = "chat.completion.chunk" + usage: ChatCompletionUsage | None = None + system_fingerprint: str | None = None + service_tier: str | None = None + + +class NucliaChatCompletionsError(RuntimeError): + def __init__( + self, message: str, *, provider_data: dict[str, Any] | None = None + ) -> None: + super().__init__(message) + self.provider_data = provider_data or {} + + +class NucliaChatCompletionsClient: + """Chat-completions transport backed by Hyperforge's shared NUA client.""" + + def __init__( + self, + nua: AsyncNuaClient, + *, + timeout: float = 5 * 60, + owns_client: bool = False, + ) -> None: + self.nua = nua + self.timeout = timeout + self._owns_client = owns_client + + @classmethod + async def from_api_key( + cls, + api_key: str, + *, + base_url: str | None = None, + timeout: float = 5 * 60, + ) -> NucliaChatCompletionsClient: + nua = await NUAConnection(key=api_key).connect(base_url=base_url) + return cls(nua, timeout=timeout, owns_client=True) + + @classmethod + def in_cluster( + cls, + *, + url: str = "http://predict.learning.svc.cluster.local:8080", + account: str | None = None, + headers: Mapping[str, str] | None = None, + timeout: float = 5 * 60, + ) -> NucliaChatCompletionsClient: + return cls( + AsyncNuaClient.internal( + url=url, account=account, headers=dict(headers or {}) + ), + timeout=timeout, + owns_client=True, + ) + + async def aclose(self) -> None: + if self._owns_client: + await self.nua.aclose() + + async def stream( + self, request: ChatCompletionRequest + ) -> AsyncIterator[ChatCompletionChunk]: + payload = request.model_dump(exclude_none=True) + for attempt in range(1, 4): + yielded_chunk = False + try: + async for data in self.nua.chat_completions_stream( + payload, timeout=self.timeout + ): + if data.get("type") == "status" and data.get("code") == "ERROR": + raise NucliaChatCompletionsError( + data.get("details") or "Generation failed", + provider_data=data, + ) + chunk = ChatCompletionChunk.model_validate(data) + yielded_chunk = True + yield chunk + return + except (httpx.RequestError, httpx.HTTPStatusError) as exc: + error = NucliaChatCompletionsError( + f"Nuclia chat completions request failed: {exc}", + provider_data={ + "http_status": exc.response.status_code + if isinstance(exc, httpx.HTTPStatusError) + else None + }, + ) + except NucliaChatCompletionsError as exc: + error = exc + + status = error.provider_data.get("http_status") + retryable = status is None or status in {408, 429} or 500 <= status < 600 + if yielded_chunk or not retryable or attempt == 3: + raise error + delay = 0.5 * (2 ** (attempt - 1)) + random.uniform(0, 0.5) + logger.warning( + "Retrying Nuclia generation after pre-output failure: attempt=%d/%d model=%s delay=%.2fs error=%s", + attempt, + 3, + request.model, + delay, + error, + extra={ + "attempt": attempt, + "max_attempts": 3, + "model": request.model, + "message_count": len(request.messages), + "tool_count": len(request.tools), + "provider_data": error.provider_data, + "error": str(error), + }, + ) + await asyncio.sleep(delay) + + +@dataclass +class ModelDelta: + text: str = "" + reasoning: str = "" + tool_calls: list[HarnessToolCall] = field(default_factory=list) + input_tokens: float = 0 + output_tokens: float = 0 + trace_id: str | None = None + model: str | None = None + + +class ModelClient(Protocol): + def stream( + self, + *, + model: str, + reasoning_effort: ReasoningEffort | None, + messages: list[HarnessMessage], + tools: list[HarnessTool], + execution_context: dict[str, object], + ) -> AsyncIterator[ModelDelta]: ... + + +class NucliaModelClient: + """Harness adapter for Nuclia's OpenAI-compatible chat endpoint.""" + + def __init__( + self, + client: NucliaChatCompletionsClient, + *, + reasoning_effort: ReasoningEffort = "medium", + max_tokens: int = 50_000, + temperature: float | None = None, + top_p: float | None = None, + frequency_penalty: float | None = None, + presence_penalty: float | None = None, + stop: str | list[str] | None = None, + response_format: ChatCompletionResponseFormat | None = None, + json_schema: dict[str, Any] | None = None, + tool_choice: ToolChoice = "auto", + ) -> None: + self.client = client + self.reasoning_effort = reasoning_effort + self.max_tokens = max_tokens + self.temperature = temperature + self.top_p = top_p + self.frequency_penalty = frequency_penalty + self.presence_penalty = presence_penalty + self.stop = stop + self.response_format = response_format + self.json_schema = json_schema + self.tool_choice = tool_choice + + @classmethod + async def from_api_key( + cls, + api_key: str, + *, + base_url: str | None = None, + **options: Any, + ) -> NucliaModelClient: + nua = await NUAConnection(key=api_key).connect(base_url=base_url) + return cls(NucliaChatCompletionsClient(nua, owns_client=True), **options) + + @classmethod + async def in_cluster( + cls, + *, + url: str = "http://predict.learning.svc.cluster.local:8080", + account: str | None = None, + headers: Mapping[str, str] | None = None, + **options: Any, + ) -> NucliaModelClient: + nua = AsyncNuaClient.internal( + url=url, account=account, headers=dict(headers or {}) + ) + return cls(NucliaChatCompletionsClient(nua, owns_client=True), **options) + + @property + def nua(self) -> AsyncNuaClient: + return self.client.nua + + async def aclose(self) -> None: + await self.client.aclose() + + async def stream( + self, + *, + model: str, + reasoning_effort: ReasoningEffort | None, + messages: list[HarnessMessage], + tools: list[HarnessTool], + execution_context: dict[str, object], + ) -> AsyncIterator[ModelDelta]: + converted_messages = self._messages(messages) + self._validate_tool_history(converted_messages) + request = ChatCompletionRequest( + messages=converted_messages, + user=str(execution_context.get("user_id", "system")), + model=model, + reasoning_effort=reasoning_effort or self.reasoning_effort, + max_tokens=self.max_tokens, + temperature=self.temperature, + top_p=self.top_p, + frequency_penalty=self.frequency_penalty, + presence_penalty=self.presence_penalty, + stop=self.stop, + response_format=self.response_format, + json_schema=self.json_schema, + tools=[ + { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + }, + } + for tool in tools + ], + tool_choice=self.tool_choice, + ) + pending_calls: dict[int, dict[str, str]] = {} + emitted_calls: set[int] = set() + + def completed_calls() -> list[HarnessToolCall]: + completed: list[HarnessToolCall] = [] + for index, pending in sorted(pending_calls.items()): + if index in emitted_calls or not pending["name"]: + continue + try: + arguments = json.loads(pending["arguments"] or "{}") + except json.JSONDecodeError as exc: + arguments = {"_tool_error": f"Malformed tool arguments: {exc}"} + if not isinstance(arguments, dict): + arguments = { + "_tool_error": "Tool arguments must decode to a JSON object" + } + completed.append( + HarnessToolCall( + id=pending["id"] or None, + name=pending["name"], + arguments=arguments, + ) + ) + emitted_calls.add(index) + return completed + + async for chunk in self.client.stream(request): + if not chunk.choices: + yield ModelDelta( + input_tokens=chunk.usage.prompt_tokens if chunk.usage else 0, + output_tokens=chunk.usage.completion_tokens if chunk.usage else 0, + trace_id=chunk.id, + model=chunk.model, + ) + continue + for choice in chunk.choices: + for call in choice.delta.tool_calls: + pending = pending_calls.setdefault( + call.index, {"id": "", "name": "", "arguments": ""} + ) + if call.id: + pending["id"] = call.id + if call.function is not None: + if call.function.name: + pending["name"] += call.function.name + if call.function.arguments: + pending["arguments"] += call.function.arguments + yield ModelDelta( + text=choice.delta.content or choice.delta.refusal or "", + reasoning=choice.delta.reasoning_content or "", + tool_calls=completed_calls() + if choice.finish_reason == "tool_calls" + else [], + input_tokens=chunk.usage.prompt_tokens if chunk.usage else 0, + output_tokens=chunk.usage.completion_tokens if chunk.usage else 0, + trace_id=chunk.id, + model=chunk.model, + ) + remaining = completed_calls() + if remaining: + yield ModelDelta(tool_calls=remaining) + + @staticmethod + def _messages(messages: Sequence[HarnessMessage]) -> list[dict[str, Any]]: + converted: list[dict[str, Any]] = [] + for message in messages: + if message.role == "assistant": + value: dict[str, Any] = { + "role": "assistant", + "content": message.content or None, + } + if message.tool_calls: + value["tool_calls"] = [ + { + "id": call.id, + "type": "function", + "function": { + "name": call.name, + "arguments": json.dumps(call.arguments), + }, + } + for call in message.tool_calls + ] + converted.append(value) + elif message.role == "tool": + converted.append( + { + "role": "tool", + "content": message.content, + "tool_call_id": message.tool_call_id, + } + ) + else: + converted.append({"role": message.role, "content": message.content}) + return converted + + @staticmethod + def _validate_tool_history(messages: Sequence[dict[str, Any]]) -> None: + pending: set[str] = set() + for index, message in enumerate(messages): + role = message.get("role") + if role == "assistant": + if pending: + raise ValueError( + f"Incomplete tool history before message {index}; missing outputs for {sorted(pending)}" + ) + pending = { + str(call["id"]) + for call in message.get("tool_calls", []) + if isinstance(call, dict) and call.get("id") + } + elif role == "tool": + tool_call_id = message.get("tool_call_id") + if tool_call_id not in pending: + raise ValueError( + f"Unexpected tool output at message {index}: {tool_call_id}" + ) + pending.remove(str(tool_call_id)) + elif pending: + raise ValueError( + f"Incomplete tool history before message {index}; missing outputs for {sorted(pending)}" + ) + if pending: + raise ValueError( + f"Incomplete tool history at end of messages; missing outputs for {sorted(pending)}" + ) diff --git a/hyperforge/src/hyperforge/harness_sdk/context.py b/hyperforge/src/hyperforge/harness_sdk/context.py new file mode 100644 index 00000000..d4e0e85c --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/context.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Callable + +from pydantic import BaseModel + +from .models import HarnessContextReference, HarnessContextType + +type ContextFormatter = Callable[[BaseModel], str] + + +@dataclass(frozen=True) +class ContextDefinition: + schema: type[BaseModel] + formatter: ContextFormatter | None = None + + +_DEFINITIONS: dict[HarnessContextType, ContextDefinition] = {} + + +def register_context( + context_type: HarnessContextType, + schema: type[BaseModel], + formatter: ContextFormatter | None = None, +) -> None: + _DEFINITIONS[context_type] = ContextDefinition(schema=schema, formatter=formatter) + + +def make_context( + context_type: HarnessContextType, output: BaseModel +) -> HarnessContextReference: + definition = _DEFINITIONS.get(context_type) + if definition is None: + return HarnessContextReference( + type=context_type, content=output.model_dump(mode="json") + ) + value = definition.schema.model_validate(output.model_dump(mode="json")) + return HarnessContextReference( + type=context_type, content=value.model_dump(mode="json") + ) + + +def format_context(reference: HarnessContextReference) -> str: + if reference.type == HarnessContextType.RETRIEVAL: + values = reference.content.get("items", [reference.content]) + visible = [ + { + key: value + for key, value in item.items() + if key not in {"images"} and value not in (None, "", [], {}) + } + for item in values + ] + return json.dumps(visible, sort_keys=True, default=str, separators=(",", ":")) + definition = _DEFINITIONS.get(reference.type) + if definition is None: + return json.dumps( + reference.content, sort_keys=True, default=str, separators=(",", ":") + ) + value = definition.schema.model_validate(reference.content) + if definition.formatter is not None: + return definition.formatter(value) + return json.dumps( + value.model_dump(mode="json"), + sort_keys=True, + default=str, + separators=(",", ":"), + ) + + +def context_schema(context_type: HarnessContextType) -> dict[str, Any]: + return _DEFINITIONS[context_type].schema.model_json_schema() diff --git a/hyperforge/src/hyperforge/harness_sdk/harness.py b/hyperforge/src/hyperforge/harness_sdk/harness.py new file mode 100644 index 00000000..d15bda04 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/harness.py @@ -0,0 +1,1100 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import time +import uuid +from collections.abc import AsyncIterator, Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any, Literal, NamedTuple + +from pydantic import BaseModel + +from .agents import published_agent_to_tools +from .clients import ModelClient, ReasoningEffort +from .context import format_context +from .models import ( + HarnessConversation, + HarnessEvent, + HarnessEventType, + HarnessInboxItem, + HarnessMessage, + HarnessToolCall, +) +from .storage import HarnessStorageProtocol, InMemoryHarnessStorage +from .tools import HarnessTool +from .tools.core import DictOutput, SendMessageInput, SpawnAgentInput, create_core_tools +from .usage import HarnessUsage, UsageLimitExceeded, UsageLimits + +logger = logging.getLogger(__name__) + +SYSTEM_PROMPT = ( + "You are an agent running in a tool loop. Use tools when they help. " + "Incorporate tool results and new inbox messages. Return a direct answer when done." +) +DEFAULT_EVENT_QUEUE_SIZE = 1000 +POST_TOOL_RETRY_PROMPT = "Use the tool results above to answer the user's request now." +EMPTY_RESPONSE_RETRY_PROMPT = ( + "Your previous response was empty. Continue the task now: use an appropriate " + "tool when needed, otherwise answer the user's request directly." +) +MAX_EMPTY_RETRIES = 2 + + +@dataclass +class AgentResult: + text: str + reasoning: str = "" + tool_calls: list[HarnessToolCall] = field(default_factory=list) + input_tokens: float = 0 + output_tokens: float = 0 + + +@dataclass(frozen=True) +class HarnessConfig: + model: str + reasoning_effort: ReasoningEffort | None + model_client: ModelClient + tools: tuple[HarnessTool, ...] + system_prompt: str + title: str + disabled_core_tools: frozenset[str] + storage: HarnessStorageProtocol + execution_context: Mapping[str, Any] + category: str + tags: tuple[str, ...] + feedback_enabled: bool + usage_limits: UsageLimits + max_spawn_depth: int + max_concurrent_agents: int + event_queue_size: int + + +@dataclass +class TurnLoopState: + tool_result_pending: bool = False + empty_post_tool_retries: int = 0 + empty_retries: int = 0 + + def record_tool_calls(self) -> None: + self.tool_result_pending = True + self.empty_post_tool_retries = 0 + + def retry_prompt(self, result: AgentResult) -> str | None: + if result.text.strip(): + self.tool_result_pending = False + return None + if ( + self.tool_result_pending + and self.empty_post_tool_retries < MAX_EMPTY_RETRIES + ): + self.empty_post_tool_retries += 1 + return POST_TOOL_RETRY_PROMPT + if self.empty_retries < MAX_EMPTY_RETRIES: + self.empty_retries += 1 + return EMPTY_RESPONSE_RETRY_PROMPT + return None + + +class ChildAgent(NamedTuple): + harness: AgentHarness + task: asyncio.Task[str] + + +class LLMCallError(RuntimeError): + """A model-call failure with the provider context needed for diagnostics.""" + + def __init__( + self, + cause: Exception, + *, + call_id: str, + trace_id: str | None, + model: str, + ) -> None: + super().__init__(str(cause)) + self.cause = cause + self.call_id = call_id + self.trace_id = trace_id + self.model = model + + +class ChildAgentManager: + def __init__( + self, + owner: AgentHarness, + active_agent_ids: set[str] | None = None, + ) -> None: + self.owner = owner + self.children: dict[str, ChildAgent] = {} + self.results: dict[str, DictOutput] = {} + self.active_agent_ids = ( + active_agent_ids if active_agent_ids is not None else set() + ) + + def interrupt(self) -> None: + for child in self.children.values(): + child.harness.interrupt() + child.task.cancel() + + async def stop(self) -> None: + children = list(self.children.items()) + tasks = [child.task for _, child in children if not child.task.done()] + for child in self.children.values(): + child.harness.interrupt() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + for child_id, child in children: + self.results.setdefault(child_id, self._result(child_id, child.task)) + self.children.clear() + + async def spawn(self, input_value: SpawnAgentInput) -> DictOutput: + owner = self.owner + if owner.spawn_depth >= owner.max_spawn_depth: + raise ValueError(f"Maximum spawn depth reached: {owner.max_spawn_depth}") + if len(self.active_agent_ids) >= owner.max_concurrent_agents - 1: + return self._concurrency_limit_result() + + child_id = uuid.uuid4().hex + child = owner._create_child( + child_id, include_history=input_value.include_history + ) + child._child_agents.active_agent_ids = self.active_agent_ids + self.active_agent_ids.add(child_id) + + async def run_child() -> str: + result = "interrupted" + async for event in child.run(input_value.prompt): + await owner._publish_event(event) + if event.agent_id != child.agent_id: + continue + if event.type == HarnessEventType.TURN_COMPLETED: + result = str(event.payload.get("text", "")) + elif event.type == HarnessEventType.TURN_FAILED: + raise RuntimeError(str(event.payload.get("error", "Child failed"))) + return result + + task = asyncio.create_task(run_child()) + task.add_done_callback(lambda _: self.active_agent_ids.discard(child_id)) + self.children[child_id] = ChildAgent(child, task) + await self._wait_until_started(child, task) + await owner.emit( + HarnessEventType.AGENT_STARTED, + { + "agent_id": child_id, + "parent_agent_id": owner.agent_id, + "prompt": input_value.prompt, + "spawn_depth": child.spawn_depth, + "max_spawn_depth": owner.max_spawn_depth, + }, + ) + return DictOutput(value={"agent_id": child_id}) + + async def send_message(self, input_value: SendMessageInput) -> DictOutput: + child = self.children.get(input_value.agent_id) + if child is None: + raise ValueError("Agent not found") + message_id = await child.harness.steer(input_value.message, sender="agent") + return DictOutput(value={"message_id": message_id}) + + async def wait(self, agent_id: str) -> DictOutput: + completed = self.results.get(agent_id) + if completed is not None: + return completed + child = self.children.get(agent_id) + if child is None: + raise ValueError("Agent not found") + try: + await child.task + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + except Exception: + pass + self.active_agent_ids.discard(agent_id) + result = self._result(agent_id, child.task) + self.results[agent_id] = result + self.children.pop(agent_id, None) + return result + + def _concurrency_limit_result(self) -> DictOutput: + waitable_agent_ids = [ + agent_id + for agent_id, child in self.children.items() + if not child.task.done() + ] + if waitable_agent_ids: + message = ( + "The conversation is at its concurrent agent limit. " + "Call wait_agent for an existing child before trying to spawn another." + ) + else: + message = ( + "The conversation is at its concurrent agent limit, and this agent has no child " + "it can wait for. Finish the current work so the parent can wait before spawning more." + ) + return DictOutput( + value={ + "status": "concurrency_limit_reached", + "max_concurrent_agents": self.owner.max_concurrent_agents, + "active_agents": len(self.active_agent_ids) + 1, + "waitable_agent_ids": waitable_agent_ids, + "message": message, + } + ) + + @staticmethod + async def _wait_until_started(child: AgentHarness, task: asyncio.Task[str]) -> None: + if task.done(): + return + started_task = asyncio.create_task(child._model_started.wait()) + await asyncio.wait({task, started_task}, return_when=asyncio.FIRST_COMPLETED) + if not started_task.done(): + started_task.cancel() + await asyncio.gather(started_task, return_exceptions=True) + await asyncio.sleep(0) + + @staticmethod + def _result(agent_id: str, task: asyncio.Task[str]) -> DictOutput: + if task.cancelled(): + value: dict[str, Any] = {"agent_id": agent_id, "status": "cancelled"} + elif (error := task.exception()) is not None: + value = {"agent_id": agent_id, "status": "failed", "error": str(error)} + else: + result = task.result() + if result == "interrupted": + return DictOutput(value={"agent_id": agent_id, "status": "cancelled"}) + value = { + "agent_id": agent_id, + "status": "completed", + "result": result, + } + return DictOutput(value=value) + + +class AgentHarness: + def __init__( + self, + *, + model: str, + model_client: ModelClient, + reasoning_effort: ReasoningEffort | None = None, + tools: Iterable[HarnessTool | Any] = (), + system_prompt: str = SYSTEM_PROMPT, + title: str = "New conversation", + conversation_id: str | None = None, + disabled_core_tools: Iterable[str] = (), + storage: HarnessStorageProtocol | None = None, + execution_context: Mapping[str, Any] | None = None, + category: str = "assistant", + tags: Iterable[str] = (), + feedback_enabled: bool = False, + usage_limits: UsageLimits | None = None, + max_spawn_depth: int = 1, + max_concurrent_agents: int = 4, + spawn_depth: int = 0, + parent_agent_id: str | None = None, + event_queue_size: int = DEFAULT_EVENT_QUEUE_SIZE, + ) -> None: + if max_spawn_depth < 0: + raise ValueError("max_spawn_depth must be non-negative") + if max_concurrent_agents < 1: + raise ValueError("max_concurrent_agents must be positive") + if spawn_depth < 0 or spawn_depth > max_spawn_depth: + raise ValueError("spawn_depth must be between zero and max_spawn_depth") + if event_queue_size <= 0: + raise ValueError("event_queue_size must be positive") + normalized_tools = tuple( + tool if isinstance(tool, HarnessTool) else tool.as_tool() for tool in tools + ) + normalized_storage = storage or InMemoryHarnessStorage() + normalized_context = dict(execution_context or {}) + normalized_tags = tuple(dict.fromkeys(tags)) + normalized_limits = usage_limits or UsageLimits() + self._config = HarnessConfig( + model=model, + reasoning_effort=reasoning_effort, + model_client=model_client, + tools=normalized_tools, + system_prompt=system_prompt, + title=title, + disabled_core_tools=frozenset(disabled_core_tools), + storage=normalized_storage, + execution_context=normalized_context, + category=category, + tags=normalized_tags, + feedback_enabled=feedback_enabled, + usage_limits=normalized_limits, + max_spawn_depth=max_spawn_depth, + max_concurrent_agents=max_concurrent_agents, + event_queue_size=event_queue_size, + ) + self.model = self._config.model + self.reasoning_effort = self._config.reasoning_effort + self.model_client = self._config.model_client + self.system_prompt = self._config.system_prompt + self.title = self._config.title + self.conversation_id = conversation_id or uuid.uuid4().hex + self.agent_id = uuid.uuid4().hex + self.disabled_core_tools = set(self._config.disabled_core_tools) + self.storage = self._config.storage + self.execution_context = dict(self._config.execution_context) + self.category = self._config.category + self.tags = list(self._config.tags) + self.usage_limits = self._config.usage_limits + self.max_spawn_depth = self._config.max_spawn_depth + self.max_concurrent_agents = self._config.max_concurrent_agents + self.spawn_depth = spawn_depth + self.parent_agent_id = parent_agent_id + self.usage = HarnessUsage() + self._usage_started_at: float | None = None + self._owns_usage = True + self._persist_events = True + self.conversation: HarnessConversation | None = None + self.messages: list[HarnessMessage] = [] + self.inbox: asyncio.Queue[HarnessInboxItem] = asyncio.Queue() + self.event_queue_size = event_queue_size + self._event_queue: asyncio.Queue[HarnessEvent] | None = None + self._child_agents = ChildAgentManager(self) + self.feedback_requests: dict[str, asyncio.Future[str]] = {} + self._pending_messages: list[HarnessMessage] = [] + self._interrupted = asyncio.Event() + self._model_started = asyncio.Event() + self._run_lock = asyncio.Lock() + self._active_task: asyncio.Task[Any] | None = None + self._turn_id: str | None = None + self._last_event_id: str | None = None + self._external_tools = list(self._config.tools) + self._tools = {tool.name: tool for tool in self._external_tools} + self._active_lazy_tools: set[str] = set() + if len(self._tools) != len(self._external_tools): + raise ValueError("Tool names must be unique") + self._install_core_tools(self._config.feedback_enabled) + + @property + def children(self) -> dict[str, ChildAgent]: + return self._child_agents.children + + @property + def child_results(self) -> dict[str, DictOutput]: + return self._child_agents.results + + async def load(self, *, create: bool = True) -> None: + conversation = await self.storage.get_conversation(self.conversation_id) + if conversation is None: + if not create: + raise ValueError("Conversation not found") + conversation = HarnessConversation( + id=self.conversation_id, + title=self.title, + category=self.category, + tags=self.tags, + metadata=self._persisted_metadata(), + ) + await self.storage.create_conversation(conversation) + self.conversation = conversation + self.messages = [] + pending_tool_outputs: dict[str, HarnessMessage] = {} + async for event in self.storage.iter_events(self.conversation_id): + self._last_event_id = event.id + if event.parent_agent_id is not None: + continue + if event.type in { + HarnessEventType.MESSAGE_ADDED, + HarnessEventType.MESSAGES_ADDED, + }: + raw_messages = ( + [event.payload["message"]] + if event.type == HarnessEventType.MESSAGE_ADDED + else event.payload["messages"] + ) + for raw_message in raw_messages: + message = HarnessMessage.model_validate(raw_message) + if message.role == "tool" and not any( + message.tool_call_id == call.id + for existing in self.messages + for call in existing.tool_calls + ): + if message.tool_call_id is not None: + pending_tool_outputs[message.tool_call_id] = message + continue + self.messages.append(message) + for call in message.tool_calls: + if call.id in pending_tool_outputs: + self.messages.append(pending_tool_outputs.pop(call.id)) + elif event.type == HarnessEventType.COMPACTED: + self.messages = [ + HarnessMessage.model_validate(item) + for item in event.payload["messages"] + ] + pending_tool_outputs.clear() + if pending_tool_outputs: + logger.warning( + "Ignoring orphaned tool outputs while loading conversation: conversation=%s call_ids=%s", + self.conversation_id, + sorted(pending_tool_outputs), + extra={ + "conversation_id": self.conversation_id, + "call_ids": sorted(pending_tool_outputs), + }, + ) + self.messages = _complete_tool_history(self.messages) + + def _persisted_metadata(self) -> dict[str, Any]: + metadata = self.execution_context.get("conversation_metadata", {}) + return dict(metadata) if isinstance(metadata, Mapping) else {} + + def add_messages(self, messages: Iterable[HarnessMessage | dict[str, Any]]) -> None: + self._pending_messages.extend( + message + if isinstance(message, HarnessMessage) + else HarnessMessage.model_validate(message) + for message in messages + ) + + @staticmethod + def to_tools( + agent: Any, *, namespace: str | None = None, manager: Any = None + ) -> list[HarnessTool[Any, Any]]: + """Expose a legacy agent's published functions as harness tools.""" + return published_agent_to_tools(agent, namespace=namespace, manager=manager) + + def iter_tools(self, *, include_inactive: bool = False) -> Iterable[HarnessTool]: + """Iterate over registered tools available to this harness.""" + return ( + tool + for tool in self._tools.values() + if include_inactive + or not tool.lazy_load + or tool.name in self._active_lazy_tools + ) + + async def save_messages(self) -> None: + await self._flush_pending_messages() + + async def steer( + self, content: str, *, sender: Literal["user", "agent"] = "user" + ) -> str: + item = HarnessInboxItem(id=uuid.uuid4().hex, sender=sender, content=content) + await self.inbox.put(item) + await self.emit( + HarnessEventType.INBOX_ADDED, {"item": item.model_dump(mode="json")} + ) + return item.id + + def interrupt(self) -> None: + self._interrupted.set() + try: + current = asyncio.current_task() + except RuntimeError: + current = None + if self._active_task is not None and self._active_task is not current: + self._active_task.cancel() + self._child_agents.interrupt() + + def _check_limit(self, name: str, value: float) -> None: + limit = getattr(self.usage_limits, name) + if limit is not None and value > limit: + raise UsageLimitExceeded(f"{name} limit exceeded: {value} > {limit}") + + def _start_usage(self) -> None: + if self._owns_usage: + self.usage = HarnessUsage() + self._usage_started_at = time.monotonic() + + def _remaining_time(self) -> float | None: + if self.usage_limits.max_time is None: + return None + if self._usage_started_at is None: + self._usage_started_at = time.monotonic() + remaining = self.usage_limits.max_time - ( + time.monotonic() - self._usage_started_at + ) + if remaining <= 0: + raise UsageLimitExceeded("max_time limit exceeded") + return remaining + + async def _with_time_limit(self, awaitable: Any) -> Any: + try: + remaining = self._remaining_time() + except BaseException: + close = getattr(awaitable, "close", None) + if close is not None: + close() + raise + if remaining is None: + return await awaitable + try: + async with asyncio.timeout(remaining): + return await awaitable + except TimeoutError as exc: + raise UsageLimitExceeded("max_time limit exceeded") from exc + + @property + def turn_id(self) -> str | None: + return self._turn_id + + async def request_feedback( + self, + *, + question: str, + response_schema: dict[str, Any], + timeout_ms: int, + ) -> str | None: + request_id = uuid.uuid4().hex + future = asyncio.get_running_loop().create_future() + self.feedback_requests[request_id] = future + await self.emit( + HarnessEventType.FEEDBACK_REQUESTED, + { + "request_id": request_id, + "question": question, + "response_schema": response_schema, + "timeout_ms": timeout_ms, + }, + ) + try: + return await asyncio.wait_for(future, timeout=timeout_ms / 1000) + except TimeoutError: + return None + finally: + self.feedback_requests.pop(request_id, None) + + async def respond_feedback(self, request_id: str, response: str) -> None: + future = self.feedback_requests.get(request_id) + if future is None or future.done(): + raise ValueError("Feedback request not found") + future.set_result(response) + await self.emit( + HarnessEventType.FEEDBACK_RESOLVED, + {"request_id": request_id, "response": response}, + ) + + async def _publish_event(self, event: HarnessEvent) -> None: + queue = self._event_queue + if queue is None: + return + if ( + event.type + in { + HarnessEventType.TEXT_DELTA, + HarnessEventType.REASONING_DELTA, + } + and queue.qsize() >= self.event_queue_size + ): + return + queue.put_nowait(event) + + async def history(self) -> AsyncIterator[HarnessEvent]: + async for event in self.storage.iter_events(self.conversation_id): + yield event + + async def emit( + self, + event_type: HarnessEventType, + payload: dict[str, Any], + *, + persist: bool = True, + ) -> HarnessEvent: + event = HarnessEvent( + id=uuid.uuid4().hex, + conversation_id=self.conversation_id, + turn_id=self._turn_id, + agent_id=self.agent_id, + parent_agent_id=self.parent_agent_id, + category=self.category, + tags=self.tags, + metadata=self._persisted_metadata(), + type=event_type, + payload=payload, + ) + if persist and self._persist_events: + await self.storage.append_event(event) + self._last_event_id = event.id + await self._publish_event(event) + return event + + async def llm(self, tools: Iterable[HarnessTool] | None = None) -> AgentResult: + await self._flush_pending_messages() + if tools is None: + selected = list(self.iter_tools()) + else: + selected = [] + seen: set[str] = set() + for requested in tools: + registered = self._tools.get(requested.name) + if ( + registered is not None + and registered.name not in seen + and ( + not registered.lazy_load + or registered.name in self._active_lazy_tools + ) + ): + selected.append(registered) + seen.add(registered.name) + self.usage.turns += 1 + self._check_limit("max_turns", self.usage.turns) + result = AgentResult(text="") + call_id = uuid.uuid4().hex + history_head_event_id = self._last_event_id + await self.emit( + HarnessEventType.LLM_STARTED, + { + "call_id": call_id, + "history_head_event_id": history_head_event_id, + "message_count": len(self.messages), + "model": self.model, + "reasoning_effort": self.reasoning_effort, + "tools": [tool.name for tool in selected], + }, + ) + trace_id: str | None = None + resolved_model: str | None = None + usage_recorded = False + try: + self._model_started.set() + async for delta in self.model_client.stream( + model=self.model, + reasoning_effort=self.reasoning_effort, + messages=self.messages, + tools=selected, + execution_context=self.execution_context, + ): + result.text += delta.text + result.reasoning += delta.reasoning + result.tool_calls.extend(delta.tool_calls) + result.input_tokens = max(result.input_tokens, delta.input_tokens) + result.output_tokens = max(result.output_tokens, delta.output_tokens) + trace_id = delta.trace_id or trace_id + resolved_model = delta.model or resolved_model + if delta.text: + await self.emit( + HarnessEventType.TEXT_DELTA, {"text": delta.text}, persist=False + ) + if delta.reasoning: + await self.emit( + HarnessEventType.REASONING_DELTA, + {"text": delta.reasoning}, + persist=False, + ) + self.usage.input_tokens += result.input_tokens + self.usage.output_tokens += result.output_tokens + usage_recorded = True + self._check_limit("max_input_tokens", self.usage.input_tokens) + self._check_limit("max_output_tokens", self.usage.output_tokens) + except asyncio.CancelledError: + raise + except UsageLimitExceeded as exc: + await self._emit_llm_result( + HarnessEventType.LLM_FAILED, + call_id, + history_head_event_id, + result, + trace_id, + resolved_model, + error=str(exc), + ) + raise + except Exception as exc: + if not usage_recorded: + self.usage.input_tokens += result.input_tokens + self.usage.output_tokens += result.output_tokens + await self._emit_llm_result( + HarnessEventType.LLM_FAILED, + call_id, + history_head_event_id, + result, + trace_id, + resolved_model, + error=str(exc), + ) + raise LLMCallError( + exc, + call_id=call_id, + trace_id=trace_id, + model=resolved_model or self.model, + ) from exc + await self._emit_llm_result( + HarnessEventType.LLM_COMPLETED, + call_id, + history_head_event_id, + result, + trace_id, + resolved_model, + ) + return result + + async def _emit_llm_result( + self, + event_type: HarnessEventType, + call_id: str, + history_head_event_id: str | None, + result: AgentResult, + trace_id: str | None, + resolved_model: str | None, + *, + error: str | None = None, + ) -> None: + payload: dict[str, Any] = { + "call_id": call_id, + "input_tokens": result.input_tokens, + "output_tokens": result.output_tokens, + "history_head_event_id": history_head_event_id, + "trace_id": trace_id, + "model": resolved_model or self.model, + "text": result.text, + "reasoning": result.reasoning, + } + if error is not None: + payload["error"] = error + await self.emit(event_type, payload) + + async def __aenter__(self) -> AgentHarness: + return self + + async def __aexit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: + self.interrupt() + task = self._active_task + if task is not None: + await asyncio.gather(task, return_exceptions=True) + + async def run(self, prompt: str) -> AsyncIterator[HarnessEvent]: + async with self._run_lock: + if self._event_queue is not None: + raise RuntimeError("A turn is already running") + self._event_queue = asyncio.Queue() + self._model_started.clear() + task = asyncio.create_task(self._execute_turn(prompt)) + self._active_task = task + try: + terminal = False + while not terminal or not task.done(): + event_task = asyncio.create_task(self._event_queue.get()) + done, _ = await asyncio.wait( + {task, event_task}, return_when=asyncio.FIRST_COMPLETED + ) + if task in done and not event_task.done(): + event_task.cancel() + await asyncio.gather(event_task, return_exceptions=True) + break + event = event_task.result() + yield event + terminal = ( + event.agent_id == self.agent_id + and event.turn_id == self._turn_id + and event.type + in { + HarnessEventType.TURN_COMPLETED, + HarnessEventType.TURN_FAILED, + HarnessEventType.TURN_INTERRUPTED, + } + ) + await task + finally: + if not task.done(): + self.interrupt() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + self._active_task = None + self._event_queue = None + self._interrupted.clear() + + async def _execute_turn(self, prompt: str) -> str: + try: + return await self._run(prompt) + except asyncio.CancelledError: + current = asyncio.current_task() + if self._interrupted.is_set() and not ( + current is not None and current.cancelling() > 1 + ): + await self._stop_children() + await self.emit(HarnessEventType.TURN_INTERRUPTED, {}) + return "interrupted" + raise + + async def _run(self, prompt: str) -> str: + if self.conversation is None: + await self.load() + self._start_usage() + self._turn_id = uuid.uuid4().hex + await self.emit(HarnessEventType.TURN_STARTED, {"prompt": prompt}) + latest_system = next( + ( + message + for message in reversed([*self.messages, *self._pending_messages]) + if message.role == "system" + ), + None, + ) + if latest_system is None or latest_system.content != self.system_prompt: + await self._append_message( + HarnessMessage(role="system", content=self.system_prompt) + ) + await self._flush_pending_messages() + await self._append_message(HarnessMessage(role="user", content=prompt)) + loop_state = TurnLoopState() + try: + if self._interrupted.is_set(): + await self.emit(HarnessEventType.TURN_INTERRUPTED, {}) + return "interrupted" + while not self._interrupted.is_set(): + await self._consume_inbox() + result = await self._with_time_limit(self.llm()) + if result.tool_calls: + await self._with_time_limit( + self._execute_tool_calls(result.tool_calls, result.text) + ) + loop_state.record_tool_calls() + continue + if retry_prompt := loop_state.retry_prompt(result): + await self._append_message( + HarnessMessage(role="user", content=retry_prompt) + ) + continue + await self._append_message( + HarnessMessage(role="assistant", content=result.text) + ) + if not self.inbox.empty(): + continue + await self.emit(HarnessEventType.TURN_COMPLETED, {"text": result.text}) + await self._stop_children() + return result.text + await self._stop_children() + await self.emit(HarnessEventType.TURN_INTERRUPTED, {}) + return "interrupted" + except Exception as exc: + parent_agent_id = self.execution_context.get("parent_agent_id") + log = logger.warning if parent_agent_id is not None else logger.exception + failure = exc.cause if isinstance(exc, LLMCallError) else exc + context = { + "conversation_id": self.conversation_id, + "turn_id": self._turn_id, + "agent_id": self.agent_id, + "parent_agent_id": parent_agent_id, + "call_id": exc.call_id if isinstance(exc, LLMCallError) else None, + "trace_id": exc.trace_id if isinstance(exc, LLMCallError) else None, + "model": exc.model if isinstance(exc, LLMCallError) else self.model, + "message_count": len(self.messages), + "tool_count": len(self._tools), + "error_type": type(failure).__name__, + "error": str(failure), + } + if parent_agent_id is not None: + log( + "Sub-agent turn failed: conversation=%s parent=%s agent=%s turn=%s call=%s model=%s error=%s", + self.conversation_id, + parent_agent_id, + self.agent_id, + self._turn_id, + context["call_id"], + context["model"], + failure, + extra=context, + ) + else: + log( + "Agent turn failed: conversation=%s agent=%s turn=%s call=%s model=%s error=%s", + self.conversation_id, + self.agent_id, + self._turn_id, + context["call_id"], + context["model"], + failure, + extra=context, + ) + await self.emit(HarnessEventType.TURN_FAILED, {"error": str(failure)}) + await self._stop_children() + raise + + async def _stop_children(self) -> None: + await self._child_agents.stop() + + async def _append_message(self, message: HarnessMessage) -> None: + self.messages.append(message) + await self.emit( + HarnessEventType.MESSAGE_ADDED, {"message": message.model_dump(mode="json")} + ) + + async def _append_messages(self, messages: list[HarnessMessage]) -> None: + self.messages.extend(messages) + await self.emit( + HarnessEventType.MESSAGES_ADDED, + {"messages": [message.model_dump(mode="json") for message in messages]}, + ) + + async def _flush_pending_messages(self) -> None: + pending, self._pending_messages = self._pending_messages, [] + for message in pending: + await self._append_message(message) + + async def _consume_inbox(self) -> None: + items: list[HarnessInboxItem] = [] + while not self.inbox.empty(): + items.append(self.inbox.get_nowait()) + if not items: + return + rendered = "New inbox messages:\n" + "\n".join( + f"- {item.sender}: {item.content}" for item in items + ) + await self._append_message(HarnessMessage(role="user", content=rendered)) + for item in items: + await self.emit(HarnessEventType.INBOX_CONSUMED, {"id": item.id}) + + async def _execute_tool_calls( + self, calls: list[HarnessToolCall], text: str = "" + ) -> None: + self.usage.tool_calls += len(calls) + self._check_limit("max_tool_calls", self.usage.tool_calls) + calls = [ + call if call.id else call.model_copy(update={"id": uuid.uuid4().hex}) + for call in calls + ] + messages = await asyncio.gather( + *(self._execute_tool_call(call) for call in calls) + ) + await self._append_messages( + [ + HarnessMessage(role="assistant", content=text, tool_calls=calls), + *messages, + ] + ) + + async def _execute_tool_call(self, call: HarnessToolCall) -> HarnessMessage: + assert call.id is not None + await self.emit( + HarnessEventType.TOOL_REQUESTED, {"call": call.model_dump(mode="json")} + ) + tool = self._tools.get(call.name) + try: + output: BaseModel | None = None + if tool is None: + raise ValueError(f"Unknown tool: {call.name}") + elif tool.lazy_load and tool.name not in self._active_lazy_tools: + raise ValueError(f"Tool is not active: {call.name}") + else: + output = await tool.execute(self, call.arguments) + result = output.model_dump(mode="json") + reference = tool.context(output) + content = format_context(reference) + await self.emit( + HarnessEventType.TOOL_COMPLETED, + {"call_id": call.id, "tool": call.name, "result": result}, + ) + return HarnessMessage( + role="tool", + tool_call_id=call.id, + tool_name=call.name, + content=content, + context=reference, + ) + except Exception as exc: + logger.warning( + "Agent tool execution failed", + extra={ + "conversation_id": self.conversation_id, + "turn_id": self._turn_id, + "agent_id": self.agent_id, + "call_id": call.id, + "tool": call.name, + "error_type": type(exc).__name__, + "error": str(exc), + }, + ) + result = {"error": str(exc)} + await self.emit( + HarnessEventType.TOOL_FAILED, + {"call_id": call.id, "tool": call.name, "result": result}, + ) + return HarnessMessage( + role="tool", + tool_call_id=call.id, + tool_name=call.name, + content=json.dumps(result, default=str, separators=(",", ":")), + ) + + def _install_core_tools(self, feedback_enabled: bool) -> None: + core_tools = create_core_tools(feedback_enabled=feedback_enabled) + enabled_core_tools = { + name: tool + for name, tool in core_tools.items() + if name not in self.disabled_core_tools + and not (name == "spawn_agent" and self.spawn_depth >= self.max_spawn_depth) + } + conflicts = self._tools.keys() & enabled_core_tools.keys() + if conflicts: + raise ValueError( + f"Tool names conflict with core tools: {', '.join(sorted(conflicts))}" + ) + self._tools.update(enabled_core_tools) + + def _create_child(self, child_id: str, *, include_history: bool) -> AgentHarness: + child_context = { + **self.execution_context, + "parent_agent_id": self.agent_id, + "spawn_agent_id": child_id, + } + child = type(self)( + model=self._config.model, + model_client=self._config.model_client, + reasoning_effort=self._config.reasoning_effort, + tools=self._config.tools, + system_prompt=self._config.system_prompt, + title=self._config.title, + conversation_id=self.conversation_id, + disabled_core_tools=self._config.disabled_core_tools, + storage=self._config.storage, + execution_context=child_context, + category=self._config.category, + tags=self._config.tags, + feedback_enabled=self._config.feedback_enabled, + usage_limits=self._config.usage_limits, + max_spawn_depth=self._config.max_spawn_depth, + max_concurrent_agents=self._config.max_concurrent_agents, + spawn_depth=self.spawn_depth + 1, + parent_agent_id=self.agent_id, + event_queue_size=self._config.event_queue_size, + ) + child.agent_id = child_id + child.usage = self.usage + child._usage_started_at = self._usage_started_at + child._owns_usage = False + child.conversation = self.conversation + if include_history: + child.messages = list(self.messages) + return child + + +def _complete_tool_history(messages: list[HarnessMessage]) -> list[HarnessMessage]: + complete: list[HarnessMessage] = [] + index = 0 + while index < len(messages): + message = messages[index] + if message.role != "assistant" or not message.tool_calls: + complete.append(message) + index += 1 + continue + expected = {call.id for call in message.tool_calls if call.id is not None} + outputs: list[HarnessMessage] = [] + cursor = index + 1 + while cursor < len(messages) and messages[cursor].role == "tool": + outputs.append(messages[cursor]) + cursor += 1 + if expected and expected == {output.tool_call_id for output in outputs}: + complete.extend([message, *outputs]) + else: + logger.warning( + "Ignoring incomplete tool exchange while loading conversation: call_ids=%s", + sorted(expected), + ) + index = cursor + return complete + + +async def create_agent(**kwargs: Any) -> AgentHarness: + create = bool(kwargs.pop("create", True)) + harness = AgentHarness(**kwargs) + await harness.load(create=create) + return harness diff --git a/hyperforge/src/hyperforge/harness_sdk/models.py b/hyperforge/src/hyperforge/harness_sdk/models.py new file mode 100644 index 00000000..b452960f --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/models.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from enum import StrEnum +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +def utcnow() -> datetime: + return datetime.now(timezone.utc) + + +class HarnessEventType(StrEnum): + CONVERSATION_STARTED = "conversation.started" + AGENT_STARTED = "agent.started" + TURN_STARTED = "turn.started" + MESSAGE_ADDED = "message.added" + MESSAGES_ADDED = "messages.added" + REASONING_DELTA = "reasoning.delta" + TEXT_DELTA = "text.delta" + TOOL_REQUESTED = "tool.requested" + TOOL_COMPLETED = "tool.completed" + TOOL_FAILED = "tool.failed" + INBOX_ADDED = "inbox.added" + INBOX_CONSUMED = "inbox.consumed" + FEEDBACK_REQUESTED = "feedback.requested" + FEEDBACK_RESOLVED = "feedback.resolved" + MEMORY_REMEMBERED = "memory.remembered" + MEMORY_FORGOTTEN = "memory.forgotten" + COMPACTED = "context.compacted" + LLM_STARTED = "llm.started" + LLM_COMPLETED = "llm.completed" + LLM_FAILED = "llm.failed" + WORKFLOW_MESSAGE = "workflow.message" + TURN_COMPLETED = "turn.completed" + TURN_INTERRUPTED = "turn.interrupted" + TURN_FAILED = "turn.failed" + + +class HarnessContextType(StrEnum): + TOOL_RESULT = "tool_result" + RETRIEVAL = "retrieval" + STATUS = "status" + ENTITY = "entity" + COLLECTION = "collection" + QUERY = "query" + STRUCTURED = "structured" + + +class HarnessToolCall(BaseModel): + id: str | None = None + name: str + arguments: dict[str, Any] = Field(default_factory=dict) + + +class HarnessContextReference(BaseModel): + type: HarnessContextType + content: dict[str, Any] + + +class HarnessMessage(BaseModel): + role: Literal["system", "user", "assistant", "tool"] + content: str + tool_calls: list[HarnessToolCall] = Field(default_factory=list) + tool_call_id: str | None = None + tool_name: str | None = None + context: HarnessContextReference | None = None + + +class HarnessConversation(BaseModel): + id: str + title: str = "New conversation" + category: str = "assistant" + tags: list[str] = Field(default_factory=list) + metadata: dict[str, Any] = Field(default_factory=dict) + created_datetime: datetime = Field(default_factory=utcnow) + updated_datetime: datetime = Field(default_factory=utcnow) + + +class HarnessEvent(BaseModel): + id: str + conversation_id: str + turn_id: str | None = None + agent_id: str | None = None + parent_agent_id: str | None = None + category: str = "assistant" + tags: list[str] = Field(default_factory=list) + metadata: dict[str, Any] = Field(default_factory=dict) + type: HarnessEventType + created_datetime: datetime = Field(default_factory=utcnow) + payload: dict[str, Any] = Field(default_factory=dict) + + +class HarnessMemory(BaseModel): + id: str + scope: str + text: str + metadata: dict[str, Any] = Field(default_factory=dict) + created_datetime: datetime = Field(default_factory=utcnow) + + +class HarnessInboxItem(BaseModel): + id: str + sender: Literal["user", "agent"] + content: str + created_datetime: datetime = Field(default_factory=utcnow) diff --git a/hyperforge/src/hyperforge/harness_sdk/schema.py b/hyperforge/src/hyperforge/harness_sdk/schema.py new file mode 100644 index 00000000..6d4788e5 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/schema.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from copy import deepcopy +from typing import Any + + +def flatten_json_schema(schema: dict[str, Any]) -> dict[str, Any]: + """Inline local schema references and remove definition containers.""" + root = deepcopy(schema) + + def resolve(ref: str, stack: tuple[str, ...]) -> dict[str, Any]: + if not ref.startswith("#/"): + raise ValueError( + f"External JSON Schema references are not supported: {ref}" + ) + if ref in stack: + raise ValueError( + f"Recursive JSON Schema reference cannot be flattened: {ref}" + ) + value: Any = root + for raw_part in ref[2:].split("/"): + part = raw_part.replace("~1", "/").replace("~0", "~") + if not isinstance(value, dict) or part not in value: + raise ValueError(f"JSON Schema reference not found: {ref}") + value = value[part] + if not isinstance(value, dict): + raise ValueError(f"JSON Schema reference must resolve to an object: {ref}") + return walk(deepcopy(value), (*stack, ref)) + + def walk(value: Any, stack: tuple[str, ...] = ()) -> Any: + if isinstance(value, list): + return [walk(item, stack) for item in value] + if not isinstance(value, dict): + return value + current = dict(value) + ref = current.pop("$ref", None) + if ref is not None: + if not isinstance(ref, str): + raise ValueError("JSON Schema $ref must be a string") + current = {**resolve(ref, stack), **current} + current.pop("$defs", None) + current.pop("definitions", None) + return {key: walk(item, stack) for key, item in current.items()} + + return walk(root) diff --git a/hyperforge/src/hyperforge/harness_sdk/storage.py b/hyperforge/src/hyperforge/harness_sdk/storage.py new file mode 100644 index 00000000..272cfc61 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/storage.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from collections import defaultdict +from collections.abc import AsyncIterator +from typing import Protocol + +from .models import HarnessConversation, HarnessEvent, HarnessMemory + + +class HarnessStorageProtocol(Protocol): + async def create_conversation(self, conversation: HarnessConversation) -> None: ... + + async def get_conversation( + self, conversation_id: str + ) -> HarnessConversation | None: ... + + async def update_conversation(self, conversation: HarnessConversation) -> None: ... + + async def append_event(self, event: HarnessEvent) -> None: ... + + def iter_events(self, conversation_id: str) -> AsyncIterator[HarnessEvent]: ... + + async def remember(self, memory: HarnessMemory) -> None: ... + + async def recall( + self, + *, + scope: str, + query: str, + limit: int = 20, + ) -> list[HarnessMemory]: ... + + async def forget(self, memory_id: str) -> None: ... + + +class InMemoryHarnessStorage: + """Process-local storage intended for tests and ephemeral harnesses.""" + + def __init__(self) -> None: + self.conversations: dict[str, HarnessConversation] = {} + self.events: dict[str, list[HarnessEvent]] = defaultdict(list) + self.memories: dict[str, HarnessMemory] = {} + + async def create_conversation(self, conversation: HarnessConversation) -> None: + if conversation.id in self.conversations: + raise ValueError(f"Conversation already exists: {conversation.id}") + self.conversations[conversation.id] = conversation.model_copy(deep=True) + + async def get_conversation( + self, conversation_id: str + ) -> HarnessConversation | None: + conversation = self.conversations.get(conversation_id) + return conversation.model_copy(deep=True) if conversation is not None else None + + async def update_conversation(self, conversation: HarnessConversation) -> None: + if conversation.id not in self.conversations: + raise ValueError(f"Conversation not found: {conversation.id}") + self.conversations[conversation.id] = conversation.model_copy(deep=True) + + async def append_event(self, event: HarnessEvent) -> None: + self.events[event.conversation_id].append(event.model_copy(deep=True)) + conversation = self.conversations.get(event.conversation_id) + if ( + conversation is not None + and event.created_datetime > conversation.updated_datetime + ): + self.conversations[event.conversation_id] = conversation.model_copy( + update={"updated_datetime": event.created_datetime} + ) + + async def iter_events(self, conversation_id: str) -> AsyncIterator[HarnessEvent]: + for event in self.events.get(conversation_id, ()): + yield event.model_copy(deep=True) + + async def remember(self, memory: HarnessMemory) -> None: + self.memories[memory.id] = memory.model_copy(deep=True) + + async def recall( + self, + *, + scope: str, + query: str, + limit: int = 20, + ) -> list[HarnessMemory]: + needle = query.casefold().strip() + values = [ + memory + for memory in self.memories.values() + if memory.scope == scope + and (not needle or needle in memory.text.casefold()) + ] + return [memory.model_copy(deep=True) for memory in values[-limit:]] + + async def forget(self, memory_id: str) -> None: + self.memories.pop(memory_id, None) diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/__init__.py b/hyperforge/src/hyperforge/harness_sdk/tools/__init__.py new file mode 100644 index 00000000..26bf49ed --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/tools/__init__.py @@ -0,0 +1,8 @@ +from .utils import ( + ContextFactory, + HarnessTool, + ToolHandler, + tool, +) + +__all__ = ["ContextFactory", "HarnessTool", "ToolHandler", "tool"] diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py new file mode 100644 index 00000000..02190640 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -0,0 +1,110 @@ +from typing import Any + +from pydantic import BaseModel + +from hyperforge.codemode import ( + RestrictedPythonTask, + SandboxRunner, + WorkerExecutionRequest, +) +from hyperforge.codemode.sandbox import settings +from hyperforge.definition import FunctionDefinition + +from . import HarnessTool, tool + +CODEMODE_TOOL_NAME = "codemode" +OUTPUT_FUNCTION_NAME = "output" + + +class CodemodeInput(BaseModel): + code: str + + +class CodemodeOutput(BaseModel): + value: Any = None + + +@tool( + name=CODEMODE_TOOL_NAME, + description=( + "Execute restricted Python code. Registered tools are available as functions; " + "call output(value) to return a result." + ), +) +async def codemode(harness: Any, input_value: CodemodeInput) -> CodemodeOutput: + tools = { + tool.name: tool + for tool in harness.iter_tools() + if tool.name != CODEMODE_TOOL_NAME + } + result = CodemodeOutput() + + async def dispatch(task: RestrictedPythonTask) -> Any: + if task.function == OUTPUT_FUNCTION_NAME: + if task.args and task.keyword_args: + raise ValueError("output accepts either a positional or keyword value") + if len(task.args) > 1: + raise ValueError("output accepts one value") + result.value = task.args[0] if task.args else task.keyword_args.get("value") + return None + + tool = tools.get(task.function) + if tool is None: + raise ValueError(f"Unknown codemode function: {task.function}") + harness.usage.tool_calls += 1 + harness._check_limit("max_tool_calls", harness.usage.tool_calls) + arguments = _tool_arguments(tool, task.args, task.keyword_args) + output = await tool.execute(harness, arguments) + return output.model_dump(mode="json") + + runner = ( + SandboxRunner.remote(settings.sandbox_socket, dispatch) + if settings.sandbox_socket is not None + else SandboxRunner.isolated_process(dispatch) + ) + definitions = { + name: FunctionDefinition( + name=name, + description=tool.description, + parameters=tool.parameters.get("properties", {}), + ) + for name, tool in tools.items() + } + definitions[OUTPUT_FUNCTION_NAME] = FunctionDefinition( + name=OUTPUT_FUNCTION_NAME, + description="Set the value returned by codemode.", + parameters={"value": {}}, + ) + await runner.run( + WorkerExecutionRequest( + code=input_value.code, + local_vars={}, + global_vars={}, + function_names={"harness": definitions}, + max_runtime_seconds=harness.usage_limits.max_codemode_runtime_seconds, + max_memory_bytes=harness.usage_limits.max_codemode_memory_bytes, + ) + ) + return result + + +def _tool_arguments( + tool: HarnessTool[Any, Any], + args: tuple[Any, ...], + keyword_args: dict[str, Any], +) -> dict[str, Any]: + if not args: + return keyword_args + names = list(tool.input_model.model_fields) + if len(args) > len(names): + raise ValueError(f"Too many positional arguments for {tool.name}") + arguments = dict(zip(names, args, strict=False)) + duplicates = arguments.keys() & keyword_args.keys() + if duplicates: + duplicate = next(iter(duplicates)) + raise ValueError(f"Multiple values for argument {duplicate!r}") + arguments.update(keyword_args) + return arguments + + +__all__ = ["CodemodeInput", "CodemodeOutput", "codemode"] diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/core.py b/hyperforge/src/hyperforge/harness_sdk/tools/core.py new file mode 100644 index 00000000..cedd1ebb --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/tools/core.py @@ -0,0 +1,230 @@ +from __future__ import annotations + +import uuid +from typing import TYPE_CHECKING, Any + +from pydantic import BaseModel + +from ..models import HarnessEventType, HarnessMemory, HarnessMessage +from . import HarnessTool, tool + +if TYPE_CHECKING: + from ..harness import AgentHarness + + +class RememberInput(BaseModel): + text: str + scope: str = "user_project" + + +class RecallInput(BaseModel): + query: str = "" + scope: str = "user_project" + + +class ForgetInput(BaseModel): + id: str + + +class SpawnAgentInput(BaseModel): + prompt: str + include_history: bool = False + + +class AgentIdInput(BaseModel): + agent_id: str + + +class SendMessageInput(AgentIdInput): + message: str + + +class CompactInput(BaseModel): + summary: str + + +class FeedbackInput(BaseModel): + question: str + + +class SearchToolsInput(BaseModel): + query: str + limit: int = 10 + + +class ActivateToolsInput(BaseModel): + names: list[str] + + +class DictOutput(BaseModel): + value: dict[str, Any] + + +class ListOutput(BaseModel): + items: list[dict[str, Any]] + + +@tool(description="Search for additional tools that can be activated.") +async def search_tools( + harness: AgentHarness, input_value: SearchToolsInput +) -> ListOutput: + terms = input_value.query.casefold().split() + candidates: list[tuple[int, HarnessTool[Any, Any]]] = [] + for candidate in harness._external_tools: + if not candidate.lazy_load or candidate.name in harness._active_lazy_tools: + continue + searchable = f"{candidate.name} {candidate.description}".casefold() + if terms and not all(term in searchable for term in terms): + continue + score = sum(2 if term in candidate.name.casefold() else 1 for term in terms) + candidates.append((score, candidate)) + candidates.sort(key=lambda item: (-item[0], item[1].name)) + limit = max(1, min(input_value.limit, 50)) + return ListOutput( + items=[ + {"name": candidate.name, "description": candidate.description} + for _, candidate in candidates[:limit] + ] + ) + + +@tool(description="Activate additional tools by their exact names.") +async def activate_tools( + harness: AgentHarness, input_value: ActivateToolsInput +) -> DictOutput: + names = list(dict.fromkeys(input_value.names)) + unknown = [ + name + for name in names + if name not in harness._tools or not harness._tools[name].lazy_load + ] + if unknown: + raise ValueError(f"Unknown lazy tools: {', '.join(unknown)}") + harness._active_lazy_tools.update(names) + return DictOutput(value={"activated": names}) + + +@tool() +async def remember(harness: AgentHarness, input_value: RememberInput) -> DictOutput: + memory = HarnessMemory( + id=uuid.uuid4().hex, + text=input_value.text, + scope=input_value.scope, + metadata=harness._persisted_metadata(), + ) + await harness.storage.remember(memory) + await harness.emit( + HarnessEventType.MEMORY_REMEMBERED, + {"memory": memory.model_dump(mode="json")}, + ) + return DictOutput(value={"id": memory.id}) + + +@tool() +async def recall(harness: AgentHarness, input_value: RecallInput) -> ListOutput: + memories = await harness.storage.recall( + scope=input_value.scope, + query=input_value.query, + ) + return ListOutput(items=[memory.model_dump(mode="json") for memory in memories]) + + +@tool() +async def forget(harness: AgentHarness, input_value: ForgetInput) -> DictOutput: + await harness.storage.forget(input_value.id) + await harness.emit(HarnessEventType.MEMORY_FORGOTTEN, {"id": input_value.id}) + return DictOutput(value={"id": input_value.id}) + + +@tool() +async def spawn_agent( + harness: AgentHarness, input_value: SpawnAgentInput +) -> DictOutput: + return await harness._child_agents.spawn(input_value) + + +@tool() +async def send_message( + harness: AgentHarness, input_value: SendMessageInput +) -> DictOutput: + return await harness._child_agents.send_message(input_value) + + +@tool() +async def wait_agent(harness: AgentHarness, input_value: AgentIdInput) -> DictOutput: + return await harness._child_agents.wait(input_value.agent_id) + + +@tool() +async def compact(harness: AgentHarness, input_value: CompactInput) -> DictOutput: + harness.messages = [ + HarnessMessage(role="system", content=harness.system_prompt), + HarnessMessage( + role="user", content=f"Conversation summary:\n{input_value.summary}" + ), + ] + await harness.emit( + HarnessEventType.COMPACTED, + {"messages": [message.model_dump(mode="json") for message in harness.messages]}, + ) + return DictOutput(value={"status": "compacted"}) + + +@tool() +async def feedback(harness: AgentHarness, input_value: FeedbackInput) -> DictOutput: + response = await harness.request_feedback( + question=input_value.question, + response_schema={ + "type": "object", + "properties": {"response": {"type": "string"}}, + }, + timeout_ms=5 * 60 * 1000, + ) + return DictOutput(value={"response": response}) + + +def create_core_tools(*, feedback_enabled: bool) -> dict[str, HarnessTool[Any, Any]]: + tools = { + item.name: item + for item in ( + remember, + recall, + forget, + spawn_agent, + send_message, + wait_agent, + compact, + search_tools, + activate_tools, + ) + } + if feedback_enabled: + tools[feedback.name] = feedback + return tools + + +__all__ = [ + "ActivateToolsInput", + "AgentIdInput", + "CompactInput", + "DictOutput", + "FeedbackInput", + "ForgetInput", + "ListOutput", + "RecallInput", + "RememberInput", + "SearchToolsInput", + "SendMessageInput", + "SpawnAgentInput", + "activate_tools", + "compact", + "create_core_tools", + "feedback", + "forget", + "recall", + "remember", + "search_tools", + "send_message", + "spawn_agent", + "wait_agent", +] diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/utils.py b/hyperforge/src/hyperforge/harness_sdk/tools/utils.py new file mode 100644 index 00000000..1cb158d1 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/tools/utils.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import inspect +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Awaitable, Callable, cast, get_type_hints + +from jsonschema import ValidationError as JsonSchemaValidationError +from jsonschema import validate as validate_json_schema +from pydantic import BaseModel, ValidationError + +from ..context import make_context +from ..models import HarnessContextReference, HarnessContextType +from ..schema import flatten_json_schema + +if TYPE_CHECKING: + from ..harness import AgentHarness + +type ToolHandler[InputT: BaseModel, OutputT: BaseModel] = Callable[ + ["AgentHarness", InputT], Awaitable[OutputT] +] +type ContextFactory[OutputT: BaseModel] = Callable[[OutputT], HarnessContextReference] + + +@dataclass(frozen=True) +class HarnessTool[InputT: BaseModel, OutputT: BaseModel]: + name: str + handler: ToolHandler[InputT, OutputT] + description: str = "" + context_type: HarnessContextType = HarnessContextType.TOOL_RESULT + parameters_schema: dict[str, Any] | None = None + context_factory: ContextFactory[OutputT] | None = None + lazy_load: bool = False + input_model: type[InputT] = field(init=False) + output_model: type[OutputT] = field(init=False) + _parameters: dict[str, Any] = field(init=False, repr=False) + + def __post_init__(self) -> None: + parameters = list(inspect.signature(self.handler).parameters.values()) + if len(parameters) != 2: + raise TypeError( + f"Tool handler {self.name!r} must accept harness and input parameters" + ) + globalns = dict(getattr(self.handler, "__globals__", {})) + globalns.setdefault("AgentHarness", object) + hints = get_type_hints(self.handler, globalns=globalns) + input_model = hints.get(parameters[1].name) + output_model = hints.get("return") + if not _is_model_type(input_model): + raise TypeError( + f"Tool handler {self.name!r} input must be annotated with a Pydantic model" + ) + if not _is_model_type(output_model): + raise TypeError( + f"Tool handler {self.name!r} return must be annotated with a Pydantic model" + ) + input_model = cast(type[InputT], input_model) + output_model = cast(type[OutputT], output_model) + object.__setattr__(self, "input_model", input_model) + object.__setattr__(self, "output_model", output_model) + schema = self.parameters_schema or input_model.model_json_schema() + object.__setattr__(self, "_parameters", flatten_json_schema(schema)) + + @property + def parameters(self) -> dict[str, Any]: + return self._parameters + + def __call__( + self, harness: AgentHarness, input_value: InputT + ) -> Awaitable[OutputT]: + return self.handler(harness, input_value) + + async def execute( + self, harness: AgentHarness, arguments: dict[str, Any] + ) -> OutputT: + if error := arguments.get("_tool_error"): + raise ValueError(str(error)) + if self.parameters_schema is not None: + try: + validate_json_schema(arguments, self.parameters_schema) + except JsonSchemaValidationError as exc: + raise ValueError( + f"Invalid {self.name} arguments: {exc.message}" + ) from exc + try: + value = self.input_model.model_validate(arguments) + except ValidationError as exc: + raise ValueError(f"Invalid {self.name} arguments: {exc}") from exc + result = await self.handler(harness, value) + return self.output_model.model_validate(result) + + def context(self, output: OutputT) -> HarnessContextReference: + if self.context_factory is not None: + return self.context_factory(output) + return make_context(self.context_type, output) + + +def tool( + name: str | None = None, + *, + description: str = "", + context_type: HarnessContextType = HarnessContextType.TOOL_RESULT, + parameters_schema: dict[str, Any] | None = None, + context_factory: ContextFactory[Any] | None = None, + lazy_load: bool = False, +) -> Callable[[ToolHandler[Any, Any]], HarnessTool[Any, Any]]: + """Create a harness tool from an annotated async handler.""" + + def decorate(handler: ToolHandler[Any, Any]) -> HarnessTool[Any, Any]: + return HarnessTool( + name=name or cast(Any, handler).__name__, + handler=handler, + description=description, + context_type=context_type, + parameters_schema=parameters_schema, + context_factory=context_factory, + lazy_load=lazy_load, + ) + + return decorate + + +def _is_model_type(value: Any) -> bool: + return isinstance(value, type) and issubclass(value, BaseModel) + + +__all__ = ["ContextFactory", "HarnessTool", "ToolHandler", "tool"] diff --git a/hyperforge/src/hyperforge/harness_sdk/usage.py b/hyperforge/src/hyperforge/harness_sdk/usage.py new file mode 100644 index 00000000..a65966ef --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/usage.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from dataclasses import dataclass + +DEFAULT_CODEMODE_RUNTIME_SECONDS = 30.0 +DEFAULT_CODEMODE_MEMORY_BYTES = 512 * 1024 * 1024 + + +@dataclass(frozen=True) +class UsageLimits: + max_tool_calls: int | None = None + max_turns: int | None = None + max_input_tokens: int | None = None + max_output_tokens: int | None = None + max_time: float | None = None + max_codemode_runtime_seconds: float = DEFAULT_CODEMODE_RUNTIME_SECONDS + max_codemode_memory_bytes: int = DEFAULT_CODEMODE_MEMORY_BYTES + + def __post_init__(self) -> None: + for name, value in vars(self).items(): + if value is not None and value <= 0: + raise ValueError(f"{name} must be greater than zero") + + +@dataclass +class HarnessUsage: + tool_calls: int = 0 + turns: int = 0 + input_tokens: float = 0 + output_tokens: float = 0 + + +class UsageLimitExceeded(RuntimeError): + pass diff --git a/hyperforge/src/hyperforge/llm.py b/hyperforge/src/hyperforge/llm.py index 05aa5c9d..00202b9b 100644 --- a/hyperforge/src/hyperforge/llm.py +++ b/hyperforge/src/hyperforge/llm.py @@ -1,6 +1,10 @@ -from typing import Any +import json +from collections.abc import AsyncIterator +from typing import Any, Self -from nuclia.lib.nua import AsyncNuaClient +import httpx +from nuclia.lib.nua import AsyncNuaClient as NucliaAsyncNuaClient +from nuclia.lib.nua import NuaEndpoint from nuclia.sdk import AsyncNucliaAuth from pydantic import BaseModel @@ -14,6 +18,60 @@ async def connect_internal(cls, kbid: str | None, account: str | None, url: str) raise NotImplementedError("Must implement connect_internal method in subclass") +class AsyncNuaClient(NucliaAsyncNuaClient): + """Hyperforge NUA client, including the chat-completions compatibility API.""" + + @classmethod + def internal( + cls, + url: str, + kbid: str | None = None, + account: str | None = None, + headers: dict[str, str] | None = None, + ) -> Self: + return cls( + region=url, + account=account or "", + headers=headers, + endpoint=NuaEndpoint.INTERNAL, + kbid=kbid, + ) + + async def chat_completions_stream( + self, payload: dict[str, Any], *, timeout: float = 5 * 60 + ) -> AsyncIterator[dict[str, Any]]: + path = ( + "/api/internal/predict/compat/chat/completions" + if self.endpoint == NuaEndpoint.INTERNAL + else "/api/v1/predict/compat/chat/completions" + ) + request = {**payload, "stream": True} + async with self.stream_client.stream( + "POST", + f"{self.url}{path}", + json=request, + headers={"accept": "text/event-stream"}, + timeout=timeout, + ) as response: + if response.status_code != 200: + detail = (await response.aread()).decode(errors="replace") + error = httpx.HTTPStatusError( + f"Nuclia chat completions API error: {response.status_code} - {detail}", + request=response.request, + response=response, + ) + raise error + async for line in response.aiter_lines(): + line = line.strip() + if not line or line.startswith(":") or line.startswith("event:"): + continue + if line == "data: [DONE]": + return + if line.startswith("data:"): + line = line[5:].lstrip() + yield json.loads(line) + + class NoopNuaClient(AsyncNuaClient): # pragma: no cover """ A no-op NUA client used when no LLM backend is configured. @@ -60,16 +118,25 @@ async def tokens(self, *args: Any, **kwargs: Any) -> Any: # type: ignore[overri async def query(self, *args: Any, **kwargs: Any) -> Any: # type: ignore[override] self._not_configured("query") + async def chat_completions_stream( + self, payload: dict[str, Any], *, timeout: float = 5 * 60 + ) -> AsyncIterator[dict[str, Any]]: + self._not_configured("chat_completions_stream") + if False: + yield {} + class NUAConnection(NuaBaseModel): key: str - async def connect(self): + async def connect(self, *, base_url: str | None = None): na = AsyncNucliaAuth() client_id, account_type, account, base_region = await na.validate_nua(self.key) if account is None or base_region is None: raise Exception("Could not connect to NUA") - return AsyncNuaClient(token=self.key, account=account, region=base_region) + return AsyncNuaClient( + token=self.key, account=account, region=base_url or base_region + ) @classmethod async def connect_internal(cls, kbid: str | None, account: str | None, url: str): diff --git a/hyperforge/src/hyperforge/server/sandbox.py b/hyperforge/src/hyperforge/server/sandbox.py index d9fd9ad1..bb7880d7 100644 --- a/hyperforge/src/hyperforge/server/sandbox.py +++ b/hyperforge/src/hyperforge/server/sandbox.py @@ -2,12 +2,12 @@ from importlib.metadata import version import sentry_sdk -from hyperforge_restricted import sandbox from nucliadb_telemetry.fastapi import application_metrics from nucliadb_telemetry.logs import setup_logging from nucliadb_telemetry.settings import LogLevel, LogSettings from sentry_sdk.integrations.excepthook import ExcepthookIntegration +from hyperforge.codemode import sandbox from hyperforge.server.settings import Settings @@ -32,6 +32,12 @@ async def run_metrics_server(port: int): await server.serve() +async def run_servers(metrics_port: int) -> None: + async with asyncio.TaskGroup() as tasks: + tasks.create_task(run_metrics_server(metrics_port)) + tasks.create_task(sandbox.run_sandbox_server()) + + def run(): # pragma: no cover settings = Settings() sandbox_settings = sandbox.SandboxSettings() @@ -49,12 +55,8 @@ def run(): # pragma: no cover ) ) - loop = asyncio.get_event_loop() - - loop.create_task(run_metrics_server(sandbox_settings.sandbox_metrics_port)) - loop.create_task(sandbox.run_sandbox_server()) - try: - loop.run_forever() - finally: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.close() + if sandbox_settings.sandbox_socket is None: + raise RuntimeError("SANDBOX_SOCKET is required for the sandbox server") + if sandbox_settings.sandbox_token is None: + raise RuntimeError("SANDBOX_TOKEN is required for the sandbox server") + asyncio.run(run_servers(sandbox_settings.sandbox_metrics_port)) diff --git a/hyperforge/tests/codemode/test_model.py b/hyperforge/tests/codemode/test_model.py new file mode 100644 index 00000000..92f86389 --- /dev/null +++ b/hyperforge/tests/codemode/test_model.py @@ -0,0 +1,29 @@ +from hyperforge.codemode import ( + SandboxMessage, + WorkerExecutionRequest, + deserialize, + serialize, +) + + +def test_nested_json_values_round_trip() -> None: + value = {"score": 0.75, "matches": [{"id": 1}, {"id": 2}]} + + assert deserialize(serialize(value)) == value + + +def test_sandbox_run_token_round_trip() -> None: + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="", + local_vars={}, + global_vars={}, + function_names={}, + ), + token="secret", + ) + + parsed = SandboxMessage.parse(SandboxMessage.serialize(message)) + + assert isinstance(parsed, SandboxMessage.Run) + assert parsed.token == "secret" diff --git a/hyperforge/tests/codemode/test_protocol.py b/hyperforge/tests/codemode/test_protocol.py new file mode 100644 index 00000000..28c75b82 --- /dev/null +++ b/hyperforge/tests/codemode/test_protocol.py @@ -0,0 +1,21 @@ +import asyncio + +import pytest + +from hyperforge.codemode.model import deserialize +from hyperforge.codemode.sandbox import MAX_PACKET_BYTES, SandboxReader + + +def test_deserialize_preserves_unknown_model_marker() -> None: + value = {"__model__": "customer-data", "value": 1} + + assert deserialize(value) == value + + +@pytest.mark.asyncio +async def test_reader_rejects_oversized_packet_before_payload() -> None: + reader = asyncio.StreamReader() + reader.feed_data((MAX_PACKET_BYTES + 1).to_bytes(4, "little")) + + with pytest.raises(ValueError, match="maximum size"): + await SandboxReader(reader)._read_packet() diff --git a/hyperforge/tests/harness/test_harness_agent.py b/hyperforge/tests/harness/test_harness_agent.py new file mode 100644 index 00000000..7ed14fed --- /dev/null +++ b/hyperforge/tests/harness/test_harness_agent.py @@ -0,0 +1,122 @@ +from collections.abc import AsyncIterator +from typing import ClassVar + +import pytest + +from hyperforge.agent import Agent, AgentConfig +from hyperforge.configure import GLOBAL_REGISTRY, AgentRegistry +from hyperforge.definition import FunctionDefinition +from hyperforge.harness import HarnessAgent, HarnessAgentConfig +from hyperforge.manager import Manager +from hyperforge.memory.memory import BaseSessionMemory +from hyperforge.models import MemoryConfig, Rules + + +class PublishedConfig(AgentConfig): + module: str = "published" + + +class PublishedAgent(Agent[PublishedConfig]): + __published_functions__: ClassVar[dict[str, FunctionDefinition]] = { + "lookup": FunctionDefinition( + name="lookup", + description="Look up a value.", + parameters={"question": {"type": "string"}}, + ) + } + + async def inner_from_config(self, config, agent_id=None): + pass + + async def lookup(self, question: str, memory=None) -> str: + suffix = memory.headers.get("request-header", "") + return f"{question.upper()}{suffix}" + + +class FakeNua: + def __init__(self) -> None: + self.calls = 0 + + async def chat_completions_stream(self, payload, **kwargs) -> AsyncIterator[dict]: + self.calls += 1 + assert any( + message["role"] == "system" and "- Be concise" in message["content"] + for message in payload["messages"] + ) + if self.calls == 1: + tool = next( + item + for item in payload["tools"] + if item["function"]["name"].endswith("__lookup") + ) + yield { + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call-1", + "function": { + "name": tool["function"]["name"], + "arguments": '{"question":"hello"}', + }, + } + ] + }, + "finish_reason": "tool_calls", + } + ] + } + else: + assert any( + message["role"] == "tool" and "HELLO" in message["content"] + for message in payload["messages"] + ) + yield {"choices": [{"index": 0, "delta": {"content": "Final answer"}}]} + + +@pytest.mark.asyncio +async def test_harness_agent_runs_configured_legacy_agent_as_tool(monkeypatch) -> None: + monkeypatch.setitem( + GLOBAL_REGISTRY.agents, + "published", + AgentRegistry( + id="published", + agent_type="context", + title="Published", + description="Published test agent", + config_schema=PublishedConfig, + klass=PublishedAgent, + ), + ) + nua = FakeNua() + config = HarnessAgentConfig.model_validate( + { + "model": "test-model", + "agents": [{"id": "source", "module": "published"}], + "workflow": { + "id": "default", + "name": "Default", + "description": None, + "parameters": None, + }, + } + ) + agent = await HarnessAgent.from_config_class(config) + manager = Manager() + manager.nua = nua + session = BaseSessionMemory.from_config( + MemoryConfig(), + agent_id="test", + workflow_id="default", + rules=Rules(rules=["Be concise"]), + ) + session.init("session") + memory = session.start_question("Answer me", headers={"request-header": "!"}) + + await agent(memory, manager) + + assert memory.final_answer == "Final answer" + assert nua.calls == 2 diff --git a/hyperforge/tests/harness_sdk/test_clients.py b/hyperforge/tests/harness_sdk/test_clients.py new file mode 100644 index 00000000..5caab3f8 --- /dev/null +++ b/hyperforge/tests/harness_sdk/test_clients.py @@ -0,0 +1,241 @@ +from collections.abc import AsyncIterator + +import pytest + +from hyperforge.harness_sdk import ( + ChatCompletionChunk, + ChatCompletionRequest, + HarnessMessage, + HarnessToolCall, + NucliaChatCompletionsClient, + NucliaModelClient, +) + + +class FakeResponse: + status_code = 200 + + async def aiter_lines(self): + yield 'data: {"id":"chunk","model":"model","choices":[{"index":0,"delta":{"content":"hello"}}]}' + yield "data: [DONE]" + + +class ErrorResponse: + status_code = 200 + + async def aiter_lines(self): + yield 'data: {"type":"status","code":"ERROR","details":"transient generation failure","request_id":"req-1"}' + + +@pytest.mark.asyncio +async def test_chat_completions_transport_parses_sse() -> None: + class FakeNua: + async def chat_completions_stream(self, payload, **kwargs): + assert payload["stream"] is True + assert payload["stream_options"] == {"include_usage": True} + yield { + "id": "chunk", + "model": "model", + "choices": [{"index": 0, "delta": {"content": "hello"}}], + } + + client = NucliaChatCompletionsClient(FakeNua()) + + chunks = [ + chunk + async for chunk in client.stream( + ChatCompletionRequest(messages=[{"role": "user", "content": "hi"}]) + ) + ] + + assert chunks[0].choices[0].delta.content == "hello" + + +@pytest.mark.asyncio +async def test_chat_completions_retries_pre_stream_generation_errors() -> None: + class FakeNua: + def __init__(self): + self.attempts = 0 + + async def chat_completions_stream(self, payload, **kwargs): + self.attempts += 1 + if self.attempts < 3: + yield { + "type": "status", + "code": "ERROR", + "details": "transient", + "request_id": "req-1", + } + else: + yield { + "id": "chunk", + "model": "model", + "choices": [{"index": 0, "delta": {"content": "hello"}}], + } + + nua = FakeNua() + client = NucliaChatCompletionsClient(nua) + + chunks = [ + chunk + async for chunk in client.stream( + ChatCompletionRequest( + messages=[{"role": "user", "content": "hi"}], model="model" + ) + ) + ] + + assert nua.attempts == 3 + assert [chunk.choices[0].delta.content for chunk in chunks] == ["hello"] + + +@pytest.mark.asyncio +async def test_harness_client_keeps_usage_only_chunk() -> None: + class UsageClient: + async def stream(self, request): + yield ChatCompletionChunk.model_validate( + { + "id": "usage", + "model": "model", + "choices": [], + "usage": {"prompt_tokens": 11, "completion_tokens": 5}, + } + ) + + deltas = [ + delta + async for delta in NucliaModelClient(UsageClient()).stream( + model="model", + reasoning_effort=None, + messages=[HarnessMessage(role="user", content="hi")], + tools=[], + execution_context={}, + ) + ] + assert (deltas[0].input_tokens, deltas[0].output_tokens) == (11, 5) + + +class FakeCompletionsClient: + async def stream( + self, request: ChatCompletionRequest + ) -> AsyncIterator[ChatCompletionChunk]: + assert request.messages[0] == {"role": "system", "content": "instructions"} + assert request.reasoning_effort == "high" + yield ChatCompletionChunk.model_validate( + { + "id": "chunk", + "model": "model", + "choices": [ + { + "index": 0, + "delta": { + "reasoning_content": "thinking", + "tool_calls": [ + { + "index": 0, + "id": "call-1", + "function": { + "name": "lookup", + "arguments": '{"query":', + }, + } + ], + }, + } + ], + } + ) + yield ChatCompletionChunk.model_validate( + { + "id": "chunk", + "model": "model", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": '"x"}'}} + ] + }, + "finish_reason": "tool_calls", + } + ], + "usage": { + "prompt_tokens": 7, + "completion_tokens": 3, + "total_tokens": 10, + }, + } + ) + + async def aclose(self) -> None: + pass + + +@pytest.mark.asyncio +async def test_harness_client_assembles_openai_tool_deltas() -> None: + client = NucliaModelClient(FakeCompletionsClient()) + deltas = [ + delta + async for delta in client.stream( + model="model", + reasoning_effort="high", + messages=[HarnessMessage(role="system", content="instructions")], + tools=[], + execution_context={"user_id": "user"}, + ) + ] + + assert deltas[0].reasoning == "thinking" + assert deltas[-1].tool_calls == [ + HarnessToolCall(id="call-1", name="lookup", arguments={"query": "x"}) + ] + assert deltas[-1].input_tokens == 7 + assert deltas[-1].output_tokens == 3 + + +def test_harness_client_preserves_structured_tool_history() -> None: + messages = NucliaModelClient._messages( + [ + HarnessMessage( + role="assistant", + content="", + tool_calls=[ + HarnessToolCall( + id="call-1", name="lookup", arguments={"query": "x"} + ) + ], + ), + HarnessMessage( + role="tool", content="result", tool_call_id="call-1", tool_name="lookup" + ), + ] + ) + + assert messages[0]["tool_calls"][0]["function"]["name"] == "lookup" + assert messages[1] == { + "role": "tool", + "content": "result", + "tool_call_id": "call-1", + } + NucliaModelClient._validate_tool_history(messages) + + +def test_harness_client_rejects_unmatched_tool_history() -> None: + messages = [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + {"role": "user", "content": "continue"}, + ] + + with pytest.raises(ValueError, match="missing outputs.*call-1"): + NucliaModelClient._validate_tool_history(messages) diff --git a/hyperforge/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py new file mode 100644 index 00000000..c1acabce --- /dev/null +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -0,0 +1,182 @@ +import sys +from collections.abc import AsyncIterator +from contextvars import ContextVar + +import pytest +from pydantic import BaseModel + +from hyperforge.harness_sdk import ( + AgentHarness, + CodemodeInput, + ModelDelta, + UsageLimits, + codemode, + tool, +) + + +class UpperInput(BaseModel): + value: str + + +class UpperOutput(BaseModel): + value: str + + +@tool(description="Uppercase a value") +async def upper(_: AgentHarness, input_value: UpperInput) -> UpperOutput: + return UpperOutput(value=input_value.value.upper()) + + +class UnusedModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + del kwargs + yield ModelDelta(text="unused") + + +@pytest.mark.asyncio +async def test_codemode_calls_registered_tools_and_returns_output() -> None: + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[upper], + disabled_core_tools=[ + "remember", + "recall", + "forget", + "spawn_agent", + "send_message", + "wait_agent", + "compact", + "search_tools", + "activate_tools", + ], + ) + + result = await codemode.execute( + harness, + CodemodeInput( + code="result = upper(value='hello')\noutput(result['value'])" + ).model_dump(), + ) + + assert result.value == "HELLO" + + +@pytest.mark.asyncio +async def test_codemode_counts_nested_tool_calls() -> None: + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[upper], + usage_limits=UsageLimits(max_tool_calls=1), + ) + + with pytest.raises(RuntimeError, match="max_tool_calls"): + await codemode.execute( + harness, + CodemodeInput(code="upper(value='one')\nupper(value='two')").model_dump(), + ) + + +@pytest.mark.asyncio +async def test_codemode_propagates_tool_validation_errors() -> None: + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[upper], + ) + + with pytest.raises(RuntimeError, match="Invalid upper arguments"): + await codemode.execute( + harness, + CodemodeInput(code="upper(missing='value')").model_dump(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("code", "expected"), + [ + ("output(sum(i * i for i in range(1, 101)))", 338350), + ( + "total = 0\nfor i in range(1, 101):\n total += i * i\noutput(total)", + 338350, + ), + ], +) +async def test_codemode_supports_common_aggregation_patterns( + code: str, expected: int +) -> None: + harness = AgentHarness(model="test", model_client=UnusedModel()) + + result = await codemode.execute(harness, CodemodeInput(code=code).model_dump()) + + assert result.value == expected + + +@pytest.mark.asyncio +async def test_codemode_preserves_context_when_calling_tools() -> None: + request_context = ContextVar("request_context", default="missing") + + @tool() + async def read_context(_: AgentHarness, _input_value: UpperInput) -> UpperOutput: + return UpperOutput(value=request_context.get()) + + harness = AgentHarness( + model="test", model_client=UnusedModel(), tools=[read_context] + ) + token = request_context.set("available") + try: + result = await codemode.execute( + harness, + CodemodeInput( + code="result = read_context(value='unused')\noutput(result['value'])" + ).model_dump(), + ) + finally: + request_context.reset(token) + + assert result.value == "available" + + +@pytest.mark.asyncio +async def test_codemode_enforces_runtime_limit() -> None: + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + usage_limits=UsageLimits(max_codemode_runtime_seconds=0.1), + ) + + with pytest.raises(RuntimeError, match="timed out"): + await codemode.execute( + harness, + CodemodeInput(code="while True: pass").model_dump(), + ) + + +@pytest.mark.asyncio +async def test_codemode_blocks_process_control_exceptions() -> None: + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(RuntimeError, match="SystemExit.*not defined"): + await codemode.execute( + harness, + CodemodeInput(code="raise SystemExit(1)").model_dump(), + ) + + +@pytest.mark.asyncio +@pytest.mark.skipif(sys.platform == "darwin", reason="RLIMIT_AS is unreliable on macOS") +async def test_codemode_enforces_memory_limit() -> None: + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + usage_limits=UsageLimits(max_codemode_memory_bytes=512 * 1024 * 1024), + ) + + with pytest.raises(RuntimeError): + await codemode.execute( + harness, + CodemodeInput(code="output('x' * (1024 * 1024 * 1024))").model_dump(), + ) diff --git a/hyperforge/tests/harness_sdk/test_harness.py b/hyperforge/tests/harness_sdk/test_harness.py new file mode 100644 index 00000000..5cc6dfd4 --- /dev/null +++ b/hyperforge/tests/harness_sdk/test_harness.py @@ -0,0 +1,1525 @@ +import asyncio +from collections.abc import AsyncIterator +from typing import ClassVar + +import pytest +from nuclia_models.predict.generative_responses import GenerativeFullResponse +from pydantic import BaseModel, Field + +from hyperforge.definition import FunctionDefinition +from hyperforge.harness_sdk import ( + AgentHarness, + HarnessEvent, + HarnessEventType, + HarnessMessage, + HarnessTool, + HarnessToolCall, + InMemoryHarnessStorage, + LLMCallError, + ModelDelta, + NucliaModelClient, + UsageLimitExceeded, + UsageLimits, + codemode, + tool, +) +from hyperforge.harness_sdk.harness import ( + EMPTY_RESPONSE_RETRY_PROMPT, + AgentResult, + TurnLoopState, +) +from hyperforge.harness_sdk.tools.core import ( + AgentIdInput, + SpawnAgentInput, + spawn_agent, + wait_agent, +) +from hyperforge.manager import Manager + + +class ToolInput(BaseModel): + value: str + + +class ToolOutput(BaseModel): + value: str + + +class Model: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + del kwargs + self.calls += 1 + if self.calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="call-1", name="upper", arguments={"value": "hi"} + ) + ] + ) + else: + yield ModelDelta(text="HI", input_tokens=4, output_tokens=1) + + +async def run_harness(harness: AgentHarness, prompt: str) -> str: + result = "interrupted" + async for event in harness.run(prompt): + if ( + event.agent_id == harness.agent_id + and event.type == HarnessEventType.TURN_COMPLETED + ): + result = str(event.payload.get("text", "")) + return result + + +def test_codemode_is_not_available_by_default() -> None: + harness = AgentHarness(model="test-model", model_client=Model()) + + assert "codemode" not in {tool.name for tool in harness.iter_tools()} + + +def test_codemode_can_be_registered_explicitly() -> None: + harness = AgentHarness( + model="test-model", + model_client=Model(), + tools=[codemode], + ) + + assert "codemode" in {tool.name for tool in harness.iter_tools()} + + +def test_turn_loop_clears_pending_tool_result_after_non_empty_response() -> None: + state = TurnLoopState() + state.record_tool_calls() + + assert state.retry_prompt(AgentResult(text="Answer")) is None + assert state.retry_prompt(AgentResult(text="")) == EMPTY_RESPONSE_RETRY_PROMPT + + +@pytest.mark.asyncio +async def test_llm_wraps_model_failure_with_call_context() -> None: + error = ValueError("provider failed") + + class FailingModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + yield ModelDelta(trace_id="trace-1", model="resolved-model") + raise error + + harness = AgentHarness(model="requested-model", model_client=FailingModel()) + + with pytest.raises(LLMCallError) as captured: + await harness.llm() + + assert captured.value.cause is error + assert captured.value.__cause__ is error + assert captured.value.call_id + assert captured.value.trace_id == "trace-1" + assert captured.value.model == "resolved-model" + + +@pytest.mark.asyncio +async def test_run_preserves_wrapped_model_failure_context(caplog) -> None: + class FailingModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + yield ModelDelta(trace_id="trace-1", model="resolved-model") + raise ValueError("provider failed") + + harness = AgentHarness(model="requested-model", model_client=FailingModel()) + + with pytest.raises(LLMCallError): + await run_harness(harness, "fail") + + record = next( + record + for record in caplog.records + if record.message.startswith("Agent turn failed") + ) + assert record.call_id + assert record.trace_id == "trace-1" + assert record.model == "resolved-model" + assert record.error_type == "ValueError" + assert record.error == "provider failed" + + +@pytest.mark.asyncio +async def test_compact_preserves_tool_call_pair_for_next_generation() -> None: + class CompactingModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + if self.calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="compact-1", + name="compact", + arguments={ + "summary": "The user asked about quarterly filings." + }, + ) + ] + ) + return + + messages = kwargs["messages"] + assert [message.role for message in messages] == [ + "system", + "user", + "assistant", + "tool", + ] + assert messages[-2].tool_calls[0].id == "compact-1" + assert messages[-1].tool_call_id == "compact-1" + assert messages[-1].tool_name == "compact" + yield ModelDelta(text="Compaction succeeded") + + harness = AgentHarness(model="test-model", model_client=CompactingModel()) + + assert await run_harness(harness, "Find the filings") == "Compaction succeeded" + + +@pytest.mark.asyncio +async def test_compact_twice_preserves_only_current_tool_call_pair() -> None: + class CompactingModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + if self.calls <= 2: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id=f"compact-{self.calls}", + name="compact", + arguments={"summary": f"Summary {self.calls}"}, + ) + ] + ) + return + messages = kwargs["messages"] + assert [message.role for message in messages] == [ + "system", + "user", + "assistant", + "tool", + ] + assert messages[-2].tool_calls[0].id == "compact-2" + assert messages[-1].tool_call_id == "compact-2" + yield ModelDelta(text="Compaction succeeded") + + harness = AgentHarness(model="test-model", model_client=CompactingModel()) + + assert await run_harness(harness, "Find the filings") == "Compaction succeeded" + + +@pytest.mark.asyncio +async def test_agent_loop_streams_direct_answer_and_records_usage() -> None: + class StreamingModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + assert [message.role for message in kwargs["messages"]] == [ + "system", + "user", + ] + yield ModelDelta(text="Hello") + yield ModelDelta( + text=" world", + input_tokens=5, + output_tokens=2, + model="resolved-model", + ) + + harness = AgentHarness(model="requested-model", model_client=StreamingModel()) + events = [] + + async for event in harness.run("Say hello"): + events.append(event) + + assert [ + event.payload["text"] + for event in events + if event.type == HarnessEventType.TEXT_DELTA + ] == [ + "Hello", + " world", + ] + assert events[-1].type == HarnessEventType.TURN_COMPLETED + assert harness.usage.turns == 1 + assert harness.usage.tool_calls == 0 + assert harness.usage.input_tokens == 5 + assert harness.usage.output_tokens == 2 + + +@pytest.mark.asyncio +async def test_agent_reasoning_effort_is_passed_to_model_and_children() -> None: + class ReasoningModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + assert kwargs["reasoning_effort"] == "high" + yield ModelDelta(text="Answer") + + harness = AgentHarness( + model="test-model", + model_client=ReasoningModel(), + reasoning_effort="high", + ) + + assert await run_harness(harness, "Think carefully") == "Answer" + assert ( + harness._create_child("child", include_history=False).reasoning_effort == "high" + ) + + +@pytest.mark.asyncio +async def test_iter_events_discards_events_from_unconsumed_previous_turn() -> None: + class DirectModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + yield ModelDelta(text=f"answer-{self.calls}") + + harness = AgentHarness(model="test-model", model_client=DirectModel()) + assert await run_harness(harness, "first") == "answer-1" + events = [] + + async for event in harness.run("second"): + events.append(event) + + assert {event.turn_id for event in events} == {harness._turn_id} + assert events[-1].payload == {"text": "answer-2"} + + +@pytest.mark.asyncio +async def test_unconsumed_turn_does_not_buffer_live_events() -> None: + harness = AgentHarness(model="test-model", model_client=Model()) + + assert await run_harness(harness, "first") == "HI" + + assert harness._event_queue is None + + +@pytest.mark.asyncio +async def test_slow_consumer_drops_deltas_without_blocking_producer() -> None: + harness = AgentHarness(model="test-model", model_client=Model(), event_queue_size=1) + harness._event_queue = asyncio.Queue() + + await harness.emit(HarnessEventType.TEXT_DELTA, {"text": "first"}, persist=False) + await harness.emit(HarnessEventType.TEXT_DELTA, {"text": "dropped"}, persist=False) + await harness.emit(HarnessEventType.TURN_COMPLETED, {"text": "complete"}) + + assert [harness._event_queue.get_nowait().type for _ in range(2)] == [ + HarnessEventType.TEXT_DELTA, + HarnessEventType.TURN_COMPLETED, + ] + + +@pytest.mark.asyncio +async def test_abandoning_run_stream_cancels_turn() -> None: + cancelled = asyncio.Event() + + class SlowModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + yield ModelDelta(text="started") + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cancelled.set() + raise + + harness = AgentHarness(model="test-model", model_client=SlowModel()) + + async for event in harness.run("slow"): + if event.type == HarnessEventType.TEXT_DELTA: + break + + await harness.__aexit__(None, None, None) + await asyncio.sleep(0) + assert cancelled.is_set() + assert harness._event_queue is None + + +@pytest.mark.asyncio +async def test_agent_loop_executes_parallel_tools_and_synthesizes_answer() -> None: + calls = [] + + async def upper(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + calls.append(value.value) + await asyncio.sleep(0) + return ToolOutput(value=value.value.upper()) + + class ParallelToolModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + if self.calls == 1: + assert "upper" in {tool.name for tool in kwargs["tools"]} + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="call-a", name="upper", arguments={"value": "a"} + ), + HarnessToolCall( + id="call-b", name="upper", arguments={"value": "b"} + ), + ], + input_tokens=4, + output_tokens=2, + ) + return + + tool_messages = [ + message for message in kwargs["messages"] if message.role == "tool" + ] + assert [message.tool_call_id for message in tool_messages] == [ + "call-a", + "call-b", + ] + assert [message.context.content for message in tool_messages] == [ + {"value": "A"}, + {"value": "B"}, + ] + yield ModelDelta(text="A and B", input_tokens=8, output_tokens=3) + + harness = AgentHarness( + model="test-model", + model_client=ParallelToolModel(), + tools=[HarnessTool("upper", upper)], + ) + + assert await run_harness(harness, "Uppercase both values") == "A and B" + assert calls == ["a", "b"] + assert harness.usage.turns == 2 + assert harness.usage.tool_calls == 2 + assert harness.usage.input_tokens == 12 + assert harness.usage.output_tokens == 5 + assert [message.role for message in harness.messages] == [ + "system", + "user", + "assistant", + "tool", + "tool", + "assistant", + ] + + +@pytest.mark.asyncio +async def test_agent_loop_applies_live_steering_before_completion() -> None: + first_call_started = asyncio.Event() + release_first_call = asyncio.Event() + + class SteeringModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + if self.calls == 1: + first_call_started.set() + await release_first_call.wait() + yield ModelDelta(text="Draft") + return + + assert any( + "Prefer a concise answer" in message.content + for message in kwargs["messages"] + ) + yield ModelDelta(text="Concise answer") + + harness = AgentHarness(model="test-model", model_client=SteeringModel()) + run_task = asyncio.create_task(run_harness(harness, "Answer the question")) + await first_call_started.wait() + message_id = await harness.steer("Prefer a concise answer") + release_first_call.set() + + assert await run_task == "Concise answer" + assert message_id + event_types = [event.type async for event in harness.history()] + assert HarnessEventType.INBOX_ADDED in event_types + assert HarnessEventType.INBOX_CONSUMED in event_types + + +@pytest.mark.asyncio +async def test_agent_loop_feedback_round_trip() -> None: + class FeedbackModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + if self.calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="feedback-call", + name="feedback", + arguments={"question": "Which format?"}, + ) + ] + ) + return + assert any( + message.role == "tool" and "markdown" in message.content + for message in kwargs["messages"] + ) + yield ModelDelta(text="I will use markdown.") + + harness = AgentHarness( + model="test-model", model_client=FeedbackModel(), feedback_enabled=True + ) + run_task = asyncio.create_task(run_harness(harness, "Prepare the report")) + + while not harness.feedback_requests: + await asyncio.sleep(0) + request_id = next(iter(harness.feedback_requests)) + await harness.respond_feedback(request_id, "markdown") + + assert await run_task == "I will use markdown." + event_types = [event.type async for event in harness.history()] + assert HarnessEventType.FEEDBACK_REQUESTED in event_types + assert HarnessEventType.FEEDBACK_RESOLVED in event_types + + +@pytest.mark.asyncio +async def test_agent_loop_spawns_and_waits_for_sub_agent() -> None: + class DelegationModel: + def __init__(self) -> None: + self.root_calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + messages = kwargs["messages"] + if any(message.content == "research topic" for message in messages): + yield ModelDelta( + text="research result", input_tokens=2, output_tokens=2 + ) + return + + self.root_calls += 1 + if self.root_calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="spawn-call", + name="spawn_agent", + arguments={"prompt": "research topic"}, + ) + ] + ) + return + if self.root_calls == 2: + spawn_message = next( + message + for message in reversed(messages) + if message.tool_name == "spawn_agent" + ) + child_id = spawn_message.context.content["value"]["agent_id"] + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="wait-call", + name="wait_agent", + arguments={"agent_id": child_id}, + ) + ] + ) + return + + assert any( + message.tool_name == "wait_agent" + and "research result" in message.content + for message in messages + ) + yield ModelDelta(text="Final: research result") + + harness = AgentHarness(model="test-model", model_client=DelegationModel()) + + assert await run_harness(harness, "Delegate research") == "Final: research result" + assert harness.children == {} + assert len(harness.child_results) == 1 + assert harness.usage.turns == 4 + assert harness.usage.tool_calls == 2 + + +@pytest.mark.asyncio +async def test_parallel_spawn_with_history_excludes_pending_spawn_calls() -> None: + class DelegationModel: + def __init__(self) -> None: + self.root_calls = 0 + self.child_histories: list[list[HarnessMessage]] = [] + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + messages = kwargs["messages"] + if any( + message.content in {"first child", "second child"} + for message in messages + ): + self.child_histories.append(list(messages)) + assert any(message.content == "delegate" for message in messages) + yield ModelDelta(text="child result") + return + self.root_calls += 1 + if self.root_calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="spawn-1", + name="spawn_agent", + arguments={ + "prompt": "first child", + "include_history": True, + }, + ), + HarnessToolCall( + id="spawn-2", + name="spawn_agent", + arguments={ + "prompt": "second child", + "include_history": True, + }, + ), + ] + ) + return + yield ModelDelta(text="done") + + model = DelegationModel() + harness = AgentHarness(model="test-model", model_client=model) + + assert await run_harness(harness, "delegate") == "done" + assert len(model.child_histories) == 2 + for messages in model.child_histories: + NucliaModelClient._validate_tool_history(NucliaModelClient._messages(messages)) + assert not any( + call.name == "spawn_agent" + for message in messages + for call in message.tool_calls + ) + + +@pytest.mark.asyncio +async def test_spawn_agent_returns_wait_result_at_conversation_limit() -> None: + release = asyncio.Event() + + class SlowChildModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + await release.wait() + yield ModelDelta(text="done") + + harness = AgentHarness( + model="test-model", + model_client=SlowChildModel(), + max_concurrent_agents=2, + ) + first = await spawn_agent( + harness, SpawnAgentInput(prompt="first", include_history=False) + ) + first_id = first.value["agent_id"] + + blocked = await spawn_agent( + harness, SpawnAgentInput(prompt="second", include_history=False) + ) + + assert blocked.value == { + "status": "concurrency_limit_reached", + "max_concurrent_agents": 2, + "active_agents": 2, + "waitable_agent_ids": [first_id], + "message": ( + "The conversation is at its concurrent agent limit. " + "Call wait_agent for an existing child before trying to spawn another." + ), + } + + release.set() + await wait_agent(harness, AgentIdInput(agent_id=first_id)) + next_spawn = await spawn_agent( + harness, SpawnAgentInput(prompt="second", include_history=False) + ) + assert "agent_id" in next_spawn.value + await harness._stop_children() + + +@pytest.mark.asyncio +async def test_concurrent_agent_limit_is_shared_with_descendants() -> None: + release = asyncio.Event() + + class SlowChildModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + await release.wait() + yield ModelDelta(text="done") + + harness = AgentHarness( + model="test-model", + model_client=SlowChildModel(), + max_spawn_depth=2, + max_concurrent_agents=2, + ) + spawned = await spawn_agent( + harness, SpawnAgentInput(prompt="child", include_history=False) + ) + child, _ = harness.children[spawned.value["agent_id"]] + + blocked = await spawn_agent( + child, SpawnAgentInput(prompt="grandchild", include_history=False) + ) + + assert blocked.value["status"] == "concurrency_limit_reached" + assert blocked.value["active_agents"] == 2 + assert blocked.value["waitable_agent_ids"] == [] + assert "Finish the current work" in blocked.value["message"] + release.set() + await harness._stop_children() + + +@pytest.mark.asyncio +async def test_load_repairs_tool_output_returned_before_its_call() -> None: + storage = InMemoryHarnessStorage() + harness = AgentHarness( + model="test-model", + model_client=Model(), + storage=storage, + conversation_id="conversation-1", + ) + await harness.load() + tool_output = HarnessMessage( + role="tool", + content="result", + tool_call_id="call-1", + tool_name="lookup", + ) + tool_call = HarnessMessage( + role="assistant", + content="", + tool_calls=[HarnessToolCall(id="call-1", name="lookup", arguments={})], + ) + await storage.append_event( + HarnessEvent( + id="event-output", + conversation_id="conversation-1", + type=HarnessEventType.MESSAGE_ADDED, + payload={"message": tool_output.model_dump(mode="json")}, + ) + ) + await storage.append_event( + HarnessEvent( + id="event-call", + conversation_id="conversation-1", + type=HarnessEventType.MESSAGE_ADDED, + payload={"message": tool_call.model_dump(mode="json")}, + ) + ) + + await harness.load(create=False) + + assert harness.messages == [tool_call, tool_output] + NucliaModelClient._validate_tool_history( + NucliaModelClient._messages(harness.messages) + ) + + +@pytest.mark.asyncio +async def test_wait_agent_returns_stable_failure_result() -> None: + class FailingModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + raise RuntimeError("child failed") + yield + + harness = AgentHarness(model="test-model", model_client=FailingModel()) + spawned = await spawn_agent( + harness, + SpawnAgentInput(prompt="fail", include_history=False), + ) + child_id = spawned.value["agent_id"] + input_value = AgentIdInput(agent_id=child_id) + + first = await wait_agent(harness, input_value) + second = await wait_agent(harness, input_value) + + assert first == second + assert first.value == { + "agent_id": child_id, + "status": "failed", + "error": "child failed", + } + + +@pytest.mark.asyncio +async def test_wait_agent_returns_when_child_fails_before_terminal_event() -> None: + class FailingStorage(InMemoryHarnessStorage): + async def append_event(self, event: HarnessEvent) -> None: + if ( + event.parent_agent_id is not None + and event.type == HarnessEventType.TURN_STARTED + ): + raise RuntimeError("storage failed") + await super().append_event(event) + + harness = AgentHarness( + model="test-model", model_client=Model(), storage=FailingStorage() + ) + spawned = await spawn_agent( + harness, + SpawnAgentInput(prompt="fail", include_history=False), + ) + child_id = spawned.value["agent_id"] + + result = await asyncio.wait_for( + wait_agent(harness, AgentIdInput(agent_id=child_id)), timeout=1 + ) + + assert result.value == { + "agent_id": child_id, + "status": "failed", + "error": "storage failed", + } + + +@pytest.mark.asyncio +async def test_harness_runs_tool_loop_and_resumes_history() -> None: + async def upper(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value.upper()) + + storage = InMemoryHarnessStorage() + harness = AgentHarness( + model="test-model", + model_client=Model(), + tools=[HarnessTool("upper", upper)], + conversation_id="conversation", + storage=storage, + ) + + assert await run_harness(harness, "say hi") == "HI" + event_types = [event.type async for event in harness.history()] + assert HarnessEventType.TOOL_COMPLETED in event_types + assert HarnessEventType.LLM_COMPLETED in event_types + + resumed = AgentHarness( + model="test-model", + model_client=Model(), + conversation_id="conversation", + storage=storage, + ) + await resumed.load(create=False) + assert [message.role for message in resumed.messages] == [ + "system", + "user", + "assistant", + "tool", + "assistant", + ] + + +@pytest.mark.asyncio +async def test_default_storage_is_ephemeral() -> None: + harness = AgentHarness(model="test-model", model_client=Model()) + await harness.load() + assert await harness.storage.get_conversation(harness.conversation_id) is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("limits", "model"), + [ + (UsageLimits(max_turns=1), Model()), + (UsageLimits(max_input_tokens=3), Model()), + ], +) +async def test_usage_limits(limits: UsageLimits, model: Model) -> None: + async def upper(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value.upper()) + + harness = AgentHarness( + model="test-model", + model_client=model, + tools=[HarnessTool("upper", upper)], + usage_limits=limits, + ) + + with pytest.raises(UsageLimitExceeded): + await run_harness(harness, "say hi") + + +@pytest.mark.asyncio +async def test_tool_call_and_output_token_limits() -> None: + class LimitModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + del kwargs + self.calls += 1 + if self.calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall(name="upper", arguments={"value": "a"}), + HarnessToolCall(name="upper", arguments={"value": "b"}), + ] + ) + else: + yield ModelDelta(text="done", output_tokens=2) + + async def upper(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value.upper()) + + tool = HarnessTool("upper", upper) + tool_limited = AgentHarness( + model="test-model", + model_client=LimitModel(), + tools=[tool], + usage_limits=UsageLimits(max_tool_calls=1), + ) + with pytest.raises(UsageLimitExceeded): + await run_harness(tool_limited, "call twice") + + output_limited = AgentHarness( + model="test-model", + model_client=LimitModel(), + tools=[tool], + usage_limits=UsageLimits(max_output_tokens=1), + ) + with pytest.raises(UsageLimitExceeded): + await run_harness(output_limited, "answer") + + +@pytest.mark.asyncio +async def test_max_time_interrupts_active_model_work() -> None: + cancelled = asyncio.Event() + + class SlowModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + del kwargs + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cancelled.set() + raise + yield ModelDelta(text="late") + + harness = AgentHarness( + model="test-model", + model_client=SlowModel(), + usage_limits=UsageLimits(max_time=0.01), + ) + + with pytest.raises(UsageLimitExceeded): + await run_harness(harness, "wait") + assert cancelled.is_set() + + +@pytest.mark.asyncio +async def test_interrupt_cancels_active_tool_work() -> None: + started = asyncio.Event() + cancelled = asyncio.Event() + + async def slow_tool(_harness: AgentHarness, _value: ToolInput) -> ToolOutput: + started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cancelled.set() + raise + return ToolOutput(value="late") + + harness = AgentHarness( + model="test-model", + model_client=Model(), + tools=[HarnessTool("upper", slow_tool)], + ) + task = asyncio.create_task(run_harness(harness, "wait")) + await started.wait() + assert harness._active_task is not task + harness.interrupt() + + assert await task == "interrupted" + assert cancelled.is_set() + assert task.cancelling() == 0 + + +@pytest.mark.asyncio +async def test_interrupt_mid_tool_can_run_again_and_reload() -> None: + started = asyncio.Event() + + async def slow(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + started.set() + await asyncio.sleep(60) + return ToolOutput(value=value.value) + + class RecoveringModel: + def __init__(self): + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + if self.calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="slow", name="slow", arguments={"value": "x"} + ) + ] + ) + else: + NucliaModelClient._validate_tool_history( + NucliaModelClient._messages(kwargs["messages"]) + ) + yield ModelDelta(text="recovered") + + storage = InMemoryHarnessStorage() + model = RecoveringModel() + harness = AgentHarness( + model="model", + model_client=model, + tools=[HarnessTool("slow", slow)], + storage=storage, + conversation_id="recover", + ) + run = asyncio.create_task(run_harness(harness, "first")) + await started.wait() + harness.interrupt() + assert await run == "interrupted" + assert await run_harness(harness, "second") == "recovered" + + resumed = AgentHarness( + model="model", + model_client=RecoveringModel(), + storage=storage, + conversation_id="recover", + ) + await resumed.load(create=False) + NucliaModelClient._validate_tool_history( + NucliaModelClient._messages(resumed.messages) + ) + + +@pytest.mark.asyncio +async def test_pre_start_interrupt_is_not_lost() -> None: + class NeverModel: + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + raise AssertionError("model should not run") + yield + + harness = AgentHarness(model="model", model_client=NeverModel()) + + async def run() -> str: + return await run_harness(harness, "stop") + + task = asyncio.create_task(run()) + harness.interrupt() + assert await task == "interrupted" + + +@pytest.mark.asyncio +async def test_spawn_depth_defaults_to_one() -> None: + harness = AgentHarness(model="model", model_client=Model()) + current = harness + for expected_depth in range(1, 2): + spawned = await spawn_agent( + current, SpawnAgentInput(prompt=f"depth {expected_depth}") + ) + child, _ = current.children[spawned.value["agent_id"]] + assert child.spawn_depth == expected_depth + current = child + assert "spawn_agent" not in current._tools + with pytest.raises(ValueError, match="Maximum spawn depth"): + await spawn_agent(current, SpawnAgentInput(prompt="too deep")) + await harness._stop_children() + + +@pytest.mark.asyncio +async def test_interrupt_cancels_sub_agents() -> None: + child_started = asyncio.Event() + child_cancelled = asyncio.Event() + + class SpawnModel: + def __init__(self) -> None: + self.root_calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + messages = kwargs["messages"] + if any(message.content == "child work" for message in messages): + child_started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + child_cancelled.set() + raise + yield ModelDelta(text="late") + return + self.root_calls += 1 + if self.root_calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + name="spawn_agent", + arguments={"prompt": "child work"}, + ) + ] + ) + else: + await asyncio.sleep(60) + + harness = AgentHarness(model="test-model", model_client=SpawnModel()) + task = asyncio.create_task(run_harness(harness, "delegate")) + await child_started.wait() + harness.interrupt() + + assert await task == "interrupted" + assert child_cancelled.is_set() + assert harness.children == {} + assert all( + result.value["status"] == "cancelled" + for result in harness.child_results.values() + ) + + +@pytest.mark.asyncio +async def test_unwaited_child_is_stopped_when_parent_completes() -> None: + child_started = asyncio.Event() + child_cancelled = asyncio.Event() + + class SpawnWithoutWaitModel: + def __init__(self): + self.root_calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + if any(message.content == "child" for message in kwargs["messages"]): + child_started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + child_cancelled.set() + raise + self.root_calls += 1 + if self.root_calls == 1: + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + id="spawn", + name="spawn_agent", + arguments={"prompt": "child"}, + ) + ] + ) + else: + await child_started.wait() + yield ModelDelta(text="parent done") + + harness = AgentHarness(model="model", model_client=SpawnWithoutWaitModel()) + assert await run_harness(harness, "root") == "parent done" + assert child_cancelled.is_set() + assert harness.children == {} + assert all( + result.value["status"] == "cancelled" + for result in harness.child_results.values() + ) + + +@pytest.mark.asyncio +async def test_llm_only_passes_registered_tools() -> None: + class CapturingModel: + def __init__(self) -> None: + self.tools = [] + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.tools = kwargs["tools"] + yield ModelDelta(text="ok") + + async def execute(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value) + + registered = HarnessTool("registered", execute) + unknown = HarnessTool("unknown", execute) + model = CapturingModel() + harness = AgentHarness(model="test-model", model_client=model, tools=[registered]) + + await harness.llm([registered, unknown]) + assert model.tools == [registered] + + +def test_agent_rejects_tools_that_conflict_with_core_tools() -> None: + async def execute(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value) + + with pytest.raises(ValueError, match="conflict with core tools: remember"): + AgentHarness( + model="test-model", + model_client=Model(), + tools=[HarnessTool("remember", execute)], + ) + + +def test_agent_allows_disabled_core_tool_name() -> None: + async def execute(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value) + + tool = HarnessTool("remember", execute) + harness = AgentHarness( + model="test-model", + model_client=Model(), + tools=[tool], + disabled_core_tools=["remember"], + ) + + assert harness._tools["remember"] is tool + + +@pytest.mark.asyncio +async def test_llm_does_not_expose_inactive_lazy_tool_explicitly() -> None: + class CapturingModel: + def __init__(self) -> None: + self.tools = [] + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.tools = kwargs["tools"] + yield ModelDelta(text="ok") + + async def execute(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value) + + lazy = HarnessTool("lazy", execute, lazy_load=True) + model = CapturingModel() + harness = AgentHarness(model="test-model", model_client=model, tools=[lazy]) + + await harness.llm([lazy]) + assert model.tools == [] + + +@pytest.mark.asyncio +async def test_lazy_tool_can_be_searched_activated_and_used() -> None: + class LazyModel: + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + names = {tool.name for tool in kwargs["tools"]} + self.calls += 1 + if self.calls == 1: + assert "upper" not in names + assert {"search_tools", "activate_tools"} <= names + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + name="search_tools", + arguments={"query": "uppercase"}, + ) + ] + ) + elif self.calls == 2: + assert "upper" not in names + assert any( + message.tool_name == "search_tools" + and '"name":"upper"' in message.content + for message in kwargs["messages"] + ) + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + name="activate_tools", arguments={"names": ["upper"]} + ) + ] + ) + elif self.calls == 3: + assert "upper" in names + yield ModelDelta( + tool_calls=[ + HarnessToolCall(name="upper", arguments={"value": "hi"}) + ] + ) + else: + yield ModelDelta(text="HI") + + async def upper(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value.upper()) + + harness = AgentHarness( + model="test-model", + model_client=LazyModel(), + tools=[ + HarnessTool( + "upper", + upper, + description="Convert text to uppercase.", + lazy_load=True, + ) + ], + ) + + assert await run_harness(harness, "Uppercase hi") == "HI" + + +@pytest.mark.asyncio +async def test_tool_decorator_configures_lazy_loading_and_preserves_calls() -> None: + @tool(description="Convert text", lazy_load=True) + async def upper(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value.upper()) + + assert isinstance(upper, HarnessTool) + assert upper.name == "upper" + assert upper.lazy_load is True + harness = AgentHarness(model="test-model", model_client=Model()) + assert await upper(harness, ToolInput(value="hi")) == ToolOutput(value="HI") + + +def test_tool_schema_preserves_descriptions() -> None: + class DescribedInput(BaseModel): + value: str = Field(description="Value to transform") + + async def execute(_harness: AgentHarness, _value: DescribedInput) -> ToolOutput: + return ToolOutput(value="ok") + + tool = HarnessTool("described", execute) + assert tool.input_model is DescribedInput + assert tool.output_model is ToolOutput + assert tool.parameters["properties"]["value"]["description"] == "Value to transform" + + +def test_tool_schema_flattens_nested_references() -> None: + class Nested(BaseModel): + value: str + + class NestedInput(BaseModel): + nested: Nested + values: list[Nested] + + async def execute(_harness: AgentHarness, _value: NestedInput) -> ToolOutput: + return ToolOutput(value="ok") + + schema = HarnessTool("nested", execute).parameters + rendered = str(schema) + assert "$ref" not in rendered + assert "$defs" not in rendered + assert schema["properties"]["nested"]["properties"]["value"]["type"] == "string" + + +def test_published_schema_preserves_constraints_and_namespaces() -> None: + class PublishedAgent: + __published_functions__: ClassVar[dict[str, FunctionDefinition]] = { + "lookup": FunctionDefinition( + name="lookup", + description="Lookup", + parameters={ + "kind": {"type": "string", "enum": ["a", "b"]}, + "limit": {"type": "integer", "minimum": 1}, + }, + lazy_load=True, + ) + } + + async def lookup(self, kind: str, **kwargs): + return {"kind": kind, **kwargs} + + tool = AgentHarness.to_tools(PublishedAgent(), namespace="source")[0] + assert tool.name == "source__lookup" + assert tool.parameters["properties"]["kind"]["enum"] == ["a", "b"] + assert tool.parameters["properties"]["limit"]["minimum"] == 1 + assert tool.parameters["required"] == ["kind"] + assert tool.lazy_load is True + + +@pytest.mark.asyncio +async def test_published_tool_validates_the_advertised_input_schema() -> None: + called = False + + class PublishedAgent: + __published_functions__: ClassVar[dict[str, FunctionDefinition]] = { + "lookup": FunctionDefinition( + name="lookup", + description="Lookup", + parameters={}, + input_schema={ + "type": "object", + "properties": { + "kind": {"type": "string", "enum": ["a", "b"]}, + "limit": {"type": "integer", "minimum": 1}, + }, + "required": ["kind", "limit"], + "additionalProperties": False, + }, + ) + } + + async def lookup(self, **kwargs): + nonlocal called + called = True + return kwargs + + tool = AgentHarness.to_tools(PublishedAgent())[0] + harness = AgentHarness(model="test-model", model_client=Model()) + + with pytest.raises(ValueError, match="Invalid lookup arguments"): + await tool.execute(harness, {"kind": "invalid", "limit": 1}) + with pytest.raises(ValueError, match="Invalid lookup arguments"): + await tool.execute(harness, {"kind": "a"}) + + assert called is False + + +@pytest.mark.asyncio +async def test_malformed_tool_arguments_do_not_run_handler() -> None: + called = False + + class OptionalInput(BaseModel): + value: str = "default" + + async def execute(_harness: AgentHarness, _value: OptionalInput) -> ToolOutput: + nonlocal called + called = True + return ToolOutput(value="ok") + + tool = HarnessTool("optional", execute) + harness = AgentHarness(model="test-model", model_client=Model()) + + with pytest.raises(ValueError, match="Malformed tool arguments"): + await tool.execute( + harness, {"_tool_error": "Malformed tool arguments: invalid JSON"} + ) + + assert called is False + + +def test_tool_requires_pydantic_handler_annotations() -> None: + async def missing_annotations(_harness, _value): + return ToolOutput(value="ok") + + with pytest.raises(TypeError, match="input must be annotated"): + HarnessTool("invalid", missing_annotations) + + +@pytest.mark.asyncio +async def test_published_agent_functions_convert_to_tools() -> None: + nua_client = object() + seen = [] + + class PublishedAgent: + __published_functions__: ClassVar[dict[str, FunctionDefinition]] = { + "lookup": FunctionDefinition( + name="knowledge_lookup", + description="Look up a question in the knowledge source.", + parameters={ + "question": { + "type": "string", + "description": "Question to look up.", + }, + "limit": { + "type": "integer", + "description": "Maximum number of results.", + }, + }, + ) + } + + async def lookup( + self, + question: str, + memory, + manager, + limit: int = 3, + ) -> dict[str, object]: + seen.append( + { + "question": question, + "limit": limit, + "conversation": memory.headers["conversation-id"], + "nua": manager.nua, + } + ) + return {"answer": question.upper(), "limit": limit} + + class PublishedModel: + client = nua_client + + def __init__(self) -> None: + self.calls = 0 + + async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: + self.calls += 1 + if self.calls == 1: + tool = next( + tool for tool in kwargs["tools"] if tool.name == "knowledge_lookup" + ) + assert tool.description == "Look up a question in the knowledge source." + assert ( + tool.parameters["properties"]["question"]["description"] + == "Question to look up." + ) + assert tool.parameters["properties"]["limit"]["default"] == 3 + yield ModelDelta( + tool_calls=[ + HarnessToolCall( + name="knowledge_lookup", + arguments={"question": "hyperforge"}, + ) + ] + ) + return + assert any( + message.tool_name == "knowledge_lookup" + and "HYPERFORGE" in message.content + for message in kwargs["messages"] + ) + yield ModelDelta(text="Found it") + + model = PublishedModel() + manager = Manager() + manager.nua = nua_client + tools = AgentHarness.to_tools(PublishedAgent(), manager=manager) + harness = AgentHarness( + model="test-model", + model_client=model, + tools=tools, + conversation_id="published-conversation", + ) + + assert await run_harness(harness, "Look it up") == "Found it" + assert seen == [ + { + "question": "hyperforge", + "limit": 3, + "conversation": "published-conversation", + "nua": nua_client, + } + ] + + +@pytest.mark.asyncio +async def test_published_agent_manager_uses_existing_nua_api() -> None: + class FakeNua: + async def generate(self, **kwargs): + return GenerativeFullResponse(answer="legacy answer") + + class LegacyAgent: + __published_functions__: ClassVar[dict[str, FunctionDefinition]] = { + "ask": FunctionDefinition( + name="ask", + description="Ask", + parameters={"question": {"type": "string"}}, + ) + } + + async def ask(self, question, manager, **kwargs): + from nuclia.lib.nua_responses import ChatModel + + response, _, _ = await manager.execute_raw(ChatModel(question=question)) + return response.answer + + manager = Manager() + manager.nua = FakeNua() + tool = AgentHarness.to_tools(LegacyAgent(), manager=manager)[0] + harness = AgentHarness(model="model", model_client=Model(), tools=[tool]) + output = await tool.execute(harness, {"question": "hello"}) + assert output.value == "legacy answer" + + +def test_agent_without_published_functions_is_not_supported() -> None: + class UnsupportedAgent: + pass + + with pytest.raises(TypeError, match="does not define any published functions"): + AgentHarness.to_tools(UnsupportedAgent()) diff --git a/uv.lock b/uv.lock index 80f1cab8..bc1ac40f 100644 --- a/uv.lock +++ b/uv.lock @@ -945,6 +945,7 @@ dependencies = [ { name = "httpx" }, { name = "hyperforge-restricted" }, { name = "jinja2" }, + { name = "jsonschema" }, { name = "lru-dict" }, { name = "mcp" }, { name = "mmh3" }, @@ -955,10 +956,11 @@ dependencies = [ { name = "nucliadb-sdk" }, { name = "nucliadb-telemetry", extra = ["otel"] }, { name = "prometheus-client" }, - { name = "psycopg2" }, + { name = "psycopg2-binary" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "redis", extra = ["hiredis"] }, + { name = "restrictedpython" }, { name = "ruff" }, { name = "sentry-sdk" }, { name = "sqlalchemy" }, @@ -995,6 +997,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.28.1,<0.29" }, { name = "hyperforge-restricted", editable = "agents/restricted" }, { name = "jinja2" }, + { name = "jsonschema" }, { name = "lru-dict" }, { name = "mcp", specifier = ">=1.26.0" }, { name = "mmh3" }, @@ -1005,10 +1008,11 @@ requires-dist = [ { name = "nucliadb-sdk", specifier = ">=6.13.1.post6414" }, { name = "nucliadb-telemetry", extras = ["otel"], specifier = ">=6.15.0" }, { name = "prometheus-client" }, - { name = "psycopg2" }, + { name = "psycopg2-binary" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "redis", extras = ["hiredis"], specifier = ">=6.4.0" }, + { name = "restrictedpython" }, { name = "ruff", specifier = ">=0.15.10" }, { name = "sentry-sdk" }, { name = "sqlalchemy" }, @@ -2807,17 +2811,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/57/bf/2086963c69bdac3d7cff1cc7ff79b8ce5ea0bec6797a017e1be338a46248/protobuf-6.33.5-py3-none-any.whl", hash = "sha256:69915a973dd0f60f31a08b8318b73eab2bd6a392c79184b3612226b0a3f8ec02", size = 170687, upload-time = "2026-01-29T21:51:32.557Z" }, ] -[[package]] -name = "psycopg2" -version = "2.9.10" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/62/51/2007ea29e605957a17ac6357115d0c1a1b60c8c984951c19419b3474cdfd/psycopg2-2.9.10.tar.gz", hash = "sha256:12ec0b40b0273f95296233e8750441339298e6a572f7039da5b260e3c8b60e11", size = 385672, upload-time = "2024-10-16T11:24:54.832Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/16/4623fad6076448df21c1a870c93a9774ad8a7b4dd1660223b59082dd8fec/psycopg2-2.9.10-cp312-cp312-win32.whl", hash = "sha256:65a63d7ab0e067e2cdb3cf266de39663203d38d6a8ed97f5ca0cb315c73fe067", size = 1025113, upload-time = "2024-10-16T11:18:40.148Z" }, - { url = "https://files.pythonhosted.org/packages/66/de/baed128ae0fc07460d9399d82e631ea31a1f171c0c4ae18f9808ac6759e3/psycopg2-2.9.10-cp312-cp312-win_amd64.whl", hash = "sha256:4a579d6243da40a7b3182e0430493dbd55950c493d8c68f4eec0b302f6bbf20e", size = 1163951, upload-time = "2024-10-16T11:18:44.377Z" }, - { url = "https://files.pythonhosted.org/packages/ae/49/a6cfc94a9c483b1fa401fbcb23aca7892f60c7269c5ffa2ac408364f80dc/psycopg2-2.9.10-cp313-cp313-win_amd64.whl", hash = "sha256:91fd603a2155da8d0cfcdbf8ab24a2d54bca72795b90d2a3ed2b6da8d979dee2", size = 2569060, upload-time = "2025-01-04T20:09:15.28Z" }, -] - [[package]] name = "psycopg2-binary" version = "2.9.10" From dfad4635c74ba422f9a2b08d2801efc95b034e91 Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:08:34 -0400 Subject: [PATCH 2/4] disable some tools --- hyperforge/src/hyperforge/harness/agent.py | 10 +++++++++- hyperforge/tests/harness/test_harness_agent.py | 3 +++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/hyperforge/src/hyperforge/harness/agent.py b/hyperforge/src/hyperforge/harness/agent.py index fe667513..a8746fdb 100644 --- a/hyperforge/src/hyperforge/harness/agent.py +++ b/hyperforge/src/hyperforge/harness/agent.py @@ -59,7 +59,15 @@ async def __call__(self, memory: QuestionMemory, manager: Manager) -> None: system_prompt=system_prompt, title=self.config.workflow.name, conversation_id=memory.get_session_id(), - disabled_core_tools=self.config.disabled_core_tools, + disabled_core_tools=[ + *self.config.disabled_core_tools, + # Since we're not using the storage adapter implementation, + # we need to disable these tools that require storage access + # across sessions. + "remember", + "recall", + "forget", + ], execution_context={ "user_id": memory.session.user_info.get("id", "system"), "conversation_metadata": { diff --git a/hyperforge/tests/harness/test_harness_agent.py b/hyperforge/tests/harness/test_harness_agent.py index 7ed14fed..7725d711 100644 --- a/hyperforge/tests/harness/test_harness_agent.py +++ b/hyperforge/tests/harness/test_harness_agent.py @@ -43,6 +43,9 @@ async def chat_completions_stream(self, payload, **kwargs) -> AsyncIterator[dict message["role"] == "system" and "- Be concise" in message["content"] for message in payload["messages"] ) + assert not {"remember", "recall", "forget"} & { + tool["function"]["name"] for tool in payload["tools"] + } if self.calls == 1: tool = next( item From 1ccfc903acff468fa0c623b9ea26420262a0e1ff Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:47:44 -0400 Subject: [PATCH 3/4] fixes --- hyperforge/src/hyperforge/harness/agent.py | 3 -- hyperforge/src/hyperforge/harness/config.py | 4 --- .../src/hyperforge/harness_sdk/clients.py | 33 +++++++++++++---- .../src/hyperforge/harness_sdk/harness.py | 26 +++++--------- .../src/hyperforge/harness_sdk/usage.py | 6 +++- .../tests/harness/test_harness_agent.py | 6 ---- hyperforge/tests/harness_sdk/test_clients.py | 35 +++++++++++++++++++ hyperforge/tests/harness_sdk/test_harness.py | 23 ++++++++++-- 8 files changed, 96 insertions(+), 40 deletions(-) diff --git a/hyperforge/src/hyperforge/harness/agent.py b/hyperforge/src/hyperforge/harness/agent.py index a8746fdb..56259195 100644 --- a/hyperforge/src/hyperforge/harness/agent.py +++ b/hyperforge/src/hyperforge/harness/agent.py @@ -57,7 +57,6 @@ async def __call__(self, memory: QuestionMemory, manager: Manager) -> None: reasoning_effort=self.config.reasoning_effort, tools=tools, system_prompt=system_prompt, - title=self.config.workflow.name, conversation_id=memory.get_session_id(), disabled_core_tools=[ *self.config.disabled_core_tools, @@ -79,8 +78,6 @@ async def __call__(self, memory: QuestionMemory, manager: Manager) -> None: }, feedback_enabled=self.config.feedback_enabled, usage_limits=self.config.usage_limits, - max_spawn_depth=self.config.max_spawn_depth, - max_concurrent_agents=self.config.max_concurrent_agents, ) history = await memory.get_chat_history() harness.add_messages( diff --git a/hyperforge/src/hyperforge/harness/config.py b/hyperforge/src/hyperforge/harness/config.py index c7e25be1..c38dfbc7 100644 --- a/hyperforge/src/hyperforge/harness/config.py +++ b/hyperforge/src/hyperforge/harness/config.py @@ -7,7 +7,6 @@ from hyperforge.driver import DriverConfig from hyperforge.harness_sdk import SYSTEM_PROMPT, ReasoningEffort, UsageLimits from hyperforge.models import MemoryConfig, Rules -from hyperforge.workflows import WorkflowData class HarnessAgentConfig(BaseModel): @@ -19,11 +18,8 @@ class HarnessAgentConfig(BaseModel): disabled_core_tools: list[str] = Field(default_factory=list) feedback_enabled: bool = False usage_limits: UsageLimits = Field(default_factory=UsageLimits) - max_spawn_depth: int = 1 - max_concurrent_agents: int = 4 rules: Rules = Field(default_factory=Rules) memory: MemoryConfig = Field(default_factory=MemoryConfig) - workflow: WorkflowData @field_validator("agents", mode="before") @classmethod diff --git a/hyperforge/src/hyperforge/harness_sdk/clients.py b/hyperforge/src/hyperforge/harness_sdk/clients.py index 6ab82fbe..c8bca71a 100644 --- a/hyperforge/src/hyperforge/harness_sdk/clients.py +++ b/hyperforge/src/hyperforge/harness_sdk/clients.py @@ -106,6 +106,30 @@ def __init__( self.provider_data = provider_data or {} +def _request_error_detail( + exc: httpx.RequestError | httpx.HTTPStatusError, +) -> tuple[str, dict[str, Any]]: + response = exc.response if isinstance(exc, httpx.HTTPStatusError) else None + status = response.status_code if response is not None else None + response_body = response.text.strip()[:2000] if response is not None else "" + request = exc.request + url = str(request.url) if request is not None else "unknown URL" + detail = str(exc).strip() or repr(exc) + parts = [f"{type(exc).__name__} for {url}"] + if status is not None: + parts.append(f"status={status}") + if response_body: + parts.append(f"response={response_body}") + parts.append(f"error={detail}") + return "; ".join(parts), { + "http_status": status, + "url": url, + "response_body": response_body or None, + "error_type": type(exc).__name__, + "error": detail, + } + + class NucliaChatCompletionsClient: """Chat-completions transport backed by Hyperforge's shared NUA client.""" @@ -172,13 +196,10 @@ async def stream( yield chunk return except (httpx.RequestError, httpx.HTTPStatusError) as exc: + detail, provider_data = _request_error_detail(exc) error = NucliaChatCompletionsError( - f"Nuclia chat completions request failed: {exc}", - provider_data={ - "http_status": exc.response.status_code - if isinstance(exc, httpx.HTTPStatusError) - else None - }, + f"Nuclia chat completions request failed: {detail}", + provider_data=provider_data, ) except NucliaChatCompletionsError as exc: error = exc diff --git a/hyperforge/src/hyperforge/harness_sdk/harness.py b/hyperforge/src/hyperforge/harness_sdk/harness.py index d15bda04..ffc065f5 100644 --- a/hyperforge/src/hyperforge/harness_sdk/harness.py +++ b/hyperforge/src/hyperforge/harness_sdk/harness.py @@ -66,8 +66,6 @@ class HarnessConfig: tags: tuple[str, ...] feedback_enabled: bool usage_limits: UsageLimits - max_spawn_depth: int - max_concurrent_agents: int event_queue_size: int @@ -292,18 +290,10 @@ def __init__( tags: Iterable[str] = (), feedback_enabled: bool = False, usage_limits: UsageLimits | None = None, - max_spawn_depth: int = 1, - max_concurrent_agents: int = 4, spawn_depth: int = 0, parent_agent_id: str | None = None, event_queue_size: int = DEFAULT_EVENT_QUEUE_SIZE, ) -> None: - if max_spawn_depth < 0: - raise ValueError("max_spawn_depth must be non-negative") - if max_concurrent_agents < 1: - raise ValueError("max_concurrent_agents must be positive") - if spawn_depth < 0 or spawn_depth > max_spawn_depth: - raise ValueError("spawn_depth must be between zero and max_spawn_depth") if event_queue_size <= 0: raise ValueError("event_queue_size must be positive") normalized_tools = tuple( @@ -313,6 +303,8 @@ def __init__( normalized_context = dict(execution_context or {}) normalized_tags = tuple(dict.fromkeys(tags)) normalized_limits = usage_limits or UsageLimits() + if spawn_depth < 0 or spawn_depth > normalized_limits.max_spawn_depth: + raise ValueError("spawn_depth must be between zero and max_spawn_depth") self._config = HarnessConfig( model=model, reasoning_effort=reasoning_effort, @@ -327,8 +319,6 @@ def __init__( tags=normalized_tags, feedback_enabled=feedback_enabled, usage_limits=normalized_limits, - max_spawn_depth=max_spawn_depth, - max_concurrent_agents=max_concurrent_agents, event_queue_size=event_queue_size, ) self.model = self._config.model @@ -344,8 +334,8 @@ def __init__( self.category = self._config.category self.tags = list(self._config.tags) self.usage_limits = self._config.usage_limits - self.max_spawn_depth = self._config.max_spawn_depth - self.max_concurrent_agents = self._config.max_concurrent_agents + self.max_spawn_depth = self.usage_limits.max_spawn_depth + self.max_concurrent_agents = self.usage_limits.max_concurrent_agents self.spawn_depth = spawn_depth self.parent_agent_id = parent_agent_id self.usage = HarnessUsage() @@ -992,7 +982,11 @@ async def _execute_tool_call(self, call: HarnessToolCall) -> HarnessMessage: ) except Exception as exc: logger.warning( - "Agent tool execution failed", + "Agent tool execution failed: tool=%s call_id=%s error_type=%s error=%s", + call.name, + call.id, + type(exc).__name__, + exc, extra={ "conversation_id": self.conversation_id, "turn_id": self._turn_id, @@ -1051,8 +1045,6 @@ def _create_child(self, child_id: str, *, include_history: bool) -> AgentHarness tags=self._config.tags, feedback_enabled=self._config.feedback_enabled, usage_limits=self._config.usage_limits, - max_spawn_depth=self._config.max_spawn_depth, - max_concurrent_agents=self._config.max_concurrent_agents, spawn_depth=self.spawn_depth + 1, parent_agent_id=self.agent_id, event_queue_size=self._config.event_queue_size, diff --git a/hyperforge/src/hyperforge/harness_sdk/usage.py b/hyperforge/src/hyperforge/harness_sdk/usage.py index a65966ef..71b73238 100644 --- a/hyperforge/src/hyperforge/harness_sdk/usage.py +++ b/hyperforge/src/hyperforge/harness_sdk/usage.py @@ -13,13 +13,17 @@ class UsageLimits: max_input_tokens: int | None = None max_output_tokens: int | None = None max_time: float | None = None + max_spawn_depth: int = 1 + max_concurrent_agents: int = 4 max_codemode_runtime_seconds: float = DEFAULT_CODEMODE_RUNTIME_SECONDS max_codemode_memory_bytes: int = DEFAULT_CODEMODE_MEMORY_BYTES def __post_init__(self) -> None: for name, value in vars(self).items(): - if value is not None and value <= 0: + if value is not None and value <= 0 and name != "max_spawn_depth": raise ValueError(f"{name} must be greater than zero") + if self.max_spawn_depth < 0: + raise ValueError("max_spawn_depth must be non-negative") @dataclass diff --git a/hyperforge/tests/harness/test_harness_agent.py b/hyperforge/tests/harness/test_harness_agent.py index 7725d711..48d3873d 100644 --- a/hyperforge/tests/harness/test_harness_agent.py +++ b/hyperforge/tests/harness/test_harness_agent.py @@ -99,12 +99,6 @@ async def test_harness_agent_runs_configured_legacy_agent_as_tool(monkeypatch) - { "model": "test-model", "agents": [{"id": "source", "module": "published"}], - "workflow": { - "id": "default", - "name": "Default", - "description": None, - "parameters": None, - }, } ) agent = await HarnessAgent.from_config_class(config) diff --git a/hyperforge/tests/harness_sdk/test_clients.py b/hyperforge/tests/harness_sdk/test_clients.py index 5caab3f8..47955527 100644 --- a/hyperforge/tests/harness_sdk/test_clients.py +++ b/hyperforge/tests/harness_sdk/test_clients.py @@ -1,5 +1,6 @@ from collections.abc import AsyncIterator +import httpx import pytest from hyperforge.harness_sdk import ( @@ -8,6 +9,7 @@ HarnessMessage, HarnessToolCall, NucliaChatCompletionsClient, + NucliaChatCompletionsError, NucliaModelClient, ) @@ -89,6 +91,39 @@ async def chat_completions_stream(self, payload, **kwargs): assert [chunk.choices[0].delta.content for chunk in chunks] == ["hello"] +@pytest.mark.asyncio +async def test_chat_completions_http_error_includes_response_detail() -> None: + request = httpx.Request("POST", "https://predict.example/chat") + response = httpx.Response( + 408, request=request, json={"detail": "Timeout connecting to generative API"} + ) + + class FakeNua: + async def chat_completions_stream(self, payload, **kwargs): + raise httpx.HTTPStatusError("", request=request, response=response) + yield {} + + client = NucliaChatCompletionsClient(FakeNua()) + + with pytest.raises(NucliaChatCompletionsError) as raised: + _ = [ + chunk + async for chunk in client.stream( + ChatCompletionRequest( + messages=[{"role": "user", "content": "hi"}], model="model" + ) + ) + ] + + message = str(raised.value) + assert "status=408" in message + assert "Timeout connecting to generative API" in message + assert ( + raised.value.provider_data["response_body"] + == '{"detail":"Timeout connecting to generative API"}' + ) + + @pytest.mark.asyncio async def test_harness_client_keeps_usage_only_chunk() -> None: class UsageClient: diff --git a/hyperforge/tests/harness_sdk/test_harness.py b/hyperforge/tests/harness_sdk/test_harness.py index 5cc6dfd4..4a3a9865 100644 --- a/hyperforge/tests/harness_sdk/test_harness.py +++ b/hyperforge/tests/harness_sdk/test_harness.py @@ -615,7 +615,7 @@ async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: harness = AgentHarness( model="test-model", model_client=SlowChildModel(), - max_concurrent_agents=2, + usage_limits=UsageLimits(max_concurrent_agents=2), ) first = await spawn_agent( harness, SpawnAgentInput(prompt="first", include_history=False) @@ -658,8 +658,7 @@ async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: harness = AgentHarness( model="test-model", model_client=SlowChildModel(), - max_spawn_depth=2, - max_concurrent_agents=2, + usage_limits=UsageLimits(max_spawn_depth=2, max_concurrent_agents=2), ) spawned = await spawn_agent( harness, SpawnAgentInput(prompt="child", include_history=False) @@ -1384,6 +1383,24 @@ async def execute(_harness: AgentHarness, _value: OptionalInput) -> ToolOutput: assert called is False +@pytest.mark.asyncio +async def test_tool_failure_log_includes_actionable_detail(caplog) -> None: + async def fail(_harness: AgentHarness, _value: ToolInput) -> ToolOutput: + raise ValueError("identity field is invalid") + + harness = AgentHarness( + model="test-model", model_client=Model(), tools=[HarnessTool("validate", fail)] + ) + call = HarnessToolCall(id="call-42", name="validate", arguments={"value": "x"}) + + await harness._execute_tool_call(call) + + assert ( + "Agent tool execution failed: tool=validate call_id=call-42 " + "error_type=ValueError error=identity field is invalid" + ) in caplog.text + + def test_tool_requires_pydantic_handler_annotations() -> None: async def missing_annotations(_harness, _value): return ToolOutput(value="ok") From 02b9b77b4695381a3e2485688213dfdbc9e6f790 Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:07:52 -0400 Subject: [PATCH 4/4] Fix harness logging test isolation --- hyperforge/tests/harness_sdk/test_harness.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/hyperforge/tests/harness_sdk/test_harness.py b/hyperforge/tests/harness_sdk/test_harness.py index 4a3a9865..a3cc6188 100644 --- a/hyperforge/tests/harness_sdk/test_harness.py +++ b/hyperforge/tests/harness_sdk/test_harness.py @@ -1384,7 +1384,7 @@ async def execute(_harness: AgentHarness, _value: OptionalInput) -> ToolOutput: @pytest.mark.asyncio -async def test_tool_failure_log_includes_actionable_detail(caplog) -> None: +async def test_tool_failure_log_includes_actionable_detail(mocker) -> None: async def fail(_harness: AgentHarness, _value: ToolInput) -> ToolOutput: raise ValueError("identity field is invalid") @@ -1392,13 +1392,18 @@ async def fail(_harness: AgentHarness, _value: ToolInput) -> ToolOutput: model="test-model", model_client=Model(), tools=[HarnessTool("validate", fail)] ) call = HarnessToolCall(id="call-42", name="validate", arguments={"value": "x"}) + warning = mocker.patch("hyperforge.harness_sdk.harness.logger.warning") await harness._execute_tool_call(call) - assert ( - "Agent tool execution failed: tool=validate call_id=call-42 " - "error_type=ValueError error=identity field is invalid" - ) in caplog.text + assert warning.call_count == 1 + assert warning.call_args.args[:4] == ( + "Agent tool execution failed: tool=%s call_id=%s error_type=%s error=%s", + "validate", + "call-42", + "ValueError", + ) + assert str(warning.call_args.args[4]) == "identity field is invalid" def test_tool_requires_pydantic_handler_annotations() -> None: