jx-react-agent 0.1.0__py3-none-any.whl

rewardkit_react/agent.py

@@ -0,0 +1,537 @@
+"""A small native-tool-calling ReAct backend for RewardKit 0.2.x."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import copy
+import io
+import json
+import mimetypes
+import os
+import tempfile
+import time
+import uuid
+from collections import Counter
+from pathlib import Path
+from typing import Any
+
+import litellm
+from rewardkit import AgentAttempt, AgentBackend
+
+from rewardkit_react.registry import ToolRegistry, load_registry
+from rewardkit_react.tool import ToolContext, ToolResult
+
+_DEFAULT_MAX_STEPS = 20
+_DEFAULT_MAX_TOOL_CHARS = 40_000
+_DEFAULT_MAX_IMAGE_BYTES = 8 * 1024 * 1024
+
+
+class ReActBackend(AgentBackend):
+    """Explore a workspace through model tool calls, then submit a score."""
+
+    name = "react"
+
+    def __init__(self, judge: Any, cwd: str | None) -> None:
+        super().__init__(judge, cwd)
+        self.registry, self.load_warnings = load_registry()
+
+    async def run(
+        self, prompt: str, schema: dict[str, Any], label: str
+    ) -> AgentAttempt:
+        model = self.judge.model or os.environ.get("REACT_JUDGE_MODEL")
+        if not model:
+            return AgentAttempt(
+                error=ValueError(
+                    "the react judge needs [judge].model or REACT_JUDGE_MODEL"
+                ),
+                warnings=list(self.load_warnings),
+            )
+
+        workspace = Path(self.cwd or os.getcwd()).resolve()
+        if not workspace.is_dir():
+            return AgentAttempt(
+                error=ValueError(f"judge workspace does not exist: {workspace}"),
+                warnings=list(self.load_warnings),
+            )
+
+        max_steps = _positive_int_env("REACT_JUDGE_MAX_STEPS", _DEFAULT_MAX_STEPS)
+        max_tool_chars = _positive_int_env(
+            "REACT_JUDGE_MAX_TOOL_CHARS", _DEFAULT_MAX_TOOL_CHARS
+        )
+        max_image_bytes = _positive_int_env(
+            "REACT_JUDGE_MAX_IMAGE_BYTES", _DEFAULT_MAX_IMAGE_BYTES
+        )
+        timeout = float(self.judge.timeout)
+        started = time.monotonic()
+        usage = _Usage(model)
+        trace: list[dict[str, Any]] = []
+        last_request: dict[str, Any] | None = None
+        last_response: dict[str, Any] | None = None
+        warnings = list(self.load_warnings)
+
+        messages: list[dict[str, Any]] = [
+            {
+                "role": "system",
+                "content": _system_prompt(workspace, max_steps),
+            },
+            {"role": "user", "content": prompt},
+        ]
+        tools = [tool.function_schema() for tool in self.registry.values()]
+        tools.append(_submit_score_tool(schema))
+
+        try:
+            with tempfile.TemporaryDirectory(prefix="rewardkit-react-") as temp_dir:
+                async with asyncio.timeout(timeout):
+                    for step in range(1, max_steps + 1):
+                        remaining = max(0.1, timeout - (time.monotonic() - started))
+                        completion_request = {
+                            "model": model,
+                            "messages": messages,
+                            "tools": tools,
+                            "timeout": remaining,
+                            **_completion_overrides(),
+                        }
+                        last_request = _trace_request(completion_request)
+                        last_response = None
+                        response = await litellm.acompletion(
+                            **completion_request,
+                        )
+                        usage.add_response(response)
+                        message = response.choices[0].message
+                        content = _field(message, "content")
+                        tool_calls = list(_field(message, "tool_calls") or [])
+                        assistant_message = _assistant_message(content, tool_calls)
+                        last_response = copy.deepcopy(assistant_message)
+                        messages.append(assistant_message)
+                        trace.append(
+                            {
+                                "step": step,
+                                "event": "assistant",
+                                "content": content,
+                                "tool_calls": assistant_message.get("tool_calls", []),
+                            }
+                        )
+
+                        if not tool_calls:
+                            trace.append(
+                                {
+                                    "step": step,
+                                    "event": "missing_tool_call",
+                                    "content": _shorten(str(content or ""), 1000),
+                                }
+                            )
+                            messages.append(
+                                {
+                                    "role": "user",
+                                    "content": (
+                                        "Continue by calling one of the provided tools. "
+                                        "When the evidence is sufficient, call submit_score."
+                                    ),
+                                }
+                            )
+                            continue
+
+                        image_observations: list[tuple[str, ToolResult]] = []
+                        for tool_call in tool_calls:
+                            call_id, name, arguments, argument_error = _parse_tool_call(
+                                tool_call
+                            )
+                            usage.tool_requests[name] += 1
+                            if argument_error:
+                                result = ToolResult(error=argument_error)
+                            elif name == "submit_score":
+                                trace.append(
+                                    {
+                                        "step": step,
+                                        "tool": name,
+                                        "arguments": arguments,
+                                    }
+                                )
+                                return AgentAttempt(
+                                    output=json.dumps(arguments, ensure_ascii=False),
+                                    usage=usage.as_rewardkit(),
+                                    logs=_save_trace(
+                                        label=label,
+                                        model=model,
+                                        workspace=workspace,
+                                        outcome="completed",
+                                        request=last_request,
+                                        response=last_response,
+                                        trace=trace,
+                                        usage=usage.as_rewardkit(),
+                                        warnings=warnings,
+                                    ),
+                                    warnings=warnings,
+                                )
+                            else:
+                                result = await _execute_tool(
+                                    self.registry,
+                                    name,
+                                    arguments,
+                                    ToolContext(
+                                        workspace=workspace,
+                                        temp_dir=Path(temp_dir),
+                                        timeout=min(remaining, 120.0),
+                                        max_output_chars=max_tool_chars,
+                                    ),
+                                )
+
+                            observation = result.observation(max_tool_chars)
+                            messages.append(
+                                {
+                                    "role": "tool",
+                                    "tool_call_id": call_id,
+                                    "name": name,
+                                    "content": observation,
+                                }
+                            )
+                            trace.append(
+                                {
+                                    "step": step,
+                                    "tool": name,
+                                    "arguments": arguments,
+                                    "observation": observation,
+                                    "images": [str(path) for path in result.images],
+                                }
+                            )
+                            if result.images:
+                                image_observations.append((name, result))
+
+                        if image_observations:
+                            blocks, image_warnings = _image_blocks(
+                                image_observations, max_image_bytes
+                            )
+                            warnings.extend(image_warnings)
+                            if blocks:
+                                messages.append({"role": "user", "content": blocks})
+
+        except (TimeoutError, asyncio.TimeoutError, litellm.Timeout):
+            return AgentAttempt(
+                usage=usage.as_rewardkit(),
+                logs=_save_trace(
+                    label=label,
+                    model=model,
+                    workspace=workspace,
+                    outcome="timed_out",
+                    request=last_request,
+                    response=last_response,
+                    trace=trace,
+                    usage=usage.as_rewardkit(),
+                    warnings=warnings,
+                ),
+                warnings=warnings,
+                timed_out=True,
+            )
+        except Exception as exc:
+            trace.append(
+                {
+                    "event": "error",
+                    "error": f"{type(exc).__name__}: {exc}",
+                }
+            )
+            return AgentAttempt(
+                usage=usage.as_rewardkit(),
+                logs=_save_trace(
+                    label=label,
+                    model=model,
+                    workspace=workspace,
+                    outcome="error",
+                    request=last_request,
+                    response=last_response,
+                    trace=trace,
+                    usage=usage.as_rewardkit(),
+                    warnings=warnings,
+                ),
+                warnings=warnings,
+                error=exc,
+            )
+
+        trace.append(
+            {
+                "event": "error",
+                "error": f"maximum model turns reached: {max_steps}",
+            }
+        )
+        return AgentAttempt(
+            usage=usage.as_rewardkit(),
+            logs=_save_trace(
+                label=label,
+                model=model,
+                workspace=workspace,
+                outcome="max_steps",
+                request=last_request,
+                response=last_response,
+                trace=trace,
+                usage=usage.as_rewardkit(),
+                warnings=warnings,
+            ),
+            warnings=warnings,
+            error=RuntimeError(
+                f"react judge reached {max_steps} model turns without submit_score"
+            ),
+        )
+
+
+async def _execute_tool(
+    registry: ToolRegistry,
+    name: str,
+    arguments: dict[str, Any],
+    context: ToolContext,
+) -> ToolResult:
+    tool = registry.get(name)
+    if tool is None:
+        return ToolResult(error=f"unknown tool: {name}")
+    try:
+        return await tool.execute(arguments, context)
+    except Exception as exc:
+        return ToolResult(error=f"{type(exc).__name__}: {exc}")
+
+
+def _system_prompt(workspace: Path, max_steps: int) -> str:
+    return f"""You are an evaluation agent judging artifacts in this workspace:
+{workspace}
+
+Use the provided tools to gather concrete evidence before scoring. Prefer targeted
+inspection over reading every file. Do not modify the judged workspace. You have at
+most {max_steps} model turns. When you have enough evidence, call submit_score exactly
+once with the final scoring object. Do not merely print the score in assistant text;
+the run ends only when you call submit_score. Include concise, evidence-based reasoning
+in every criterion's reasoning field."""
+
+
+def _submit_score_tool(schema: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "type": "function",
+        "function": {
+            "name": "submit_score",
+            "description": (
+                "Submit the final RewardKit score after inspecting enough evidence. "
+                "Calling this tool ends the judge run."
+            ),
+            "parameters": schema,
+        },
+    }
+
+
+def _completion_overrides() -> dict[str, Any]:
+    result: dict[str, Any] = {}
+    base_url = os.environ.get("REACT_JUDGE_API_BASE")
+    api_key = os.environ.get("REACT_JUDGE_API_KEY")
+    if base_url:
+        result["base_url"] = base_url
+    if api_key:
+        result["api_key"] = api_key
+    return result
+
+
+def _assistant_message(content: Any, tool_calls: list[Any]) -> dict[str, Any]:
+    result: dict[str, Any] = {"role": "assistant", "content": content}
+    if tool_calls:
+        result["tool_calls"] = []
+        for call in tool_calls:
+            function = _field(call, "function")
+            result["tool_calls"].append(
+                {
+                    "id": str(_field(call, "id") or ""),
+                    "type": "function",
+                    "function": {
+                        "name": str(_field(function, "name") or ""),
+                        "arguments": _field(function, "arguments") or "{}",
+                    },
+                }
+            )
+    return result
+
+
+def _parse_tool_call(
+    call: Any,
+) -> tuple[str, str, dict[str, Any], str | None]:
+    call_id = str(_field(call, "id") or f"call-{id(call)}")
+    function = _field(call, "function")
+    name = str(_field(function, "name") or "")
+    raw_arguments = _field(function, "arguments")
+    if isinstance(raw_arguments, dict):
+        arguments = raw_arguments
+    else:
+        try:
+            arguments = json.loads(raw_arguments or "{}")
+        except (TypeError, json.JSONDecodeError) as exc:
+            return call_id, name, {}, f"invalid JSON tool arguments: {exc}"
+    if not isinstance(arguments, dict):
+        return call_id, name, {}, "tool arguments must decode to a JSON object"
+    if not name:
+        return call_id, name, arguments, "tool call has no function name"
+    return call_id, name, arguments, None
+
+
+def _field(value: Any, name: str, default: Any = None) -> Any:
+    if isinstance(value, dict):
+        return value.get(name, default)
+    return getattr(value, name, default)
+
+
+def _image_blocks(
+    observations: list[tuple[str, ToolResult]], max_bytes: int
+) -> tuple[list[dict[str, Any]], list[str]]:
+    blocks: list[dict[str, Any]] = []
+    warnings: list[str] = []
+    for tool_name, result in observations:
+        for path in result.images:
+            try:
+                data, mime_type = _read_image(path, max_bytes)
+            except Exception as exc:
+                warnings.append(f"could not attach image {path}: {exc}")
+                continue
+            blocks.append(
+                {
+                    "type": "text",
+                    "text": f"Visual observation from {tool_name}: {path.name}",
+                }
+            )
+            blocks.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:{mime_type};base64,{base64.b64encode(data).decode('ascii')}"
+                    },
+                }
+            )
+    return blocks, warnings
+
+
+def _read_image(path: Path, max_bytes: int) -> tuple[bytes, str]:
+    data = path.read_bytes()
+    mime_type = mimetypes.guess_type(path.name)[0] or "image/png"
+    if len(data) <= max_bytes:
+        return data, mime_type
+
+    try:
+        from PIL import Image
+    except ImportError as exc:
+        raise ValueError(
+            f"image is {len(data)} bytes (limit {max_bytes}) and Pillow is unavailable"
+        ) from exc
+
+    with Image.open(path) as image:
+        image.thumbnail((1800, 1800))
+        if image.mode not in {"RGB", "L"}:
+            image = image.convert("RGB")
+        output = io.BytesIO()
+        image.save(output, format="JPEG", quality=85, optimize=True)
+    data = output.getvalue()
+    if len(data) > max_bytes:
+        raise ValueError(f"resized image is still larger than {max_bytes} bytes")
+    return data, "image/jpeg"
+
+
+def _positive_int_env(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    value = int(raw)
+    if value <= 0:
+        raise ValueError(f"{name} must be positive")
+    return value
+
+
+def _shorten(value: str, limit: int) -> str:
+    return value if len(value) <= limit else value[:limit] + "...[truncated]"
+
+
+def _trace_request(request: dict[str, Any]) -> dict[str, Any]:
+    """Copy the model request for the trace while excluding credentials."""
+    return {
+        key: copy.deepcopy(value)
+        for key, value in request.items()
+        if key.lower() not in {"api_key", "authorization"}
+    }
+
+
+def _save_trace(
+    *,
+    label: str,
+    model: str,
+    workspace: Path,
+    outcome: str,
+    request: dict[str, Any] | None,
+    response: dict[str, Any] | None,
+    trace: list[dict[str, Any]],
+    usage: dict[str, Any],
+    warnings: list[str],
+) -> list[str]:
+    configured = os.environ.get("REACT_JUDGE_LOG_DIR")
+    log_dir = (
+        Path(configured).expanduser()
+        if configured
+        else Path("/logs/verifier/judge_agents_trajs")
+    )
+    safe_label = "".join(
+        char if char.isalnum() or char in "-_" else "_" for char in label
+    )
+    path = log_dir / f"trajectory.react.{safe_label}.{uuid.uuid4().hex}.json"
+    payload = {
+        "version": 2,
+        "backend": "react",
+        "label": label,
+        "model": model,
+        "workspace": str(workspace),
+        "outcome": outcome,
+        "usage": usage,
+        "request": request,
+        "response": response,
+        "events": trace,
+    }
+    try:
+        log_dir.mkdir(parents=True, exist_ok=True)
+        path.write_text(
+            json.dumps(payload, ensure_ascii=False, indent=2, default=str) + "\n"
+        )
+    except Exception as exc:
+        warnings.append(f"could not save react judge trajectory to {path}: {exc}")
+        return []
+    return [str(path)]
+
+
+class _Usage:
+    def __init__(self, model: str) -> None:
+        self.model = model
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self.reasoning_output_tokens = 0
+        self.total_tokens = 0
+        self.tool_requests: Counter[str] = Counter()
+
+    def add_response(self, response: Any) -> None:
+        raw = _field(response, "usage")
+        if raw is None:
+            return
+        self.input_tokens += int(
+            _field(raw, "prompt_tokens", _field(raw, "input_tokens", 0)) or 0
+        )
+        self.output_tokens += int(
+            _field(raw, "completion_tokens", _field(raw, "output_tokens", 0)) or 0
+        )
+        self.total_tokens += int(_field(raw, "total_tokens", 0) or 0)
+        details = _field(raw, "completion_tokens_details") or {}
+        self.reasoning_output_tokens += int(_field(details, "reasoning_tokens", 0) or 0)
+
+    def as_rewardkit(self) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.input_tokens:
+            result["input_tokens"] = self.input_tokens
+        if self.output_tokens:
+            result["output_tokens"] = self.output_tokens
+        if self.reasoning_output_tokens:
+            result["reasoning_output_tokens"] = self.reasoning_output_tokens
+        total = self.total_tokens or self.input_tokens + self.output_tokens
+        if total:
+            result["total_tokens"] = total
+        if self.tool_requests:
+            result["tool_requests"] = dict(self.tool_requests)
+        if result:
+            model_usage = {
+                key: value for key, value in result.items() if key != "models"
+            }
+            result["models"] = {self.model: model_usage}
+        return result

rewardkit_react/registry.py

@@ -0,0 +1,152 @@
+"""Tool registry and plugin discovery."""
+
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import os
+from importlib.metadata import entry_points
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Iterable
+
+from rewardkit_react.tool import Tool
+
+ENTRY_POINT_GROUP = "rewardkit_react.tools"
+
+
+class ToolRegistry:
+    def __init__(self) -> None:
+        self._tools: dict[str, Tool] = {}
+
+    def register(self, tool: Tool, *, replace: bool = False) -> None:
+        name = getattr(tool, "name", "")
+        if not isinstance(name, str) or not name:
+            raise TypeError("tool.name must be a non-empty string")
+        if name == "submit_score":
+            raise ValueError("submit_score is reserved by the agent")
+        if not callable(getattr(tool, "execute", None)):
+            raise TypeError(f"tool {name!r} has no execute() method")
+        if name in self._tools and not replace:
+            raise ValueError(f"tool already registered: {name}")
+        self._tools[name] = tool
+
+    def get(self, name: str) -> Tool | None:
+        return self._tools.get(name)
+
+    def values(self) -> tuple[Tool, ...]:
+        return tuple(self._tools[name] for name in sorted(self._tools))
+
+    def retain(self, names: set[str]) -> None:
+        unknown = names.difference(self._tools)
+        if unknown:
+            raise ValueError(f"unknown tools in REACT_JUDGE_TOOLS: {sorted(unknown)}")
+        self._tools = {
+            name: tool for name, tool in self._tools.items() if name in names
+        }
+
+
+def _materialize(candidate: Any) -> Tool:
+    if isinstance(candidate, type):
+        candidate = candidate()
+    elif callable(candidate) and not hasattr(candidate, "execute"):
+        candidate = candidate()
+    if not hasattr(candidate, "name") or not callable(
+        getattr(candidate, "execute", None)
+    ):
+        raise TypeError(
+            "plugin must be a Tool instance, Tool class, or zero-argument factory"
+        )
+    return candidate
+
+
+def _register_collection(registry: ToolRegistry, values: Iterable[Any]) -> None:
+    for value in values:
+        registry.register(_materialize(value))
+
+
+def _load_module_from_path(path: Path) -> ModuleType:
+    module_name = f"rewardkit_react_local_{abs(hash(path.resolve()))}"
+    spec = importlib.util.spec_from_file_location(module_name, path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"cannot load tool module: {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _register_module(registry: ToolRegistry, module: ModuleType) -> None:
+    hook = getattr(module, "register_tools", None)
+    if callable(hook):
+        hook(registry)
+        return
+    factory = getattr(module, "get_tools", None)
+    if callable(factory):
+        _register_collection(registry, factory())
+        return
+    tools = getattr(module, "TOOLS", None)
+    if tools is not None:
+        _register_collection(registry, tools)
+        return
+    raise ValueError(
+        f"{module.__name__} must expose register_tools(registry), get_tools(), or TOOLS"
+    )
+
+
+def load_registry() -> tuple[ToolRegistry, list[str]]:
+    """Load built-ins, installed entry points, and task-local modules."""
+    from rewardkit_react.tools import (
+        InspectFileTool,
+        ListFilesTool,
+        RenderDocumentTool,
+        SearchFilesTool,
+        ShellTool,
+        ViewTool,
+    )
+
+    registry = ToolRegistry()
+    for tool in (
+        ShellTool(),
+        ViewTool(),
+        ListFilesTool(),
+        SearchFilesTool(),
+        InspectFileTool(),
+        RenderDocumentTool(),
+    ):
+        registry.register(tool)
+
+    warnings: list[str] = []
+    for entry_point in entry_points(group=ENTRY_POINT_GROUP):
+        try:
+            registry.register(_materialize(entry_point.load()))
+        except Exception as exc:
+            warnings.append(
+                f"failed to load tool entry point {entry_point.name}: {exc}"
+            )
+
+    raw_modules = os.environ.get("REACT_JUDGE_TOOL_MODULES", "")
+    module_names = [
+        item.strip() for item in raw_modules.replace(",", os.pathsep).split(os.pathsep)
+    ]
+    for module_name in filter(None, module_names):
+        try:
+            path = Path(module_name)
+            module = (
+                _load_module_from_path(path)
+                if path.suffix == ".py" or path.exists()
+                else importlib.import_module(module_name)
+            )
+            _register_module(registry, module)
+        except Exception as exc:
+            warnings.append(f"failed to load tool module {module_name}: {exc}")
+
+    enabled = os.environ.get("REACT_JUDGE_TOOLS")
+    if enabled:
+        registry.retain({name.strip() for name in enabled.split(",") if name.strip()})
+    elif os.environ.get("REACT_JUDGE_ENABLE_SHELL", "").lower() not in {
+        "1",
+        "true",
+        "yes",
+    }:
+        registry.retain({tool.name for tool in registry.values()} - {"shell"})
+    return registry, warnings