From 6cf416c12881344c94bd8a681e116930bf6de005 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Fri, 18 Oct 2024 13:00:17 -0700 Subject: [PATCH] feat: Add synopisis core loop (#166) Co-authored-by: Michael Neale --- packages/exchange/src/exchange/content.py | 13 ++ packages/exchange/src/exchange/message.py | 4 + pyproject.toml | 4 + src/goose/cli/session.py | 69 +++--- src/goose/synopsis/developer.md | 32 +++ src/goose/synopsis/moderator.py | 103 +++++++++ src/goose/synopsis/plan.md | 59 ++++++ src/goose/synopsis/summarize.md | 26 +++ src/goose/synopsis/synopsis.md | 25 +++ src/goose/synopsis/system.py | 148 +++++++++++++ src/goose/synopsis/toolkit.py | 244 ++++++++++++++++++++++ src/goose/toolkit/developer.py | 109 +--------- src/goose/toolkit/prompts/developer.jinja | 64 ++---- src/goose/toolkit/utils.py | 8 +- src/goose/utils/check_shell_command.py | 34 --- src/goose/utils/session_file.py | 24 +-- src/goose/utils/shell.py | 149 +++++++++++++ tests/cli/test_session.py | 60 +----- tests/synopsis/test_moderator.py | 33 +++ tests/synopsis/test_process_management.py | 62 ++++++ tests/synopsis/test_system.py | 93 +++++++++ tests/synopsis/test_toolkit.py | 95 +++++++++ tests/utils/test_check_shell_command.py | 2 +- tests/utils/test_session_file.py | 39 ---- 24 files changed, 1166 insertions(+), 333 deletions(-) create mode 100644 src/goose/synopsis/developer.md create mode 100644 src/goose/synopsis/moderator.py create mode 100644 src/goose/synopsis/plan.md create mode 100644 src/goose/synopsis/summarize.md create mode 100644 src/goose/synopsis/synopsis.md create mode 100644 src/goose/synopsis/system.py create mode 100644 src/goose/synopsis/toolkit.py delete mode 100644 src/goose/utils/check_shell_command.py create mode 100644 src/goose/utils/shell.py create mode 100644 tests/synopsis/test_moderator.py create mode 100644 tests/synopsis/test_process_management.py create mode 100644 tests/synopsis/test_system.py create mode 100644 tests/synopsis/test_toolkit.py diff --git a/packages/exchange/src/exchange/content.py b/packages/exchange/src/exchange/content.py index 66957b7c6e..e5b7ff0576 100644 --- a/packages/exchange/src/exchange/content.py +++ b/packages/exchange/src/exchange/content.py @@ -1,5 +1,6 @@ from typing import Optional +import json from attrs import define, asdict @@ -21,6 +22,10 @@ class Content: class Text(Content): text: str + @property + def summary(self) -> str: + return "content:text\n" + self.text + @define class ToolUse(Content): @@ -30,9 +35,17 @@ class ToolUse(Content): is_error: bool = False error_message: Optional[str] = None + @property + def summary(self) -> str: + return f"content:tool_use:{self.name}\nparameters:{json.dumps(self.parameters)}" + @define class ToolResult(Content): tool_use_id: str output: str is_error: bool = False + + @property + def summary(self) -> str: + return f"content:tool_result:error={self.is_error}\noutput:{self.output}" diff --git a/packages/exchange/src/exchange/message.py b/packages/exchange/src/exchange/message.py index 5edff692cc..e072a86020 100644 --- a/packages/exchange/src/exchange/message.py +++ b/packages/exchange/src/exchange/message.py @@ -119,3 +119,7 @@ class Message: @classmethod def assistant(cls: type["Message"], text: str) -> "Message": return cls(role="assistant", content=[Text(text)]) + + @property + def summary(self) -> str: + return f"message:{self.role}\n" + "\n".join(content.summary for content in self.content) diff --git a/pyproject.toml b/pyproject.toml index 1c5ed50ef1..af13a3b8cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ jira = "goose.toolkit.jira:Jira" screen = "goose.toolkit.screen:Screen" reasoner = "goose.toolkit.reasoner:Reasoner" repo_context = "goose.toolkit.repo_context.repo_context:RepoContext" +synopsis = "goose.synopsis.toolkit:SynopsisDeveloper" [project.entry-points."goose.profile"] default = "goose.profile:default_profile" @@ -42,6 +43,9 @@ goose = "goose.cli.main:goose_cli" [project.entry-points."goose.cli.group_option"] +[project.entry-points."exchange.moderator"] +synopsis = "goose.synopsis.moderator:Synopsis" + [project.scripts] goose = "goose.cli.main:cli" diff --git a/src/goose/cli/session.py b/src/goose/cli/session.py index 8cdb2f646e..3678484c24 100644 --- a/src/goose/cli/session.py +++ b/src/goose/cli/session.py @@ -21,7 +21,7 @@ from goose.profile import Profile from goose.utils import droid, load_plugins from goose.utils._cost_calculator import get_total_cost_message from goose.utils._create_exchange import create_exchange -from goose.utils.session_file import is_empty_session, is_existing_session, read_or_create_file, save_latest_session +from goose.utils.session_file import is_empty_session, is_existing_session, read_or_create_file, log_messages RESUME_MESSAGE = "I see we were interrupted. How can I help you?" @@ -122,13 +122,13 @@ class Session: def setup_plan(self, plan: dict) -> None: if len(self.exchange.messages): raise ValueError("The plan can only be set on an empty session.") - self.exchange.messages.append(Message.user(plan["kickoff_message"])) - tasks = [] - if "tasks" in plan: - tasks = [dict(description=task, status="planned") for task in plan["tasks"]] - plan_tool_use = ToolUse(id="initialplan", name="update_plan", parameters=dict(tasks=tasks)) - self.exchange.add_tool_use(plan_tool_use) + # we append the plan to the kickoff message for now. We should + # revisit this if we intend plans to be handled in a consistent way across toolkits + plan_steps = "\n" + "\n".join(f"{i}. {t}" for i, t in enumerate(plan["tasks"])) + + message = Message.user(plan["kickoff_message"] + plan_steps) + self.exchange.add(message) def process_first_message(self) -> Optional[Message]: # Get a first input unless it has been specified, such as by a plan @@ -158,9 +158,6 @@ class Session: self.exchange.add(message) self.reply() # Process the user message - save_latest_session(self.session_file_path, self.exchange.messages) - print() # Print a newline for separation. - print(f"[dim]ended run | name:[cyan]{self.name}[/] profile:[cyan]{profile}[/]") print(f"[dim]to resume: [magenta]goose session resume {self.name} --profile {profile}[/][/]") @@ -184,8 +181,6 @@ class Session: try: self.exchange.add(message) self.reply() # Process the user message. - except KeyboardInterrupt: - self.interrupt_reply() except Exception: # rewind to right before the last user message self.exchange.rewind() @@ -197,7 +192,6 @@ class Session: + " - [yellow]depending on the error you may be able to continue[/]" ) self.notifier.stop() - save_latest_session(self.session_file_path, self.exchange.messages) print() # Print a newline for separation. user_input = self.prompt_session.get_user_input() message = Message.user(text=user_input.text) if user_input.to_continue() else None @@ -208,31 +202,48 @@ class Session: @observe_wrapper() def reply(self) -> None: """Reply to the last user message, calling tools as needed""" - self.status_indicator.update("responding") - response = self.exchange.generate() + # These are the *raw* messages, before the moderator rewrites things + committed = [self.exchange.messages[-1]] - if response.text: - print(Markdown(response.text)) - - while response.tool_use: - content = [] - for tool_use in response.tool_use: - tool_result = self.exchange.call_function(tool_use) - content.append(tool_result) - self.exchange.add(Message(role="user", content=content)) + try: self.status_indicator.update("responding") response = self.exchange.generate() + committed.append(response) if response.text: print(Markdown(response.text)) - def interrupt_reply(self) -> None: + while response.tool_use: + content = [] + for tool_use in response.tool_use: + tool_result = self.exchange.call_function(tool_use) + content.append(tool_result) + message = Message(role="user", content=content) + committed.append(message) + self.exchange.add(message) + self.status_indicator.update("responding") + response = self.exchange.generate() + committed.append(response) + + if response.text: + print(Markdown(response.text)) + except KeyboardInterrupt: + # The interrupt reply modifies the message history, + # and we sync those changes to committed + self.interrupt_reply(committed) + + # we log the committed messages only once the reply completes + # this prevents messages related to uncaught errors from being recorded + log_messages(self.session_file_path, committed) + + def interrupt_reply(self, committed: list[Message]) -> None: """Recover from an interruption at an arbitrary state""" # Default recovery message if no user message is pending. recovery = "We interrupted before the next processing started." if self.exchange.messages and self.exchange.messages[-1].role == "user": # If the last message is from the user, remove it. self.exchange.messages.pop() + committed.pop() recovery = "We interrupted before the model replied and removed the last message." if ( @@ -250,9 +261,13 @@ class Session: is_error=True, ) ) - self.exchange.add(Message(role="user", content=content)) + message = Message(role="user", content=content) + self.exchange.add(message) + committed.append(message) recovery = f"We interrupted the existing call to {tool_use.name}. How would you like to proceed?" - self.exchange.add(Message.assistant(recovery)) + message = Message.assistant(recovery) + self.exchange.add(message) + committed.append(message) # Print the recovery message with markup for visibility. print(f"[yellow]{recovery}[/]") diff --git a/src/goose/synopsis/developer.md b/src/goose/synopsis/developer.md new file mode 100644 index 0000000000..2526b287da --- /dev/null +++ b/src/goose/synopsis/developer.md @@ -0,0 +1,32 @@ +Your role is a developer agent. You build software and solve problems by editing files and +running commands on the shell. + +You can use the shell tool to run any command that would work on the relevant operating system. + +You are an expert with ripgrep - `rg`. When you need to locate content in the code base, use +`rg` exclusively. It will respect ignored files for efficiency. + +To locate files by name, use + +```bash +rg --files | rg example.py +``` + +To locate content inside files, use +```bash +rg 'class Example' +``` + + +If you need to edit files, use either the write_file tool or the patch tool. +Make sure to read existing content before attempting to edit. + +The write file tool will do a full overwrite of the existing file, while the patch tool +will edit it using a find and replace. Choose the tool which will make the edit as simple +as possible to execute. + + +# Instructions + +You'll receive a summary and a plan, and can immediately start using your tools and can directly +reply to the user as needed. diff --git a/src/goose/synopsis/moderator.py b/src/goose/synopsis/moderator.py new file mode 100644 index 0000000000..d853881950 --- /dev/null +++ b/src/goose/synopsis/moderator.py @@ -0,0 +1,103 @@ +import os +from goose.toolkit.utils import render_template +from pathlib import Path +from exchange.content import Text +from exchange.exchange import Exchange +from exchange.message import Message +from exchange.moderators import Moderator +from exchange.moderators.passive import PassiveModerator +from exchange.moderators.truncate import ContextTruncate +from goose.synopsis.system import system + + +class Synopsis(Moderator): + """Synopsis rewrites the chat into a single input message after every reply + + The goal of synopsis is to remove history artifacts that eat up the context budget + and provide opportunities for the model to attend to the wrong content. + + When? + - After every message (mostly tool use), we apply automatic context management + - After every user input, we use an LLM to curate the context + This is a compromise for speed, where we do fast updates as frequently as possible + + What? + - [Automatic] Current sytem state + - [Automatic] Current file context + - (TODO) [Automatic] Loaded memories + - [Curated] Summary of the discussion so far + - [Curated] Summary of the plan, next step to solve + + At the moment, this is tightly coupled to the SynopsisDeveloper toolkit as base. We + could revisit how that works, because this application shows some limitations in how the + goose state is managed. + """ + + def __init__(self) -> None: + super().__init__() + self.current_summary = "" + self.current_plan = "" + self.originals = [] + + hints = [] + hints_path = Path(".goosehints") + home_hints_path = Path.home() / ".config/goose/.goosehints" + if hints_path.is_file(): + hints.append(render_template(hints_path)) + if home_hints_path.is_file(): + hints.append(render_template(home_hints_path)) + self.hints = "\n".join(hints) + + def rewrite(self, exchange: Exchange) -> None: + # Get the last message, which would be either a user text or a user tool use + last_message: Message = exchange.messages[-1] + + # The first message is the synopsis, there could be multiple tool usages until we encounter a new user message + # [synopsis, tool_use, tool_result, tool_use, tool_result, ..., assistant_message, user_message] + + if isinstance(last_message.content[0], Text): + # We're in the state: + # [synopsis, tool_use, tool_result, ..., user_message] + # and we'll reset it to: + # [synopsis] + + # keep track of the original messages before we reset + if not self.originals: + self.originals.extend(exchange.messages) + if len(exchange.messages) > 1: + # we are resuming an existing session, and need to restore state + system.restore(exchange.messages) + else: + self.originals.extend(exchange.messages[1:]) + + exchange.messages.clear() + exchange.add(self.get_synopsis(exchange, summarize=True, plan=True)) + else: + # We're in the state + # [synopsis, ..., tool_use, tool_result] + # and we'll keep going but updated the synopsis + # [new_synopsis, ..., tool_use, tool_result] + exchange.messages[0] = self.get_synopsis(exchange) + + def get_synopsis(self, exchange: Exchange, summarize: bool = False, plan: bool = False) -> Message: + if summarize: + self.current_summary = self.summarize(exchange) + + if plan: + self.current_plan = self.plan(exchange) + + return Message.load("synopsis.md", synopsis=self, system=system) + + def summarize(self, exchange: Exchange) -> str: + message = Message.load("summarize.md", synopsis=self, messages=self.originals, exchange=exchange, system=system) + model = os.environ.get("SUMMARIZER", exchange.model) + new_exchange = exchange.replace(moderator=ContextTruncate(), tools=(), system="", messages=[], model=model) + new_exchange.add(message) + return new_exchange.generate().content[0].text + + def plan(self, exchange: Exchange) -> str: + message = Message.load("plan.md", synopsis=self, messages=self.originals, exchange=exchange, system=system) + model = os.environ.get("PLANNER", exchange.model) + new_exchange = exchange.replace(moderator=PassiveModerator(), tools=(), system="", messages=[], model=model) + new_exchange.add(message) + return new_exchange.generate().content[0].text diff --git a/src/goose/synopsis/plan.md b/src/goose/synopsis/plan.md new file mode 100644 index 0000000000..8c7fe098e0 --- /dev/null +++ b/src/goose/synopsis/plan.md @@ -0,0 +1,59 @@ +{{synopsis.current_summary}} + +# Instructions + +Prepare a plan that an agent will use to followup to the request described above. Your plan +will be carried out by an agent with access to the following tools: + +{% for tool in exchange.tools %} +{{tool.name}}: {{tool.description}}{% endfor %} + +If the request is simple, such as a greeting or a request for information or advice, the plan can simply be: +"reply to the user". + +However for anything more complex, reflect on the available tools and describe a step by step +solution that the agent can follow using their tools. + +Your plan needs to use the following format, but can have any number of tasks. + +```json +[ + {"description": "the first task here"}, + {"description": "the second task here"}, +] +``` + +# Hints + +{{synopsis.hints}} + +# Context + +The current state of the agent is: + +{{system.info()}} + +The agent already has access to content of the following files, your plan does not need to include finding or reading these. +However if a file or code object is not here but needs to be viewed to complete the goal, include plan steps to +use `rg` (ripgrep) to find the relevant references. + +{% for file in system.active_files %} +{{file.path}}{% endfor %} + +# Examples + +These examples show the format you should follow + +```json +[ + {"description": "reply to the user"}, +] +``` + +```json +[ + {"description": "create a directory 'demo'"}, + {"description": "write a file at 'demo/fibonacci.py' with a function fibonacci implementation"}, + {"description": "run python demo/fibonacci.py"}, +] +``` diff --git a/src/goose/synopsis/summarize.md b/src/goose/synopsis/summarize.md new file mode 100644 index 0000000000..bea9f3e9a6 --- /dev/null +++ b/src/goose/synopsis/summarize.md @@ -0,0 +1,26 @@ +The following is a conversation between a human and an AI agent, +who is taking actions on behalf of the human. + +Messages sent by the human are marked with the "user" role. +Messages sent by the AI are marked with the "assistant" role. +When the AI requests a tool be called, there will be a tool_use and a corresponding tool_result. + +These are the tools available to the agent: +{% for tool in exchange.tools %} +{{tool.name}}: {{tool.description}}{% endfor %} + +# Messages +{% for message in messages %} +{{message.summary}} +{% endfor %} + +# Instructions + +Summarize the conversation so far. Highlight the relevant details, making sure to include any +relevant tool use and result content that is important for the conversation. + +To preserve space, you can omit some details: +- File content details from tools like read and write file: they will be included separately +- Errors that have already been resolved: provide only the current working details + +Make sure to emphasize the most recent request from the human at the end of your summary. diff --git a/src/goose/synopsis/synopsis.md b/src/goose/synopsis/synopsis.md new file mode 100644 index 0000000000..f7ac27d0c4 --- /dev/null +++ b/src/goose/synopsis/synopsis.md @@ -0,0 +1,25 @@ +# Current System Info + +{{system.info()}} + +# Hints + +{{synopsis.hints}} + +# Relevant Files + +{% for file in system.active_files %} +{{file.path}} +```{{file.language}} +{{file.content}} +``` + +{% endfor %} + +# Summary + +{{synopsis.current_summary}} + +# Suggested Plan + +{{synopsis.current_plan}} diff --git a/src/goose/synopsis/system.py b/src/goose/synopsis/system.py new file mode 100644 index 0000000000..4a105bec59 --- /dev/null +++ b/src/goose/synopsis/system.py @@ -0,0 +1,148 @@ +import json +from tiktoken import get_encoding +from exchange import Message +import subprocess +import os +import atexit +import platform +from pathlib import Path +from typing import Dict, Iterable, List, Set + +from attrs import define, field +from exchange.content import ToolUse +from goose.toolkit.utils import get_language + + +@define +class File: + path: str + content: str + language: str + + @property + def context(self) -> str: + return f"""\ +{self.path} +```{self.language} +{self.content} +```""" + + +@define +class OperatingSystem: + """ + OperatingSystem class that can manage background processes created using subprocess.Popen. + """ + + cwd: str = os.getcwd() + platform: str = platform.system() + env: Dict[str, str] = os.environ.copy() + _active_files: Set[str] = field(init=False, factory=set) + _processes: Dict[int, subprocess.Popen] = field(init=False, factory=dict) + + def __attrs_post_init__(self) -> None: + atexit.register(self._cleanup_processes) + + def _cleanup_processes(self) -> None: + """Terminate all processes in the process queue.""" + for process_id in list(self._processes.keys()): + self.cancel_process(process_id) + + def add_process(self, process: subprocess.Popen) -> int: + """Add a new background process and return its assigned ID.""" + process_id = process.pid + os.set_blocking(process.stdout.fileno(), False) + self._processes[process_id] = process + return process_id + + def get_processes(self) -> Dict[int, str]: + """List all background processes with their IDs and commands.""" + return {pid: str(proc.args) for pid, proc in self._processes.items()} + + def view_process_output(self, process_id: int) -> str: + """View the output of a running background process.""" + if not (process := self._processes.get(process_id)): + raise ValueError(f"No process found with ID: {process_id}") + + output = [] + while line := process.stdout.readline(): + output.append(line) + + return "".join(output) + + def cancel_process(self, process_id: int) -> bool: + """Cancel the background process with the specified ID.""" + process = self._processes.pop(process_id, None) + if process: + process.terminate() + return True + return False + + def to_relative(self, path: str) -> str: + """convert string path to be relative to cwd, assuming it starts absolute""" + return os.path.relpath(path, start=self.cwd) + + def to_patho(self, path: str) -> Path: + """convert string to pathlib.Path object + + This attempts to resolve the path against the cwd + """ + patho = Path(path) + if patho.is_absolute(): + return patho + + return (self.cwd / patho).resolve() + + def remember_file(self, path: str) -> None: + """Place a file into the active files""" + path = str(self.to_patho(path)) + + # Do a size check on the file to ensure we don't overload the LLM context + with open(path, "r") as f: + content = f.read() + + max_output_chars = 2**20 + max_output_tokens = 16000 + encoder = get_encoding("cl100k_base") + + if len(content) > max_output_chars or len(encoder.encode(content)) > max_output_tokens: + raise ValueError(f"The file at {path} is too large to read directly!") + + self._active_files.add(path) + + def forget_file(self, path: str) -> None: + """Forget an existing active file""" + self._active_files.discard(str(self.to_patho(path))) + + def info(self) -> str: + """Summarize the current operating system""" + return json.dumps( + dict( + os=self.platform, + cwd=str(self.to_patho(self.cwd)), + shell=os.environ.get("SHELL", "unknown"), + ), + indent=4, + ) + + def is_active(self, path: str) -> bool: + resolved = self.to_patho(path) + return str(resolved) in self._active_files + + @property + def active_files(self) -> Iterable["File"]: + """Yield a File instance for each path in active files, with paths relative to cwd.""" + self._active_files = set(f for f in self._active_files if Path(f).exists()) + + for path in self._active_files: + yield File(path=self.to_relative(path), content=Path(path).read_text(), language=get_language(path)) + + def restore(self, messages: List[Message]) -> None: + """Restore the file content space from a previous sessions""" + for message in messages: + for content in message.content: + if isinstance(content, ToolUse) and content.name == "read_file": + self.remember_file(content.parameters["path"]) + + +system = OperatingSystem() diff --git a/src/goose/synopsis/toolkit.py b/src/goose/synopsis/toolkit.py new file mode 100644 index 0000000000..2b694d8ede --- /dev/null +++ b/src/goose/synopsis/toolkit.py @@ -0,0 +1,244 @@ +# janky global state for now, think about it +import subprocess +import os +from pathlib import Path +from typing import Dict + +from exchange import Message +from goose.synopsis.system import system +from goose.toolkit.base import Toolkit, tool +from goose.toolkit.utils import RULEPREFIX, RULESTYLE, get_language +from goose.utils.shell import is_dangerous_command, shell, keep_unsafe_command_prompt +from rich.markdown import Markdown +from rich.rule import Rule + + +class SynopsisDeveloper(Toolkit): + """Provides shell and file operation tools using OperatingSystem.""" + + def __init__(self, *args: object, **kwargs: Dict[str, object]) -> None: + super().__init__(*args, **kwargs) + + def system(self) -> str: + """Retrieve system configuration details for developer""" + system_prompt = Message.load("developer.md").text + return system_prompt + + def logshell(self, command: str, title: str = "shell") -> None: + self.notifier.log("") + self.notifier.log( + Rule(RULEPREFIX + f"{title} | [dim magenta]{os.path.abspath(system.cwd)}[/]", style=RULESTYLE, align="left") + ) + self.notifier.log(Markdown(f"```bash\n{command}\n```")) + self.notifier.log("") + + @tool + def source(self, path: str) -> str: + """Source the file at path, keeping the updates reflected in future shell commands + + Args: + path (str): The path to the file to source. + """ + source_command = f"source {path} && env" + self.logshell(f"source {path}") + result = shell(source_command, self.notifier, self.exchange_view, cwd=system.cwd, env=system.env) + env_vars = dict(line.split("=", 1) for line in result.splitlines() if "=" in line) + system.env.update(env_vars) + return f"Sourced {path}" + + @tool + def shell(self, command: str) -> str: + """Execute any command on the shell + + Args: + command (str): The shell command to run. It can support multiline statements + if you need to run more than one at a time + """ + if command.startswith("cat"): + raise ValueError("You must read files through the read_file tool.") + if command.startswith("cd"): + raise ValueError("You must change dirs through the change_dir tool.") + if command.startswith("source"): + raise ValueError("You must source files through the source tool.") + + self.logshell(command) + return shell(command, self.notifier, self.exchange_view, cwd=system.cwd, env=system.env) + + @tool + def read_file(self, path: str) -> str: + """Read the content of the file at path + + Args: + path (str): The destination file path, in the format "path/to/file.txt" + """ + system.remember_file(path) + self.logshell(f"cat {path}") + return f"The file content at {path} has been updated above." + + @tool + def write_file(self, path: str, content: str) -> str: + """ + Write a file at the specified path with the provided content. This will create any directories if they do not exist. + The content will fully overwrite the existing file. + + Args: + path (str): The destination file path, in the format "path/to/file.txt" + content (str): The raw file content. + """ # noqa: E501 + patho = system.to_patho(path) + + if patho.exists() and not system.is_active(path): + print(f"We are warning the LLM to view before write in write_file, with path={path} and patho={str(patho)}") + raise ValueError(f"You must view {path} using read_file before you overwrite it") + + patho.parent.mkdir(parents=True, exist_ok=True) + patho.write_text(content) + system.remember_file(path) + + language = get_language(path) + md = f"```{language}\n{content}\n```" + + self.notifier.log("") + self.notifier.log(Rule(RULEPREFIX + path, style=RULESTYLE, align="left")) + self.notifier.log(Markdown(md)) + self.notifier.log("") + + return f"Successfully wrote to {path}" + + @tool + def patch_file(self, path: str, before: str, after: str) -> str: + """Patch the file at the specified by replacing before with after + + Before **must** be present exactly once in the file, so that it can safely + be replaced with after. + + Args: + path (str): The path to the file, in the format "path/to/file.txt" + before (str): The content that will be replaced + after (str): The content it will be replaced with + """ + self.notifier.status(f"editing {path}") + patho = system.to_patho(path) + + if not patho.exists(): + raise ValueError(f"You can't patch {path} - it does not exist yet") + + if not system.is_active(path): + raise ValueError(f"You must view {path} using read_file before you patch it") + + language = get_language(path) + + content = patho.read_text() + + if content.count(before) > 1: + raise ValueError("The before content is present multiple times in the file, be more specific.") + if content.count(before) < 1: + raise ValueError("The before content was not found in file, be careful that you recreate it exactly.") + + content = content.replace(before, after) + system.remember_file(path) + patho.write_text(content) + + output = f""" +```{language} +{before} +``` +-> +```{language} +{after} +``` +""" + self.notifier.log("") + self.notifier.log(Rule(RULEPREFIX + path, style=RULESTYLE, align="left")) + self.notifier.log(Markdown(output)) + self.notifier.log("") + return "Succesfully replaced before with after." + + @tool + def start_process(self, command: str) -> int: + """Start a background process running the specified command + + Use this exclusively for processes that you need to run in the background + because they do not terminate, such as running a webserver. + + Args: + command (str): The shell command to run + """ + self.logshell(command, title="background") + + if is_dangerous_command(command): + self.notifier.stop() + if not keep_unsafe_command_prompt(command): + raise RuntimeError( + f"The command {command} was rejected as dangerous by the user." + " Do not proceed further, instead ask for instructions." + ) + self.notifier.start() + + process = subprocess.Popen( + command, + shell=True, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + cwd=system.cwd, + env=system.env, + ) + process_id = system.add_process(process) + return process_id + + @tool + def list_processes(self) -> Dict[int, str]: + """List all running background processes with their IDs and commands.""" + processes = system.get_processes() + process_list = "```\n" + "\n".join(f"id: {pid}, command: {cmd}" for pid, cmd in processes.items()) + "\n```" + self.notifier.log("") + self.notifier.log(Rule(RULEPREFIX + "processes", style=RULESTYLE, align="left")) + self.notifier.log(Markdown(process_list)) + self.notifier.log("") + return processes + + @tool + def view_process_output(self, process_id: int) -> str: + """View the output of a running background process + + Args: + process_id (int): The ID of the process to view output. + """ + self.notifier.log("") + self.notifier.log(Rule(RULEPREFIX + "processes", style=RULESTYLE, align="left")) + self.notifier.log(Markdown(f"```\nreading {process_id}\n```")) + self.notifier.log("") + output = system.view_process_output(process_id) + return output + + @tool + def cancel_process(self, process_id: int) -> str: + """Cancel the background process with the specified ID. + + Args: + process_id (int): The ID of the process to be cancelled. + """ + result = system.cancel_process(process_id) + self.logshell(f"kill {process_id}") + if result: + return f"process {process_id} cancelled" + else: + return f"no known process {process_id}" + + @tool + def change_dir(self, path: str) -> str: + """Change the directory to the specified path + + Args: + path (str): The new dir path, in the format "path/to/dir" + """ + patho = system.to_patho(path) + if not patho.is_dir(): + raise ValueError(f"The directory {path} does not exist") + if patho.resolve() < Path(os.getcwd()).resolve(): + raise ValueError("You can cd into subdirs but not above the directory where we started.") + self.logshell(f"cd {path}") + system.cwd = str(patho) + return path diff --git a/src/goose/toolkit/developer.py b/src/goose/toolkit/developer.py index 7b5426a5cc..4bb40a153a 100644 --- a/src/goose/toolkit/developer.py +++ b/src/goose/toolkit/developer.py @@ -1,7 +1,5 @@ import os import re -import subprocess -import time import tempfile import httpx @@ -9,26 +7,12 @@ from pathlib import Path from exchange import Message from goose.toolkit.base import Toolkit, tool -from goose.toolkit.utils import get_language, render_template -from goose.utils.ask import ask_an_ai -from goose.utils.check_shell_command import is_dangerous_command +from goose.toolkit.utils import get_language, render_template, RULEPREFIX, RULESTYLE +from goose.utils.shell import shell from rich.markdown import Markdown -from rich.prompt import Confirm from rich.table import Table -from rich.text import Text from rich.rule import Rule -RULESTYLE = "bold" -RULEPREFIX = f"[{RULESTYLE}]───[/] " - - -def keep_unsafe_command_prompt(command: str) -> bool: - command_text = Text(command, style="bold red") - message = ( - Text("\nWe flagged the command: ") + command_text + Text(" as potentially unsafe, do you want to proceed?") - ) - return Confirm.ask(message, default=True) - class Developer(Toolkit): """Provides a set of general purpose development capabilities @@ -40,6 +24,7 @@ class Developer(Toolkit): def __init__(self, *args: object, **kwargs: dict[str, object]) -> None: super().__init__(*args, **kwargs) self.timestamps: dict[str, float] = {} + self.cwd = os.getcwd() def system(self) -> str: """Retrieve system configuration details for developer""" @@ -195,93 +180,7 @@ class Developer(Toolkit): # Log the command being executed in a visually structured format (Markdown). self.notifier.log(Rule(RULEPREFIX + "shell", style=RULESTYLE, align="left")) self.notifier.log(Markdown(f"```bash\n{command}\n```")) - - if is_dangerous_command(command): - # Stop the notifications so we can prompt - self.notifier.stop() - if not keep_unsafe_command_prompt(command): - raise RuntimeError( - f"The command {command} was rejected as dangerous by the user." - " Do not proceed further, instead ask for instructions." - ) - self.notifier.start() - self.notifier.status("running shell command") - - # Define patterns that might indicate the process is waiting for input - interaction_patterns = [ - r"Do you want to", # Common prompt phrase - r"Enter password", # Password prompt - r"Are you sure", # Confirmation prompt - r"\(y/N\)", # Yes/No prompt - r"Press any key to continue", # Awaiting keypress - r"Waiting for input", # General waiting message - ] - compiled_patterns = [re.compile(pattern, re.IGNORECASE) for pattern in interaction_patterns] - - proc = subprocess.Popen( - command, - shell=True, - stdin=subprocess.DEVNULL, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - # this enables us to read lines without blocking - os.set_blocking(proc.stdout.fileno(), False) - - # Accumulate the output logs while checking if it might be blocked - output_lines = [] - last_output_time = time.time() - cutoff = 10 - while proc.poll() is None: - self.notifier.status("running shell command") - line = proc.stdout.readline() - if line: - output_lines.append(line) - last_output_time = time.time() - - # If we see a clear pattern match, we plan to abort - exit_criteria = any(pattern.search(line) for pattern in compiled_patterns) - - # and if we haven't seen a new line in 10+s, check with AI to see if it may be stuck - if not exit_criteria and time.time() - last_output_time > cutoff: - self.notifier.status("checking on shell status") - response = ask_an_ai( - input="\n".join([command] + output_lines), - prompt=( - "You will evaluate the output of shell commands to see if they may be stuck." - " Look for commands that appear to be awaiting user input, or otherwise running indefinitely (such as a web service)." # noqa - " A command that will take a while, such as downloading resources is okay." # noqa - " return [Yes] if stuck, [No] otherwise." - ), - exchange=self.exchange_view.processor, - with_tools=False, - ) - exit_criteria = "[yes]" in response.content[0].text.lower() - # We add exponential backoff for how often we check for the command being stuck - cutoff *= 10 - - if exit_criteria: - proc.terminate() - raise ValueError( - f"The command `{command}` looks like it will run indefinitely or is otherwise stuck." - f"You may be able to specify inputs if it applies to this command." - f"Otherwise to enable continued iteration, you'll need to ask the user to run this command in another terminal." # noqa - ) - - # read any remaining lines - while line := proc.stdout.readline(): - output_lines.append(line) - output = "".join(output_lines) - - # Determine the result based on the return code - if proc.returncode == 0: - result = "Command succeeded" - else: - result = f"Command failed with returncode {proc.returncode}" - - # Return the combined result and outputs if we made it this far - return "\n".join([result, output]) + return shell(command, self.notifier, self.exchange_view) @tool def write_file(self, path: str, content: str) -> str: diff --git a/src/goose/toolkit/prompts/developer.jinja b/src/goose/toolkit/prompts/developer.jinja index 6da404e9ea..16e7605f3d 100644 --- a/src/goose/toolkit/prompts/developer.jinja +++ b/src/goose/toolkit/prompts/developer.jinja @@ -1,51 +1,7 @@ -To start solving problems, you will always create a plan, and then immediately -execute the next step of the plan. Do not wait for confirmation on your plan, -always proceed to the next step. After each step, update the plan based -on the output you recieve from any actions you take, including marking -finished tasks as complete! - -To accomplish this, you will call your tools such as update_plan, shell, or write. - -Always use the update_plan tool before taking any new actions, to show the client -an up to date plan. The plan will be displayed automatically any time you update it. - -The plan should consist of as few entries as possible, and translate from the user -request into concrete tasks that you will use to get it done. These should reflect -the actions you will need to take, such as writing files or executing shell commands. - -For example, here's a plan to unstage all edited files in a git repo - -{"description": "Use the git status command to find edited files", "status": "planned"} -{"description": "For each file with changes, call git restore on the file", "status": "planned"} - -After running git status, you would update to - -{"description": "Use the git status command to find edited files", "status": "complete"} -{"description": "For each file with changes, call git restore on the file", "status": "planned"} - - -Here's another plan example to get the sum of the "payment" column in data.csv - -{"description": "Install pandas", "status": "planned"} -{"description": "Write a python script in the file sum.py that loads the csv in pandas and sums the column", "status": "planned"} -{"description": "Run the python script with bash", "status": "planned"} - -If you were to encounter an error along the way, you can update the plan to specify a new approach. -Always call update_plan before any other tool calls! You should specify the whole plan upfront as planned, -and then update status at each step. **Do not describe the plan in your text response, only communicate -the plan through the tool** - -If you need to manipulate files, always prefer the write file tool. To edit a file -that already exists, first check the content with the shell and then overwrite it - - -Some of the files that you work with will be long. When you want to edit a long file, -prefer to use the patch tool. Make sure that you always read the file using the read tool -before you call patch. - -The patch and write tools can accomplish the same operations, but patch is a more complex tool. -The patch tool is worth it when the file is large enough that it would be tedious to fully rewrite. +Your role is a developer agent. You build software and solve problems by editing files and +running commands on the shell. +You can use the shell tool to run any command that would work on the relevant operating system. You are an expert with ripgrep - `rg`. When you need to locate content in the code base, use `rg` exclusively. It will respect ignored files for efficiency. @@ -60,3 +16,17 @@ To locate content inside files, use ```bash rg 'class Example' ``` + + +If you need to manipulate files, use either the write_file tool or the patch tool. +Make sure to read existing content before attempting to edit. + +The write file tool will do a full overwrite of the existing file, while the patch tool +will edit it using a find and replace. Choose the tool which will make the edit as simple +as possible to execute. + + +# Instructions + +You'll receive a summary and a plan, and can immediately start using your tools and can directly +reply to the user as needed. diff --git a/src/goose/toolkit/utils.py b/src/goose/toolkit/utils.py index d6f0335b1d..470d1c8405 100644 --- a/src/goose/toolkit/utils.py +++ b/src/goose/toolkit/utils.py @@ -7,7 +7,11 @@ from pygments.util import ClassNotFound from jinja2 import Environment, FileSystemLoader -def get_language(filename: Path) -> str: +RULESTYLE = "bold" +RULEPREFIX = f"[{RULESTYLE}]───[/] " + + +def get_language(filename: str) -> str: """ Determine the programming language of a file based on its filename extension. @@ -19,7 +23,7 @@ def get_language(filename: Path) -> str: """ try: lexer = get_lexer_for_filename(filename) - return lexer.name + return lexer.name.lower() except ClassNotFound: return "" diff --git a/src/goose/utils/check_shell_command.py b/src/goose/utils/check_shell_command.py deleted file mode 100644 index e6c15f288b..0000000000 --- a/src/goose/utils/check_shell_command.py +++ /dev/null @@ -1,34 +0,0 @@ -import re - - -def is_dangerous_command(command: str) -> bool: - """ - Check if the command matches any dangerous patterns. - - Dangerous patterns in this function are defined as commands that may present risk to system stability. - - Args: - command (str): The shell command to check. - - Returns: - bool: True if the command is dangerous, False otherwise. - """ - dangerous_patterns = [ - # Commands that are generally unsafe - r"\brm\b", # rm command - r"\bgit\s+push\b", # git push command - r"\bsudo\b", # sudo command - r"\bmv\b", # mv command - r"\bchmod\b", # chmod command - r"\bchown\b", # chown command - r"\bmkfs\b", # mkfs command - r"\bsystemctl\b", # systemctl command - r"\breboot\b", # reboot command - r"\bshutdown\b", # shutdown command - # Target files that are unsafe - r"\b~\/\.|\/\.\w+", # commands that point to files or dirs in home that start with a dot (dotfiles) - ] - for pattern in dangerous_patterns: - if re.search(pattern, command): - return True - return False diff --git a/src/goose/utils/session_file.py b/src/goose/utils/session_file.py index 510b4aa193..f8b5e95185 100644 --- a/src/goose/utils/session_file.py +++ b/src/goose/utils/session_file.py @@ -1,6 +1,4 @@ import json -import os -import tempfile from pathlib import Path from typing import Iterator @@ -17,11 +15,6 @@ def is_empty_session(path: Path) -> bool: return path.is_file() and path.stat().st_size == 0 -def write_to_file(file_path: Path, messages: list[Message]) -> None: - with open(file_path, "w") as f: - _write_messages_to_file(f, messages) - - def read_or_create_file(file_path: Path) -> list[Message]: if file_path.exists(): return read_from_file(file_path) @@ -55,15 +48,8 @@ def session_file_exists(session_files_directory: Path) -> bool: return any(list_session_files(session_files_directory)) -def save_latest_session(file_path: Path, messages: list[Message]) -> None: - with tempfile.NamedTemporaryFile("w", delete=False) as temp_file: - _write_messages_to_file(temp_file, messages) - temp_file_path = temp_file.name - - os.replace(temp_file_path, file_path) - - -def _write_messages_to_file(file: any, messages: list[Message]) -> None: - for m in messages: - json.dump(m.to_dict(), file) - file.write("\n") +def log_messages(file_path: Path, messages: list[Message]) -> None: + with open(file_path, "a") as f: + for message in messages: + json.dump(message.to_dict(), f) + f.write("\n") diff --git a/src/goose/utils/shell.py b/src/goose/utils/shell.py new file mode 100644 index 0000000000..d53bbb0563 --- /dev/null +++ b/src/goose/utils/shell.py @@ -0,0 +1,149 @@ +import os +import re +import subprocess +import time +from typing import Mapping, Optional + +from goose.notifier import Notifier +from goose.utils.ask import ask_an_ai +from goose.view import ExchangeView +from rich.prompt import Confirm + + +def is_dangerous_command(command: str) -> bool: + """ + Check if the command matches any dangerous patterns. + + Dangerous patterns in this function are defined as commands that may present risk to system stability. + + Args: + command (str): The shell command to check. + + Returns: + bool: True if the command is dangerous, False otherwise. + """ + dangerous_patterns = [ + # Commands that are generally unsafe + r"\brm\b", # rm command + r"\bgit\s+push\b", # git push command + r"\bsudo\b", # sudo command + r"\bmv\b", # mv command + r"\bchmod\b", # chmod command + r"\bchown\b", # chown command + r"\bmkfs\b", # mkfs command + r"\bsystemctl\b", # systemctl command + r"\breboot\b", # reboot command + r"\bshutdown\b", # shutdown command + # Target files that are unsafe + r"\b~\/\.|\/\.\w+", # commands that point to files or dirs in home that start with a dot (dotfiles) + ] + for pattern in dangerous_patterns: + if re.search(pattern, command): + return True + return False + + +def keep_unsafe_command_prompt(command: str) -> bool: + message = f"\nWe flagged the command - [bold red]{command}[/] - as potentially unsafe, do you want to proceed?" + return Confirm.ask(message, default=True) + + +def shell( + command: str, + notifier: Notifier, + exchange_view: ExchangeView, + cwd: Optional[str] = None, + env: Optional[Mapping[str, str]] = None, +) -> str: + """Execute a command on the shell + + This handles + """ + if is_dangerous_command(command): + # Stop the notifications so we can prompt + notifier.stop() + if not keep_unsafe_command_prompt(command): + raise RuntimeError( + f"The command {command} was rejected as dangerous by the user." + " Do not proceed further, instead ask for instructions." + ) + notifier.start() + notifier.status("running shell command") + + # Define patterns that might indicate the process is waiting for input + interaction_patterns = [ + r"Do you want to", # Common prompt phrase + r"Enter password", # Password prompt + r"Are you sure", # Confirmation prompt + r"\(y/N\)", # Yes/No prompt + r"Press any key to continue", # Awaiting keypress + r"Waiting for input", # General waiting message + ] + compiled_patterns = [re.compile(pattern, re.IGNORECASE) for pattern in interaction_patterns] + + proc = subprocess.Popen( + command, + shell=True, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + cwd=cwd, + env=env, + ) + # this enables us to read lines without blocking + os.set_blocking(proc.stdout.fileno(), False) + + # Accumulate the output logs while checking if it might be blocked + output_lines = [] + last_output_time = time.time() + cutoff = 10 + while proc.poll() is None: + notifier.status("running shell command") + line = proc.stdout.readline() + if line: + output_lines.append(line) + last_output_time = time.time() + + # If we see a clear pattern match, we plan to abort + exit_criteria = any(pattern.search(line) for pattern in compiled_patterns) + + # and if we haven't seen a new line in 10+s, check with AI to see if it may be stuck + if not exit_criteria and time.time() - last_output_time > cutoff: + notifier.status("checking on shell status") + response = ask_an_ai( + input="\n".join([command] + output_lines), + prompt=( + "You will evaluate the output of shell commands to see if they may be stuck." + " Look for commands that appear to be awaiting user input, or otherwise running indefinitely (such as a web service)." # noqa + " A command that will take a while, such as downloading resources is okay." # noqa + " return [Yes] if stuck, [No] otherwise." + ), + exchange=exchange_view.processor, + with_tools=False, + ) + exit_criteria = "[yes]" in response.content[0].text.lower() + # We add exponential backoff for how often we check for the command being stuck + cutoff *= 10 + + if exit_criteria: + proc.terminate() + raise ValueError( + f"The command `{command}` looks like it will run indefinitely or is otherwise stuck." + f"You may be able to specify inputs if it applies to this command." + f"Otherwise to enable continued iteration, you'll need to ask the user to run this command in another terminal." # noqa + ) + + # read any remaining lines + while line := proc.stdout.readline(): + output_lines.append(line) + output = "".join(output_lines) + + # Determine the result based on the return code + if proc.returncode == 0: + result = "Command succeeded" + else: + result = f"Command failed with returncode {proc.returncode}" + + # Return the combined result and outputs if we made it this far + return "\n".join([result, output]) diff --git a/tests/cli/test_session.py b/tests/cli/test_session.py index b2eafea6c4..c605b524a7 100644 --- a/tests/cli/test_session.py +++ b/tests/cli/test_session.py @@ -1,8 +1,7 @@ -import json from unittest.mock import MagicMock, patch import pytest -from exchange import Exchange, Message, ToolResult, ToolUse +from exchange import Message, ToolResult, ToolUse from goose.cli.prompt.goose_prompt_session import GoosePromptSession from goose.cli.prompt.user_input import PromptAction, UserInput from goose.cli.session import Session @@ -124,63 +123,6 @@ def test_log_log_cost(create_session_with_mock_configs): mock_logger.info.assert_called_once_with(cost_message) -@patch.object(GoosePromptSession, "get_user_input", name="get_user_input") -@patch.object(Exchange, "generate", name="mock_generate") -@patch("goose.cli.session.save_latest_session", name="mock_save_latest_session") -def test_run_should_auto_save_session( - mock_save_latest_session, - mock_generate, - mock_get_user_input, - create_session_with_mock_configs, - mock_sessions_path, -): - def custom_exchange_generate(self, *args, **kwargs): - message = Message.assistant("Response") - self.add(message) - return message - - def mock_generate_side_effect(*args, **kwargs): - return custom_exchange_generate(session.exchange, *args, **kwargs) - - def save_latest_session(file, messages): - file.write_text("\n".join(json.dumps(m.to_dict()) for m in messages)) - - user_inputs = [ - UserInput(action=PromptAction.CONTINUE, text="Question1"), - UserInput(action=PromptAction.CONTINUE, text="Question2"), - UserInput(action=PromptAction.EXIT), - ] - session = create_session_with_mock_configs({"name": SESSION_NAME}) - - mock_get_user_input.side_effect = user_inputs - mock_generate.side_effect = mock_generate_side_effect - mock_save_latest_session.side_effect = save_latest_session - - session.run() - - session_file = mock_sessions_path / f"{SESSION_NAME}.jsonl" - - assert mock_generate.call_count == 2 - assert mock_save_latest_session.call_count == 2 - assert mock_save_latest_session.call_args_list[0][0][0] == session_file - assert session_file.exists() - - with open(session_file, "r") as f: - saved_messages = [json.loads(line) for line in f] - - expected_messages = [ - Message.user("Question1"), - Message.assistant("Response"), - Message.user("Question2"), - Message.assistant("Response"), - ] - - assert len(saved_messages) == len(expected_messages) - for saved, expected in zip(saved_messages, expected_messages): - assert saved["role"] == expected.role - assert saved["content"][0]["text"] == expected.text - - @patch("goose.cli.session.droid", return_value="generated_session_name", name="mock_droid") def test_set_generated_session_name(mock_droid, create_session_with_mock_configs, mock_sessions_path): session = create_session_with_mock_configs({"name": None}) diff --git a/tests/synopsis/test_moderator.py b/tests/synopsis/test_moderator.py new file mode 100644 index 0000000000..4d610cb4c3 --- /dev/null +++ b/tests/synopsis/test_moderator.py @@ -0,0 +1,33 @@ +from unittest.mock import patch + +import pytest +from exchange.content import Text +from exchange.message import Message +from goose.synopsis.moderator import Synopsis + + +@pytest.fixture +def mock_exchange(exchange_factory): + exchange = exchange_factory() + exchange.messages.extend( + [ + Message(role="user", content=[Text("Test content for user message")]), + Message(role="assistant", content=[Text("Test content for assistant message")]), + Message(role="user", content=[Text("Another user message")]), + ] + ) + return exchange + + +def test_rewrite_with_tool_use(mock_exchange): + tool_use_message = Message(role="user", content=[Text("Tool use message")]) + mock_exchange.messages.append(tool_use_message) + + with patch.object(Synopsis, "get_synopsis") as mock_get_synopsis: + message = Message(role="synopsis", content=[Text("Updated synopsis")]) + mock_get_synopsis.return_value = message + synopsis = Synopsis() + synopsis.rewrite(mock_exchange) + + # The first message should be replaced, and the rest are cleared + assert mock_exchange.messages == [message] diff --git a/tests/synopsis/test_process_management.py b/tests/synopsis/test_process_management.py new file mode 100644 index 0000000000..ae9483df92 --- /dev/null +++ b/tests/synopsis/test_process_management.py @@ -0,0 +1,62 @@ +import pytest +import time +import requests +from goose.synopsis.toolkit import SynopsisDeveloper +from goose.synopsis.system import system + + +class MockNotifier: + def log(self, message): + pass + + def status(self, message): + pass + + +@pytest.fixture +def toolkit(tmpdir): + original_cwd = system.cwd + system.cwd = str(tmpdir) + notifier = MockNotifier() + toolkit = SynopsisDeveloper(notifier=notifier) + + yield toolkit + + # Teardown: cancel all processes and restore original working directory + for process_id in list(system._processes.keys()): + system.cancel_process(process_id) + system.cwd = original_cwd + + +def test_start_process(toolkit): + process_id = toolkit.start_process("python -m http.server 8000") + assert process_id > 0 + time.sleep(2) # Give the server time to start + + # Check if the server is running + try: + response = requests.get("http://localhost:8000") + assert response.status_code == 200 + except requests.ConnectionError: + pytest.fail("HTTP server did not start successfully") + output = toolkit.view_process_output(process_id) + assert "200" in output + + +def test_list_processes(toolkit): + process_id = toolkit.start_process("python -m http.server 8001") + processes = toolkit.list_processes() + assert process_id in processes + assert "python -m http.server 8001" in processes[process_id] + + +def test_cancel_process(toolkit): + process_id = toolkit.start_process("python -m http.server 8003") + time.sleep(2) # Give the server time to start + + result = toolkit.cancel_process(process_id) + assert result == f"process {process_id} cancelled" + + # Verify that the process is no longer running + with pytest.raises(ValueError): + toolkit.view_process_output(process_id) diff --git a/tests/synopsis/test_system.py b/tests/synopsis/test_system.py new file mode 100644 index 0000000000..fdc94dbafb --- /dev/null +++ b/tests/synopsis/test_system.py @@ -0,0 +1,93 @@ +import os +from unittest.mock import Mock +import pytest +from goose.synopsis.system import OperatingSystem + + +@pytest.fixture +def os_instance(tmpdir): + original_cwd = os.getcwd() + os.chdir(tmpdir) + yield OperatingSystem(cwd=str(tmpdir)) + os.chdir(original_cwd) + + +def test_to_relative(os_instance, tmpdir): + abs_path = os.path.join(tmpdir, "test_file.txt") + rel_path = os_instance.to_relative(abs_path) + assert rel_path == "test_file.txt" + + +def test_remember_forget_file(os_instance, tmpdir): + test_file = tmpdir.join("test_file.txt") + test_file.write("test content") + + os_instance.remember_file(str(test_file)) + assert os_instance.is_active(str(test_file)) + + os_instance.forget_file(str(test_file)) + assert not os_instance.is_active(str(test_file)) + + +def test_active_files(os_instance, tmpdir): + test_file1 = tmpdir.join("test_file1.txt") + test_file2 = tmpdir.join("test_file2.py") + test_file1.write("test content 1") + test_file2.write("test content 2") + + os_instance.remember_file(str(test_file1)) + os_instance.remember_file(str(test_file2)) + + active_files = list(os_instance.active_files) + assert len(active_files) == 2 + assert any(f.path == "test_file1.txt" for f in active_files) + assert any(f.path == "test_file2.py" for f in active_files) + + +def test_info(os_instance): + info = os_instance.info() + assert "os" in info + assert "cwd" in info + assert "shell" in info + + +def test_add_process(os_instance): + process = Mock() + process.pid = 1234 + process.stdout = Mock() + process.stdout.fileno.return_value = 1 + process_id = os_instance.add_process(process) + assert process_id == 1234 + assert 1234 in os_instance._processes + + +def test_get_processes(os_instance): + process1 = Mock() + process1.pid = 1234 + process1.args = "python -m http.server 8000" + process1.stdout = Mock() + process1.stdout.fileno.return_value = 1 + os_instance.add_process(process1) + + process2 = Mock() + process2.pid = 5678 + process2.args = "python script.py" + process2.stdout = Mock() + process2.stdout.fileno.return_value = 2 + os_instance.add_process(process2) + + processes = os_instance.get_processes() + assert processes == {1234: "python -m http.server 8000", 5678: "python script.py"} + + +def test_cancel_process(os_instance): + process = Mock() + process.pid = 1234 + process.stdout = Mock() + process.stdout.fileno.return_value = 1 + os_instance.add_process(process) + + result = os_instance.cancel_process(1234) + assert result is True + assert 1234 not in os_instance._processes + process.terminate.assert_called_once() diff --git a/tests/synopsis/test_toolkit.py b/tests/synopsis/test_toolkit.py new file mode 100644 index 0000000000..9d213782eb --- /dev/null +++ b/tests/synopsis/test_toolkit.py @@ -0,0 +1,95 @@ +import os +import pytest +from goose.synopsis.toolkit import SynopsisDeveloper +from goose.synopsis.system import system + + +class MockNotifier: + def log(self, message): + pass + + def status(self, message): + pass + + +@pytest.fixture +def toolkit(tmpdir): + original_cwd = os.getcwd() + os.chdir(tmpdir) + system.cwd = str(tmpdir) + notifier = MockNotifier() + toolkit = SynopsisDeveloper(notifier=notifier) + + yield toolkit + + # Teardown: cancel all processes and restore original working directory + for process_id in list(system._processes.keys()): + system.cancel_process(process_id) + os.chdir(original_cwd) + system.cwd = original_cwd + + +def test_shell(toolkit, tmpdir): + result = toolkit.shell("echo 'Hello, World!'") + assert "Hello, World!" in result + + +def test_read_write_file(toolkit, tmpdir): + test_file = tmpdir.join("test_file.txt") + content = "Test content" + + toolkit.write_file(str(test_file), content) + assert test_file.read() == content + + result = toolkit.read_file(str(test_file)) + assert "The file content at" in result + assert system.is_active(str(test_file)) + + +def test_patch_file(toolkit, tmpdir): + test_file = tmpdir.join("test_file.txt") + test_file.write("Hello, World!") + + toolkit.read_file(str(test_file)) # Remember the file + result = toolkit.patch_file(str(test_file), "World", "Universe") + assert "Succesfully replaced before with after" in result + assert test_file.read() == "Hello, Universe!" + + +def test_change_dir(toolkit, tmpdir): + subdir = tmpdir.mkdir("subdir") + result = toolkit.change_dir(str(subdir)) + assert result == str(subdir) + assert system.cwd == str(subdir) + + +def test_start_process(toolkit, tmpdir): + process_id = toolkit.start_process("python -m http.server 8000") + assert process_id > 0 + + # Check if the process is in the list of running processes + processes = toolkit.list_processes() + assert process_id in processes + assert "python -m http.server 8000" in processes[process_id] + + +def test_list_processes(toolkit, tmpdir): + process_id1 = toolkit.start_process("python -m http.server 8001") + process_id2 = toolkit.start_process("python -m http.server 8002") + + processes = toolkit.list_processes() + assert process_id1 in processes + assert process_id2 in processes + assert "python -m http.server 8001" in processes[process_id1] + assert "python -m http.server 8002" in processes[process_id2] + + +def test_cancel_process(toolkit, tmpdir): + process_id = toolkit.start_process("python -m http.server 8003") + + result = toolkit.cancel_process(process_id) + assert result == f"process {process_id} cancelled" + + # Verify that the process is no longer in the list + processes = toolkit.list_processes() + assert process_id not in processes diff --git a/tests/utils/test_check_shell_command.py b/tests/utils/test_check_shell_command.py index f267d8e82b..2e1a8ae262 100644 --- a/tests/utils/test_check_shell_command.py +++ b/tests/utils/test_check_shell_command.py @@ -1,5 +1,5 @@ import pytest -from goose.utils.check_shell_command import is_dangerous_command +from goose.utils.shell import is_dangerous_command @pytest.mark.parametrize( diff --git a/tests/utils/test_session_file.py b/tests/utils/test_session_file.py index 2905645665..60b25505ac 100644 --- a/tests/utils/test_session_file.py +++ b/tests/utils/test_session_file.py @@ -3,15 +3,12 @@ from pathlib import Path from unittest.mock import patch import pytest -from exchange import Message from goose.utils.session_file import ( is_empty_session, list_sorted_session_files, read_from_file, read_or_create_file, - save_latest_session, session_file_exists, - write_to_file, ) @@ -20,17 +17,6 @@ def file_path(tmp_path): return tmp_path / "test_file.jsonl" -def test_read_write_to_file(file_path): - messages = [ - Message.user("prompt1"), - Message.user("prompt2"), - ] - write_to_file(file_path, messages) - assert file_path.exists() - - assert read_from_file(file_path) == messages - - def test_read_from_file_non_existing_file(tmp_path): with pytest.raises(FileNotFoundError): read_from_file(tmp_path / "no_existing.json") @@ -49,16 +35,6 @@ def test_read_or_create_file_when_file_not_exist(tmp_path): assert os.path.exists(file_path) -def test_read_or_create_file_when_file_exists(file_path): - messages = [ - Message.user("prompt1"), - ] - write_to_file(file_path, messages) - - assert file_path.exists() - assert read_from_file(file_path) == messages - - def test_list_sorted_session_files(tmp_path): session_files_directory = tmp_path / "session_files_dir" session_files_directory.mkdir() @@ -98,21 +74,6 @@ def test_session_file_exists_return_true_when_session_file_exists(tmp_path): assert session_file_exists(session_files_directory) -def test_save_latest_session(file_path, tmp_path): - messages = [ - Message.user("prompt1"), - Message.user("prompt2"), - ] - write_to_file(file_path, messages) - - messages.append(Message.user("prompt3")) - save_latest_session(file_path, messages) - - messages_in_file = read_from_file(file_path) - assert messages_in_file == messages - assert len(messages_in_file) == 3 - - def create_session_file(file_path, file_name) -> Path: file = file_path / f"{file_name}.jsonl" file.touch()