diff --git a/opensin_core/agent/loop.py b/opensin_core/agent/loop.py index 94d9b331..15536750 100644 --- a/opensin_core/agent/loop.py +++ b/opensin_core/agent/loop.py @@ -11,6 +11,7 @@ from opensin_core.session.session_store import SessionStore, Session from opensin_core.session.compaction import Compactor, CompactionConfig from opensin_core.session.permissions import PermissionManager, PermissionLevel +from opensin_core.recursive_mas import RecursiveMasConfig, RecursiveMasMonitor class LLMClient(Protocol): @@ -47,6 +48,7 @@ class AgentConfig: sandbox_enabled: bool = False system_prompt: str | None = None workspace_root: str = "." + recursive_mas: RecursiveMasConfig | None = None @dataclass @@ -80,6 +82,11 @@ def __init__( self.config = config or AgentConfig() self.permissions.set_default_level(self.config.permission_mode) self.permissions.set_default_level(self.config.permission_mode) + self.recursive_mas_monitor = ( + RecursiveMasMonitor(self.config.recursive_mas) + if self.config.recursive_mas + else None + ) self._current_session: Session | None = None async def run( @@ -171,6 +178,11 @@ async def run( except (json.JSONDecodeError, Exception): pass + if self.recursive_mas_monitor: + self.recursive_mas_monitor.before_tool_use( + tool_name, json.dumps(tool_args) + ) + allowed, reason = self.permissions.check_permission( tool_name, self._get_tool_permission_level(tool_name), @@ -179,10 +191,24 @@ async def run( tool_result = ToolResult.error( f"Permission denied: {reason}", code=403 ) + if self.recursive_mas_monitor: + self.recursive_mas_monitor.after_tool_use( + tool_name, + json.dumps(tool_args), + reason, + True, + ) else: tool_result = await self.registry.execute( tool_name, tool_args, context ) + if self.recursive_mas_monitor: + self.recursive_mas_monitor.after_tool_use( + tool_name, + json.dumps(tool_args), + tool_result.output, + not tool_result.success, + ) all_tool_results.append(tool_result) @@ -268,6 +294,11 @@ async def run_streaming( except (json.JSONDecodeError, Exception): pass + if self.recursive_mas_monitor: + self.recursive_mas_monitor.before_tool_use( + tool_name, json.dumps(tool_args) + ) + allowed, _ = self.permissions.check_permission( tool_name, self._get_tool_permission_level(tool_name), @@ -276,6 +307,15 @@ async def run_streaming( tool_result = await self.registry.execute( tool_name, tool_args, context ) + + if self.recursive_mas_monitor: + self.recursive_mas_monitor.after_tool_use( + tool_name, + json.dumps(tool_args), + tool_result.output, + not tool_result.success, + ) + yield TurnResult( role="tool", content=tool_result.output, diff --git a/opensin_core/recursive_mas.py b/opensin_core/recursive_mas.py new file mode 100644 index 00000000..1d0b6398 --- /dev/null +++ b/opensin_core/recursive_mas.py @@ -0,0 +1,157 @@ + +import json +import subprocess +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class RecursiveMasConfig: + enabled: bool = False + topology: str = "chain" + latent_dim: int = 64 + cli_path: str = "recursivemas" + log_traces: bool = True + + +@dataclass +class RecursiveMasSession: + enabled: bool = False + topology: str = "chain" + latent_dim: int = 64 + cli_path: str = "recursivemas" + log_traces: bool = True + recursive_depth: int = 0 + latent_traces: dict[str, list[float]] = field(default_factory=dict) + + def is_enabled(self) -> bool: + return self.enabled + + def increment_depth(self) -> None: + if not self.enabled: + return + self.recursive_depth += 1 + + def decrement_depth(self) -> None: + if not self.enabled: + return + if self.recursive_depth > 0: + self.recursive_depth -= 1 + + def current_depth(self) -> int: + return self.recursive_depth + + def log_trace(self, tool_name: str, input_str: str, output: str) -> None: + if not self.log_traces: + return + + trace_key = f"{tool_name}_{self.recursive_depth}" + trace_value = self._compute_latent_trace(tool_name, input_str, output) + self.latent_traces[trace_key] = trace_value + + def _compute_latent_trace( + self, tool_name: str, input_str: str, output: str + ) -> list[float]: + dim = self.latent_dim + trace = [] + for i in range(dim): + value = ((len(tool_name) * (i + 1)) / 100.0) * ( + (len(input_str) + len(output)) / 50.0 + ) + trace.append(value) + return trace + + def call_recursivemas_cli(self, args: list[str]) -> str | None: + if not self.enabled: + return None + + try: + result = subprocess.run( + [self.cli_path] + args, + capture_output=True, + text=True, + timeout=30, + ) + if result.returncode == 0: + return result.stdout.strip() + return result.stderr.strip() + except (subprocess.SubprocessError, FileNotFoundError): + return None + + def get_traces(self) -> dict[str, list[float]]: + return self.latent_traces + + def build_trace_message(self) -> str: + if not self.latent_traces: + return "" + + lines = ["RecursiveMAS Traces:"] + for key, trace in self.latent_traces.items(): + lines.append(f" {key}: {trace}") + return "\n".join(lines) + + +class RecursiveMasMonitor: + def __init__(self, config: RecursiveMasConfig | None = None): + self.session = RecursiveMasSession( + enabled=config.enabled if config else False, + topology=config.topology if config else "chain", + latent_dim=config.latent_dim if config else 64, + cli_path=config.cli_path if config else "recursivemas", + log_traces=config.log_traces if config else True, + ) + + def is_enabled(self) -> bool: + return self.session.is_enabled() + + def before_tool_use(self, tool_name: str, input_str: str) -> None: + if not self.is_enabled(): + return + + self.session.increment_depth() + self.session.call_recursivemas_cli( + [ + "inspect", + "--tool", tool_name, + "--input", json.dumps(input_str), + "--topology", self.session.topology, + ] + ) + + def after_tool_use( + self, tool_name: str, input_str: str, output: str, is_error: bool + ) -> None: + if not self.is_enabled(): + return + + self.session.log_trace(tool_name, input_str, output) + + status = "error" if is_error else "success" + self.session.call_recursivemas_cli( + [ + "benchmark", + "--tool", tool_name, + "--status", status, + "--depth", str(self.session.current_depth()), + ] + ) + + self.session.decrement_depth() + + def get_session(self) -> RecursiveMasSession: + return self.session + + +def monitor_tool_use( + tool_name: str, input_str: str, output: str, is_error: bool, + config: RecursiveMasConfig | None = None, +) -> str: + monitor = RecursiveMasMonitor(config) + + if not monitor.is_enabled(): + return "" + + monitor.before_tool_use(tool_name, input_str) + monitor.after_tool_use(tool_name, input_str, output, is_error) + + return monitor.get_session().build_trace_message() diff --git a/opensin_core/tests/test_agent_loop.py b/opensin_core/tests/test_agent_loop.py index e35ea6f6..a928512d 100644 --- a/opensin_core/tests/test_agent_loop.py +++ b/opensin_core/tests/test_agent_loop.py @@ -7,6 +7,7 @@ import pytest from opensin_core.agent.loop import AgentLoop, AgentConfig, TurnResult +from opensin_core.recursive_mas import RecursiveMasConfig, RecursiveMasSession, RecursiveMasMonitor from opensin_core.tools import ( ToolRegistry, ToolResult, @@ -300,3 +301,89 @@ def test_with_tool_calls(self): ) assert len(result.tool_calls) == 1 assert len(result.tool_results) == 1 + + +# --- RecursiveMAS Tests --- + + +class TestRecursiveMasConfig: + def test_defaults(self): + config = RecursiveMasConfig() + assert config.enabled is False + assert config.topology == "chain" + assert config.latent_dim == 64 + assert config.cli_path == "recursivemas" + assert config.log_traces is True + + def test_custom_values(self): + config = RecursiveMasConfig( + enabled=True, topology="star", latent_dim=128, + cli_path="/usr/bin/recursivemas", log_traces=False + ) + assert config.enabled is True + assert config.topology == "star" + assert config.latent_dim == 128 + assert config.cli_path == "/usr/bin/recursivemas" + assert config.log_traces is False + + +class TestRecursiveMasSession: + def test_depth_tracking(self): + config = RecursiveMasConfig(enabled=True) + session = RecursiveMasSession( + enabled=True, topology="chain", latent_dim=4, + cli_path="echo", log_traces=True + ) + assert session.is_enabled() is True + assert session.current_depth() == 0 + + session.increment_depth() + assert session.current_depth() == 1 + + session.increment_depth() + assert session.current_depth() == 2 + + session.decrement_depth() + assert session.current_depth() == 1 + + def test_log_traces(self): + config = RecursiveMasConfig(enabled=True, latent_dim=4) + session = RecursiveMasSession( + enabled=True, topology="chain", latent_dim=4, + cli_path="echo", log_traces=True + ) + session.log_trace("bash", "ls", "file1.txt") + traces = session.get_traces() + assert len(traces) > 0 + + def test_disabled_session(self): + session = RecursiveMasSession(enabled=False) + assert session.is_enabled() is False + session.increment_depth() # Should not crash + assert session.current_depth() == 0 + + +class TestRecursiveMasMonitor: + def test_enabled_monitor(self): + config = RecursiveMasConfig(enabled=True, latent_dim=4) + monitor = RecursiveMasMonitor(config) + assert monitor.is_enabled() is True + + def test_disabled_monitor(self): + config = RecursiveMasConfig(enabled=False) + monitor = RecursiveMasMonitor(config) + assert monitor.is_enabled() is False + + def test_before_after_calls(self): + config = RecursiveMasConfig(enabled=True, latent_dim=4) + monitor = RecursiveMasMonitor(config) + monitor.before_tool_use("bash", '{"cmd": "ls"}') + assert monitor.get_session().current_depth() == 1 + + monitor.after_tool_use("bash", '{"cmd": "ls"}', "output", False) + assert monitor.get_session().current_depth() == 0 + + traces = monitor.get_session().get_traces() + assert len(traces) > 0 + +