diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b9a356..ebb6cc2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ +## [2.3.31] - 2026-10-25 + +### Added +- **`ROCmLLMChat` tool loop**: the node can now execute tool calls instead of only relaying text. `tools` picks a toolset (`none`, `read`, `read_command`, `read_write`, `full`); the node sends them in OpenAI `tools` form, runs whatever the model calls, feeds each result back as a `tool` message, and loops until the model answers or `max_tool_steps` (default 8) is reached. Tools: `read_file` (numbered lines, `offset`/`limit`), `list_dir`, `search_files` (recursive glob), `write_file`, `run_command`. File tools are confined to `path_root` (default: the ComfyUI directory; `realpath` containment, so `..` and symlinks cannot escape). `run_command` runs in a persistent bash session where `cd` and `export` carry across calls of one execution, with the ComfyUI user's permissions and no sandbox, so it is opt-in. Results are clipped to 8000 chars, commands time out after 120 s, and tool errors come back as text so the model can retry. With a toolset enabled and an empty `system_prompt` the node injects an agent system prompt naming the real host, root and enabled tools and forbidding invented file contents or command output. `thinking` now accumulates `reasoning_content` across steps. New module `rocm_nodes/core/llm_tools.py`, no new dependency. + ## [2.3.30] - 2026-10-25 ### Added diff --git a/README.md b/README.md index 5c9c2c8..6c9e57d 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,10 @@ After running: **ROCm Ninodes** is a custom node collection tuned for AMD GPUs with ROCm (especially gfx1151). It includes optimized VAE decode, KSampler, checkpoint/diffusion/GGUF/LoRA loaders, LTX2 prompt generation, SamplerCustomAdvanced drop-in, and monitoring nodes to maximize performance on AMD hardware with mature ROCm drivers. +## 🚀 What's new in v2.3.31 + +- **ROCm LLM Chat can now use tools**: pick a toolset (`read`, `read_command`, `read_write`, `full`) and the node runs the model's tool calls itself — read/list/search files, write files, and a persistent shell where `cd` and `export` carry across calls — feeding each result back until the model answers. File tools are confined to `path_root` (default: the ComfyUI directory); the shell runs with your ComfyUI user's permissions, so it is opt-in. Category: **ROCm Ninodes → Generative AI**. + ## 🚀 What's new in v2.3.30 - **ROCm LLM Chat**: chat with a local OpenAI-compatible server (Halogen, llama.cpp, vLLM, Ollama) from a node. The HTTP call is made by the ComfyUI backend, so the `server` field is the backend-side address — a remote UI can point at an LLM on the ComfyUI host or any other machine. Outputs the answer, the model's reasoning text, and a chainable conversation history. No new dependencies. Category: **ROCm Ninodes → Generative AI**. diff --git a/comfyui_manager.json b/comfyui_manager.json index df394fe..00ce819 100644 --- a/comfyui_manager.json +++ b/comfyui_manager.json @@ -2,7 +2,7 @@ "name": "rocm-ninodes", "display_name": "ROCm Ninodes", "description": "ROCm-optimized nodes for ComfyUI: VAE decode, sampling (KSampler + SamplerCustomAdvanced), checkpoint/diffusion/GGUF/LoRA loaders, LTX2 prompt generation, Seed-VC voice conversion, and monitoring — tuned for AMD GPUs (gfx1151) with ROCm.", - "version": "2.3.30", + "version": "2.3.31", "author": "iGavroche", "author_email": "nino2k@proton.me", "repository": "https://github.com/iGavroche/rocm-ninodes", diff --git a/docs/nodes/INDEX.md b/docs/nodes/INDEX.md index 7cdb528..10df7f0 100644 --- a/docs/nodes/INDEX.md +++ b/docs/nodes/INDEX.md @@ -46,7 +46,7 @@ All nodes shipped by ROCm Ninodes (`rocm_nodes/nodes.py` `NODE_CLASS_MAPPINGS` p | Class | Display name | Category | Purpose | Doc | |---|---|---|---|---| | `ROCmTextGenerateLTX2Prompt` | ROCm Text Generate LTX2 Prompt | ROCm Ninodes/Generative AI | Drop-in `TextGenerateLTX2Prompt` with device pinning and memory cleanup around the LLM call. | [ROCmTextGenerateLTX2Prompt.md](ROCmTextGenerateLTX2Prompt.md) | -| `ROCmLLMChat` | ROCm LLM Chat | ROCm Ninodes/Generative AI | Chat with a local OpenAI-compatible server (Halogen, llama.cpp, vLLM, Ollama) from the backend; answer, reasoning and chainable history outputs. | [ROCmLLMChat.md](ROCmLLMChat.md) | +| `ROCmLLMChat` | ROCm LLM Chat | ROCm Ninodes/Generative AI | Chat with a local OpenAI-compatible server (Halogen, llama.cpp, vLLM, Ollama) from the backend; answer, reasoning and chainable history outputs, plus an optional file/shell tool loop. | [ROCmLLMChat.md](ROCmLLMChat.md) | ## Attention and step caching diff --git a/docs/nodes/ROCmLLMChat.md b/docs/nodes/ROCmLLMChat.md index 71dbc2a..b455f1e 100644 --- a/docs/nodes/ROCmLLMChat.md +++ b/docs/nodes/ROCmLLMChat.md @@ -25,6 +25,9 @@ The request is issued by the **ComfyUI backend process**, so the `server` addres | `seed` | INT | `0` | 0 … 0xFFFFFFFFFFFFFFFF | Change it to get a new answer for an unchanged prompt (ComfyUI caches identical inputs). | | `api_key` | STRING | `""` | — | Sent as `Authorization: Bearer `; falls back to `ROCM_LLM_API_KEY`. | | `timeout` | FLOAT | `300.0` | 1.0 … 3600.0 | Seconds to wait for the response, prompt processing included. | +| `tools` | combo | `none` | `none`, `read`, `read_command`, `read_write`, `full` | Toolset the node executes for the model. See [Tools](#tools). | +| `path_root` | STRING | `""` | — | Working root for file tools and commands. Empty = the ComfyUI directory; `~` is expanded. Paths outside it are rejected, `/` grants the whole filesystem. | +| `max_tool_steps` | INT | `8` | 1 … 64 | How many model turns may use tools before the node must answer. | `FUNCTION = chat`. `RETURN_TYPES = ("STRING", "STRING", "CHAT_HISTORY")`, `RETURN_NAMES = ("response", "thinking", "chat_history")`. @@ -33,8 +36,27 @@ The request is issued by the **ComfyUI backend process**, so the `server` addres | Name | Type | Notes | |---|---|---| | `response` | STRING | `choices[0].message.content`, stripped. Feed into `CLIPTextEncode`, a text node, or the next chat node. | -| `thinking` | STRING | `choices[0].message.reasoning_content` when the server emits reasoning, else empty. | -| `chat_history` | CHAT_HISTORY | The messages actually sent, plus this reply. Chain it into another node's `chat_history` for multi-turn. | +| `thinking` | STRING | `reasoning_content` from every step, joined with a blank line; empty when the server emits no reasoning. | +| `chat_history` | CHAT_HISTORY | The messages actually sent — with a toolset this includes the `assistant`/`tool` tool turns — plus this reply. Chain it into another node's `chat_history` for multi-turn. | + +## Tools + +With a toolset other than `none`, the node sends the tools in OpenAI `tools` form and executes whatever the model calls, feeding each result back as a `tool` message until the model answers or `max_tool_steps` is reached. Each executed call is logged as `[ROCm LLM Chat] -> N chars`. + +| Toolset | Tools | +|---|---| +| `read` | `read_file`, `list_dir`, `search_files` | +| `read_command` | `read` + `run_command` | +| `read_write` | `read` + `write_file` | +| `full` | all five, including `run_command` and `write_file` | + +`read_file` returns up to 400 numbered lines from `offset`; `list_dir` lists one level (directories suffixed `/`); `search_files` globs file names recursively; `write_file` replaces a file wholesale and creates parents; `run_command` runs a shell command and returns `exit ` plus merged stdout/stderr. Results are clipped to 8000 characters, commands time out after 120 s. + +`run_command` is a **persistent bash session**: `cd` and `export` carry over between calls of one node execution, and the session starts in `path_root`. It runs with the ComfyUI user's permissions and no sandbox — anything the user running ComfyUI can do, the model can do. Enable it only for prompts you control. + +File tools are confined to `path_root` (`realpath` containment, so `..` and symlinks cannot escape); `run_command` is not — the containment applies to file tools only. + +When `tools` is not `none` and `system_prompt` is empty, the node injects an agent system prompt that states the real host, root and enabled tools, and requires answers to come from tool results instead of invented file contents. A non-empty `system_prompt` replaces it entirely. ## Multi-turn @@ -44,6 +66,8 @@ There is no loop in a ComfyUI graph, so a conversation is a chain of nodes: each ROCmLLMChat (prompt "…") -> chat_history -> ROCmLLMChat (prompt "…") -> chat_history -> … ``` +A history produced with a toolset contains `tool` messages, so the node that consumes it must have a `tools` value that enables the same tools, or the server will reject the replay. + ## gfx1151 / deployment notes - No GPU work in ComfyUI: inference happens in the LLM server process. On a shared machine that means the LLM and the diffusion model contend for the same memory bus — expect slower sampling while a long chat generation runs. diff --git a/package.json b/package.json index ef9fa9a..66d2fa6 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "rocm-ninodes", - "version": "2.3.30", + "version": "2.3.31", "description": "ROCM-optimized ComfyUI nodes for AMD GPU performance", "author": "Nino", "email": "nino2k@proton.me", @@ -25,7 +25,7 @@ "display_name": "RocM Ninodes", "description": "ROCM-optimized nodes for AMD GPU performance", "author": "Nino", - "version": "2.3.30", + "version": "2.3.31", "category": "Optimization", "tags": ["rocm", "amd", "gpu", "optimization", "vae", "sampling"], "nodes": [ diff --git a/pyproject.toml b/pyproject.toml index a42389b..65d7da7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rocm-ninodes" description = "ROCM Optimized Nodes for ComfyUI - High-performance VAE decode and sampling nodes specifically tuned for AMD GPUs with ROCm support, particularly targeting gfx1151 architecture" -version = "2.3.30" +version = "2.3.31" license = {file = "LICENSE"} dependencies = [ "numpy>=1.20.0", @@ -28,7 +28,7 @@ requires-comfyui = ">=0.0.1" name = "rocm-ninodes" display_name = "ROCm Ninodes" description = "ROCm-optimized nodes for ComfyUI: VAE decode, sampling, SamplerCustomAdvanced drop-in, loaders, LTX2 prompt generation, and monitoring — tuned for AMD GPUs (gfx1151) with ROCm." -version = "2.3.30" +version = "2.3.31" author = "iGavroche" author_email = "nino2k@proton.me" repository = "https://github.com/iGavroche/rocm-ninodes" diff --git a/rocm_nodes/core/llm_chat.py b/rocm_nodes/core/llm_chat.py index 31531a3..55b1701 100644 --- a/rocm_nodes/core/llm_chat.py +++ b/rocm_nodes/core/llm_chat.py @@ -8,10 +8,54 @@ typing 127.0.0.1 reaches the ComfyUI host, not their own laptop. import json import os +import platform import urllib.error import urllib.request from typing import Any, Optional +from .llm_tools import TOOLSETS, ShellSession, run_tool, tool_payload + + +def default_root() -> str: + """Working root for tools: the ComfyUI install directory, else the process cwd.""" + try: + import folder_paths + return folder_paths.base_path + except ImportError: + return os.getcwd() + + +def resolve_root(path: str) -> str: + """Absolute working root; empty input means the ComfyUI directory.""" + return os.path.realpath(os.path.expanduser((path or "").strip() or default_root())) + + +AGENT_SYSTEM_PROMPT = """You are a technical assistant running inside a ComfyUI workflow, with tools. + +Environment: {system}. Your working root is {root}: relative paths resolve there and absolute paths outside it are rejected. Available tools: {tools}. + +Rules: +- Answer from tool results. NEVER invent file contents, command output, or paths you did not observe. +- Read before you claim. If a tool call errors, adapt (list_dir, search_files) instead of guessing a path. +- Make the narrowest call that answers the question, and put independent reads in the same step. +- Never report an action you did not perform, and never report a result you did not verify. +- Be terse: answer first, no preamble or restatement. Code in one runnable block, no narration unless asked. +- Ambiguous request: take the most plausible reading, name it in one line, continue. +- Stop as soon as the question is answered; do not explore past it. +{shell} +""" + + +def agent_system_prompt(root: str, tools: list[str]) -> str: + """System prompt describing the real root and the tools actually enabled.""" + system = f"{platform.system()} {platform.release()} ({platform.machine()})" + shell = "" + if "run_command" in tools: + shell = ("- run_command is a persistent shell: cd and export carry over between calls, so " + "run one command per call and check its exit code before continuing.") + return AGENT_SYSTEM_PROMPT.format(system=system, root=root, + tools=", ".join(tools), shell=shell) + # Halogen flash-server default; override per machine with ROCM_LLM_SERVER. DEFAULT_SERVER = "http://127.0.0.1:8731/v1" @@ -125,6 +169,22 @@ class ROCmLLMChat: "default": 300.0, "min": 1.0, "max": 3600.0, "step": 1.0, "tooltip": "Seconds to wait for the response, including prompt processing.", }), + "tools": (list(TOOLSETS), { + "default": "none", + "tooltip": "Toolset the node executes for the model. 'full' adds shell " + "commands run on the ComfyUI host; only enable it for prompts " + "you control.", + }), + "path_root": ("STRING", { + "default": "", + "tooltip": "Working root for file tools and commands. Empty = the ComfyUI " + "directory. Paths outside it are rejected; use '/' for full " + "filesystem access.", + }), + "max_tool_steps": ("INT", { + "default": 8, "min": 1, "max": 64, + "tooltip": "How many model turns may use tools before the node must answer.", + }), }, } @@ -136,14 +196,21 @@ class ROCmLLMChat: "Chat with a local OpenAI-compatible LLM server (Halogen, llama.cpp, vLLM, Ollama). " "The request is sent from the ComfyUI backend, so the server address is the backend's " "view of the network. Outputs the answer, the model's reasoning text when the server " - "provides it, and the conversation history for multi-turn chaining." + "provides it, and the conversation history for multi-turn chaining. With a toolset " + "enabled the node also runs the model's file and shell tool calls on the ComfyUI host." ) def chat(self, server: str, prompt: str, chat_history: Optional[list[dict[str, str]]] = None, system_prompt: str = "", model: str = "", max_tokens: int = 2048, temperature: float = 0.7, top_p: float = 0.95, seed: int = 0, - api_key: str = "", timeout: float = 300.0) -> tuple[str, str, list[dict[str, str]]]: + api_key: str = "", timeout: float = 300.0, tools: str = "none", + path_root: str = "", max_tool_steps: int = 8 + ) -> tuple[str, str, list[dict[str, Any]]]: endpoint = chat_endpoint(server) + enabled = TOOLSETS[tools] + root = resolve_root(path_root) + if enabled and not system_prompt.strip(): + system_prompt = agent_system_prompt(root, enabled) messages = build_messages(prompt, system_prompt, chat_history) payload: dict[str, Any] = { "messages": messages, @@ -155,22 +222,47 @@ class ROCmLLMChat: } if model: payload["model"] = model + if enabled: + payload["tools"] = tool_payload(enabled) api_key = api_key or os.environ.get("ROCM_LLM_API_KEY", "") - data = request_chat_completion(endpoint, payload, api_key, timeout) - choices = data.get("choices") or [] - if not choices: - raise ValueError(f"LLM server {endpoint} returned no choices: {str(data)[:300]}") - message = choices[0].get("message") or {} - response = (message.get("content") or "").strip() - thinking = (message.get("reasoning_content") or "").strip() - served = data.get("model") or "" - if model and served and served != model: - print(f"[ROCm LLM Chat] requested model '{model}', server answered as " - f"'{served}' - check the name against GET /v1/models.") + steps = max_tool_steps if enabled else 1 + session = ShellSession(root) + thinking: list[str] = [] + response = "" + for step in range(steps): + data = request_chat_completion(endpoint, payload, api_key, timeout) + choices = data.get("choices") or [] + if not choices: + raise ValueError(f"LLM server {endpoint} returned no choices: {str(data)[:300]}") + served = data.get("model") or "" + if step == 0 and model and served and served != model: + print(f"[ROCm LLM Chat] requested model '{model}', server answered as " + f"'{served}' - check the name against GET /v1/models.") + message = choices[0].get("message") or {} + reasoning = (message.get("reasoning_content") or "").strip() + if reasoning: + thinking.append(reasoning) + calls = message.get("tool_calls") or [] + response = (message.get("content") or "").strip() + if not calls: + break + if step == steps - 1: + print(f"[ROCm LLM Chat] the model asked for tools but no step was left " + f"(max_tool_steps={steps}); raise it.") + break + messages.append({"role": "assistant", "content": message.get("content") or "", + "tool_calls": calls}) + for call in calls: + function = call.get("function") or {} + name = function.get("name") or "" + result = run_tool(name, function.get("arguments", ""), root, enabled, session) + print(f"[ROCm LLM Chat] {name} -> {len(result)} chars") + messages.append({"role": "tool", "tool_call_id": call.get("id") or "", + "content": result}) + payload["messages"] = messages + if not response and thinking: - print(f"[ROCm LLM Chat] empty answer, all output went to reasoning " - f"(finish_reason={choices[0].get('finish_reason')}); raise max_tokens.") - + print("[ROCm LLM Chat] empty answer, all output went to reasoning; raise max_tokens.") history = messages + [{"role": "assistant", "content": response}] - return (response, thinking, history) + return (response, "\n\n".join(thinking), history) diff --git a/rocm_nodes/core/llm_tools.py b/rocm_nodes/core/llm_tools.py new file mode 100644 index 0000000..d0eeea4 --- /dev/null +++ b/rocm_nodes/core/llm_tools.py @@ -0,0 +1,253 @@ +""" +Filesystem and command tools executed by the ROCm LLM Chat agent loop. + +Every tool returns plain text, the way a terminal would show it, so the model can +read the result back directly. File tools are confined to a root directory and +the command tool is only reachable when the caller opted into it. +""" + +import fnmatch +import json +import os +import shlex +import shutil +import subprocess +from typing import Optional + + +MAX_RESULT_CHARS = 8000 +COMMAND_TIMEOUT = 120 +LIST_LIMIT = 500 +SEARCH_LIMIT = 200 +READ_LINES = 400 + + +TOOL_SPECS = { + "read_file": { + "description": "Read a text file as lines. Returns at most `limit` lines starting at `offset` (1-based).", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string", "description": "File path, absolute or relative to the working root."}, + "offset": {"type": "integer", "description": "First line to read, 1-based. Default 1."}, + "limit": {"type": "integer", "description": f"Lines to read. Default {READ_LINES}."}, + }, + "required": ["path"], + }, + }, + "list_dir": { + "description": "List one directory level. Directories end with '/', files show their size.", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string", "description": "Directory path. Default '.' (the working root)."}, + }, + }, + }, + "search_files": { + "description": f"Find files by glob pattern on the file name, recursively. Returns up to {SEARCH_LIMIT} paths.", + "parameters": { + "type": "object", + "properties": { + "pattern": {"type": "string", "description": "Glob matched against the file name, e.g. '*.json'."}, + "path": {"type": "string", "description": "Directory to search under. Default '.' (the working root)."}, + }, + "required": ["pattern"], + }, + }, + "write_file": { + "description": "Write text to a file, replacing it entirely. Creates parent directories.", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string", "description": "File path, absolute or relative to the working root."}, + "content": {"type": "string", "description": "Full file contents."}, + }, + "required": ["path", "content"], + }, + }, + "run_command": { + "description": ( + "Run a shell command in a persistent bash session and return stdout, stderr and the " + f"exit code. The working directory and exported variables carry over between calls of " + f"this run; the session starts in the working root. Times out after {COMMAND_TIMEOUT}s." + ), + "parameters": { + "type": "object", + "properties": { + "command": {"type": "string", "description": "Shell command line."}, + }, + "required": ["command"], + }, + }, +} + + +TOOLSETS = { + "none": [], + "read": ["read_file", "list_dir", "search_files"], + "read_command": ["read_file", "list_dir", "search_files", "run_command"], + "read_write": ["read_file", "list_dir", "search_files", "write_file"], + "full": ["read_file", "list_dir", "search_files", "write_file", "run_command"], +} + + +def tool_payload(names: list[str]) -> list[dict]: + """OpenAI `tools` array for the enabled tool names.""" + return [{"type": "function", "function": {"name": name, **TOOL_SPECS[name]}} for name in names] + + +def resolve_path(path: str, root: str) -> str: + """Absolute path for `path`, rejecting anything outside `root`.""" + root = os.path.realpath(root) + candidate = path if os.path.isabs(path) else os.path.join(root, path) + target = os.path.realpath(candidate) + if target != root and not target.startswith(root + os.sep): + raise ValueError(f"'{path}' resolves outside the allowed root '{root}'") + return target + + +def _clip(text: str) -> str: + if len(text) <= MAX_RESULT_CHARS: + return text + return text[:MAX_RESULT_CHARS] + f"\n... [{len(text) - MAX_RESULT_CHARS} more chars truncated]" + + +def _read_file(path: str, offset: int = 1, limit: int = READ_LINES, *, root: str) -> str: + target = resolve_path(path, root) + if not os.path.isfile(target): + return f"error: not a file: {target}" + with open(target, "r", encoding="utf-8", errors="replace") as handle: + lines = handle.readlines() + start = max(1, int(offset)) + count = max(1, int(limit)) + chunk = lines[start - 1:start - 1 + count] + if not chunk: + return f"{target}: {len(lines)} lines total, nothing at offset {start}" + header = f"{target} lines {start}-{start + len(chunk) - 1} of {len(lines)}" + numbered = "\n".join(f"{start + i}:{line.rstrip(chr(10))}" for i, line in enumerate(chunk)) + return _clip(f"{header}\n{numbered}") + + +def _list_dir(path: str = ".", *, root: str) -> str: + target = resolve_path(path, root) + if not os.path.isdir(target): + return f"error: not a directory: {target}" + entries = sorted(os.listdir(target)) + lines = [] + for name in entries[:LIST_LIMIT]: + full = os.path.join(target, name) + if os.path.isdir(full): + lines.append(f"{name}/") + else: + lines.append(f"{name} {os.path.getsize(full)}B") + if len(entries) > LIST_LIMIT: + lines.append(f"... [{len(entries) - LIST_LIMIT} more entries]") + return _clip(f"{target}: {len(entries)} entries\n" + "\n".join(lines)) + + +def _search_files(pattern: str, path: str = ".", *, root: str) -> str: + start = resolve_path(path, root) + if not os.path.isdir(start): + return f"error: not a directory: {start}" + matches = [] + for dirpath, dirnames, filenames in os.walk(start): + dirnames[:] = [d for d in dirnames if d not in (".git", "__pycache__", "node_modules", ".venv")] + for name in filenames: + if fnmatch.fnmatch(name, pattern): + matches.append(os.path.relpath(os.path.join(dirpath, name), start)) + if len(matches) >= SEARCH_LIMIT: + return _clip("\n".join(matches) + f"\n... [capped at {SEARCH_LIMIT}]") + if not matches: + return f"no files matching '{pattern}' under {start}" + return _clip("\n".join(matches)) + + +def _write_file(path: str, content: str, *, root: str) -> str: + target = resolve_path(path, root) + if os.path.isdir(target): + return f"error: is a directory: {target}" + parent = os.path.dirname(target) + if parent: + os.makedirs(parent, exist_ok=True) + with open(target, "w", encoding="utf-8") as handle: + handle.write(content) + return f"wrote {len(content)} chars to {target}" + + +class ShellSession: + """Working directory and exported variables carried across run_command calls.""" + + def __init__(self, cwd: str): + self.cwd = cwd + self.env: Optional[dict[str, str]] = None + + +_STATE_MARK = "__ROCM_LLM_STATE__" +_STATE_SCRIPT = ( + "\n} 2>&1\n" + "__rc=$?\n" + f"printf '\\n{_STATE_MARK}%s\\ncwd=%s\\n' \"$__rc\" \"$(pwd)\"\n" + "export -p\n" +) + + +def _run_command(command: str, *, session: ShellSession) -> str: + # The command runs in a brace group so its stderr joins stdout before the state + # block, and cd/export in it still affect the persistent session. + completed = subprocess.run("{\n" + command + _STATE_SCRIPT, shell=True, cwd=session.cwd, + env=session.env, capture_output=True, text=True, + timeout=COMMAND_TIMEOUT, executable=shutil.which("bash")) + out = (completed.stdout or "") + (completed.stderr or "") + head, found, state = out.rpartition(_STATE_MARK) + if not found: + return _clip(f"exit {completed.returncode}\n{out.strip() or '[no output]'}") + session.cwd, session.env, exit_code = _parse_state(state) + return _clip(f"exit {exit_code}\n{head.strip() or '[no output]'}") + + +def _parse_state(state: str) -> tuple[str, dict[str, str], int]: + """Split the sentinel block into cwd, exported environment and exit code.""" + lines = state.splitlines() + exit_code = int(lines[0]) if lines and lines[0].lstrip("-").isdigit() else 0 + cwd = "" + env: dict[str, str] = {} + for line in lines[1:]: + if line.startswith("cwd="): + cwd = line[4:].strip() + elif line.startswith("declare -x "): + name, _, value = line[len("declare -x "):].partition("=") + env[name] = shlex.split(value)[0] if value else "" + return cwd or ".", env, exit_code + + +_TOOLS = { + "read_file": _read_file, + "list_dir": _list_dir, + "search_files": _search_files, + "write_file": _write_file, + "run_command": _run_command, +} + + +def run_tool(name: str, arguments, root: str, enabled: list[str], + session: Optional[ShellSession] = None) -> str: + """Execute one tool call and return its text result, errors included.""" + if name not in enabled: + return f"error: tool '{name}' is not enabled on this node" + try: + args = json.loads(arguments) if isinstance(arguments, str) else dict(arguments or {}) + except (json.JSONDecodeError, TypeError, ValueError) as e: + return f"error: invalid JSON arguments for {name}: {e}" + if name == "run_command": + context = {"session": session or ShellSession(root)} + else: + context = {"root": root} + try: + return _TOOLS[name](**context, **args) + except TypeError as e: + return f"error: bad arguments for {name}: {e}" + except subprocess.TimeoutExpired: + return f"error: {name} timed out after {COMMAND_TIMEOUT}s" + except Exception as e: + return f"error: {type(e).__name__}: {e}" diff --git a/tests/conftest.py b/tests/conftest.py index 966bb78..a7bc8f9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -106,6 +106,7 @@ comfy.utils.ProgressBar = MockProgressBar # Need to mock folder_paths for constants import folder_paths = types.ModuleType('folder_paths') +folder_paths.base_path = '/fake' folder_paths.get_full_path_or_raise = lambda folder, name: os.path.join('/fake', folder, name) folder_paths.get_filename_list = lambda folder: [] folder_paths.get_folder_paths = lambda folder: ['/fake/' + folder] diff --git a/tests/unit/test_llm_chat.py b/tests/unit/test_llm_chat.py index 88d50fd..abe4c07 100644 --- a/tests/unit/test_llm_chat.py +++ b/tests/unit/test_llm_chat.py @@ -36,6 +36,7 @@ class Handler(BaseHTTPRequestHandler): requests = [] status = 200 raw_body = None + scripted = [] def do_POST(self): length = int(self.headers.get("Content-Length", 0)) @@ -45,7 +46,9 @@ class Handler(BaseHTTPRequestHandler): "payload": json.loads(body), "authorization": self.headers.get("Authorization"), }) - if Handler.raw_body is not None: + if Handler.scripted: + payload = json.dumps(Handler.scripted.pop(0)) + elif Handler.raw_body is not None: payload = Handler.raw_body elif Handler.status != 200: payload = json.dumps({"error": {"message": "bad request", "type": "invalid_request_error"}}) @@ -76,6 +79,7 @@ def reset_handler(): Handler.requests = [] Handler.status = 200 Handler.raw_body = None + Handler.scripted = [] yield @@ -166,7 +170,10 @@ class TestNodeSchema: inp = ROCmLLMChat.INPUT_TYPES() assert set(inp["required"]) == {"server", "prompt"} assert {"chat_history", "system_prompt", "model", "max_tokens", "temperature", - "top_p", "seed", "api_key", "timeout"} <= set(inp["optional"]) + "top_p", "seed", "api_key", "timeout", "tools", "path_root", + "max_tool_steps"} <= set(inp["optional"]) + assert inp["optional"]["tools"][0] == ["none", "read", "read_command", + "read_write", "full"] class TestChat: @@ -215,3 +222,84 @@ class TestChat: assert response == "" assert thinking == "thinking" assert history[-1] == {"role": "assistant", "content": ""} + + +def tool_call_message(name, arguments, call_id="call_1"): + return {"role": "assistant", "content": None, "tool_calls": [{ + "id": call_id, "type": "function", + "function": {"name": name, "arguments": json.dumps(arguments)}}]} + + +def completion(message, finish_reason="stop"): + return {"model": "m", "choices": [{"index": 0, "finish_reason": finish_reason, + "message": message}]} + + +class TestAgentLoop: + def test_tools_none_keeps_single_request(self, server): + ROCmLLMChat().chat(server, "hello", tools="none") + assert len(Handler.requests) == 1 + assert "tools" not in Handler.requests[0]["payload"] + assert Handler.requests[0]["payload"]["messages"][0]["role"] == "user" + + def test_tool_call_is_executed_and_fed_back(self, server, tmp_path): + (tmp_path / "a.txt").write_text("first\nsecond\n") + Handler.scripted = [ + completion(tool_call_message("read_file", {"path": "a.txt"}), "tool_calls"), + completion({"role": "assistant", "content": "two lines"}), + ] + response, _, history = ROCmLLMChat().chat( + server, "read a.txt", tools="read", path_root=str(tmp_path)) + assert response == "two lines" + assert len(Handler.requests) == 2 + assert Handler.requests[0]["payload"]["tools"][0]["function"]["name"] == "read_file" + second = Handler.requests[1]["payload"]["messages"] + assert [m["role"] for m in second] == ["system", "user", "assistant", "tool"] + assert second[1] == {"role": "user", "content": "read a.txt"} + assert second[2]["tool_calls"][0]["function"]["name"] == "read_file" + assert second[3]["tool_call_id"] == "call_1" + assert "second" in second[3]["content"] + assert history[-1] == {"role": "assistant", "content": "two lines"} + + def test_disabled_tool_error_is_returned_to_model(self, server, tmp_path): + Handler.scripted = [ + completion(tool_call_message("run_command", {"command": "id"}), "tool_calls"), + completion({"role": "assistant", "content": "blocked"}), + ] + response, _, _ = ROCmLLMChat().chat(server, "whoami", tools="read", + path_root=str(tmp_path)) + assert response == "blocked" + assert "not enabled" in Handler.requests[1]["payload"]["messages"][3]["content"] + + def test_loop_stops_at_max_tool_steps(self, server, tmp_path): + Handler.scripted = [completion(tool_call_message("list_dir", {}), "tool_calls")] + response, _, _ = ROCmLLMChat().chat(server, "ls", tools="read", + path_root=str(tmp_path), max_tool_steps=1) + assert response == "" + assert len(Handler.requests) == 1 + + def test_agent_system_prompt_describes_root_and_tools(self, server, tmp_path): + Handler.scripted = [completion({"role": "assistant", "content": "ok"})] + ROCmLLMChat().chat(server, "ls", tools="full", path_root=str(tmp_path)) + system = Handler.requests[0]["payload"]["messages"][0] + assert system["role"] == "system" + assert str(tmp_path) in system["content"] + assert "run_command" in system["content"] + + def test_explicit_system_prompt_replaces_agent_prompt(self, server, tmp_path): + Handler.scripted = [completion({"role": "assistant", "content": "ok"})] + ROCmLLMChat().chat(server, "ls", tools="full", path_root=str(tmp_path), + system_prompt="custom") + assert Handler.requests[0]["payload"]["messages"][0] == { + "role": "system", "content": "custom"} + + def test_thinking_is_accumulated_across_steps(self, server, tmp_path): + first = tool_call_message("list_dir", {}) + first["reasoning_content"] = "need a listing" + Handler.scripted = [ + completion(first, "tool_calls"), + completion({"role": "assistant", "content": "done", + "reasoning_content": "listing read"}), + ] + _, thinking, _ = ROCmLLMChat().chat(server, "ls", tools="read", path_root=str(tmp_path)) + assert thinking == "need a listing\n\nlisting read" diff --git a/tests/unit/test_llm_tools.py b/tests/unit/test_llm_tools.py new file mode 100644 index 0000000..9fbd26a --- /dev/null +++ b/tests/unit/test_llm_tools.py @@ -0,0 +1,130 @@ +""" +Unit tests for the file and shell tools executed by the ROCm LLM Chat agent loop. +""" +import json +import os +import sys + +import pytest + +sys.path.insert(0, "/home/nino/ComfyUI/custom_nodes/rocm-ninodes") + +from rocm_nodes.core.llm_tools import ( + ShellSession, + TOOLSETS, + resolve_path, + run_tool, + tool_payload, +) + +ALL = TOOLSETS["full"] + + +def call(name, arguments, root, enabled=ALL, session=None): + return run_tool(name, json.dumps(arguments), root, enabled, session) + + +class TestContainment: + def test_relative_path_resolves_under_root(self, tmp_path): + assert resolve_path("sub/f.txt", str(tmp_path)) == str(tmp_path / "sub/f.txt") + + def test_parent_escape_is_rejected(self, tmp_path): + with pytest.raises(ValueError, match="outside the allowed root"): + resolve_path("../outside.txt", str(tmp_path)) + + def test_absolute_path_outside_root_is_rejected(self, tmp_path): + with pytest.raises(ValueError, match="outside the allowed root"): + resolve_path("/etc/passwd", str(tmp_path)) + + def test_symlink_escape_is_rejected(self, tmp_path): + link = tmp_path / "link" + link.symlink_to("/etc") + with pytest.raises(ValueError, match="outside the allowed root"): + resolve_path("link/passwd", str(tmp_path)) + + def test_root_itself_is_allowed(self, tmp_path): + assert resolve_path(".", str(tmp_path)) == os.path.realpath(str(tmp_path)) + + +class TestFileTools: + def test_read_file_numbers_lines_and_pages(self, tmp_path): + (tmp_path / "f.txt").write_text("a\nb\nc\n") + assert "2:b" in call("read_file", {"path": "f.txt", "offset": 2}, str(tmp_path)) + assert "3 lines total" in call("read_file", {"path": "f.txt", "offset": 9}, str(tmp_path)) + + def test_read_file_on_directory_is_an_error(self, tmp_path): + assert "error" in call("read_file", {"path": "."}, str(tmp_path)) + + def test_list_dir_marks_directories(self, tmp_path): + (tmp_path / "d").mkdir() + (tmp_path / "f.txt").write_text("x") + out = call("list_dir", {}, str(tmp_path)) + assert "d/" in out and "f.txt" in out + + def test_search_files_matches_recursively(self, tmp_path): + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "a.json").write_text("{}") + assert "sub/a.json" in call("search_files", {"pattern": "*.json"}, str(tmp_path)) + assert "no files matching" in call("search_files", {"pattern": "*.xyz"}, str(tmp_path)) + + def test_write_file_creates_parents(self, tmp_path): + out = call("write_file", {"path": "a/b/c.txt", "content": "hello"}, str(tmp_path)) + assert "wrote 5 chars" in out + assert (tmp_path / "a/b/c.txt").read_text() == "hello" + + def test_write_file_cannot_escape(self, tmp_path): + assert "error" in call("write_file", {"path": "../x.txt", "content": "x"}, str(tmp_path)) + assert not (tmp_path.parent / "x.txt").exists() + + +class TestRunCommand: + def test_exit_code_and_output(self, tmp_path): + assert call("run_command", {"command": "echo hi"}, str(tmp_path)).startswith("exit 0\nhi") + assert call("run_command", {"command": "exit 7"}, str(tmp_path)).startswith("exit 7") + + def test_stderr_is_returned(self, tmp_path): + assert "boom" in call("run_command", {"command": "echo boom >&2"}, str(tmp_path)) + + def test_working_directory_persists(self, tmp_path): + (tmp_path / "sub").mkdir() + session = ShellSession(str(tmp_path)) + call("run_command", {"command": "cd sub"}, str(tmp_path), session=session) + assert str(tmp_path / "sub") in call("run_command", {"command": "pwd"}, + str(tmp_path), session=session) + + def test_exported_variable_persists(self, tmp_path): + session = ShellSession(str(tmp_path)) + call("run_command", {"command": "export ROCM_TEST=42"}, str(tmp_path), session=session) + assert "42" in call("run_command", {"command": "echo $ROCM_TEST"}, + str(tmp_path), session=session) + + def test_long_output_is_truncated(self, tmp_path): + out = call("run_command", {"command": "seq 1 4000"}, str(tmp_path)) + assert "more chars truncated" in out + + def test_timeout_is_reported(self, tmp_path, monkeypatch): + import rocm_nodes.core.llm_tools as tools + monkeypatch.setattr(tools, "COMMAND_TIMEOUT", 1) + assert "timed out" in call("run_command", {"command": "sleep 30"}, str(tmp_path)) + + +class TestDispatch: + def test_disabled_tool_is_refused(self, tmp_path): + out = call("run_command", {"command": "id"}, str(tmp_path), enabled=TOOLSETS["read"]) + assert "is not enabled" in out + + def test_invalid_json_arguments(self, tmp_path): + assert "invalid JSON" in run_tool("list_dir", "{", str(tmp_path), ALL) + + def test_unknown_argument_is_reported(self, tmp_path): + assert "bad arguments" in call("list_dir", {"nope": 1}, str(tmp_path)) + + def test_unknown_tool_is_reported(self, tmp_path): + assert "is not enabled" in run_tool("nope", {}, str(tmp_path), ALL) + + def test_payload_shape(self): + payload = tool_payload(TOOLSETS["read"]) + assert [p["function"]["name"] for p in payload] == ["read_file", "list_dir", + "search_files"] + assert payload[0]["type"] == "function" + assert "path" in payload[0]["function"]["parameters"]["properties"]