tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

chore: initial commit

Browse Source 
Co-authored-by: Lifei Zhou <lifei@squareup.com>
Co-authored-by: Mic Neale <micn@tbd.email>
Co-authored-by: Lily Delalande <ldelalande@squareup.com>
Co-authored-by: Bradley Axen <baxen@squareup.com>
Co-authored-by: Andy Lane <alane@squareup.com>
Co-authored-by: Elena Zherdeva <ezherdeva@squareup.com>
Co-authored-by: Zaki Ali <zaki@squareup.com>
Co-authored-by: Salman Mohammed <smohammed@squareup.com>
This commit is contained in:
Luke Alvoeiro
2024-08-23 16:39:04 -07:00
commit dd126afa6c
68 changed files with 4498 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/cli/prompt/test_goose_prompt_session.py
+47
View File
@@ -0,0 +1,47 @@
from unittest.mock import patch
import pytest
from goose.cli.prompt.goose_prompt_session import GoosePromptSession
from goose.cli.prompt.user_input import PromptAction, UserInput
@pytest.fixture
def mock_prompt_session():
with patch("prompt_toolkit.PromptSession") as mock_prompt_session:
yield mock_prompt_session
def test_get_save_session_name(mock_prompt_session):
mock_prompt_session.prompt.return_value = "my_session"
goose_prompt_session = GoosePromptSession(mock_prompt_session)
assert goose_prompt_session.get_save_session_name() == "my_session"
def test_get_user_input_to_continue(mock_prompt_session):
mock_prompt_session.prompt.return_value = "input_value"
goose_prompt_session = GoosePromptSession(mock_prompt_session)
user_input = goose_prompt_session.get_user_input()
assert user_input == UserInput(PromptAction.CONTINUE, "input_value")
@pytest.mark.parametrize("exit_input", ["exit", ":q"])
def test_get_user_input_to_exit(exit_input, mock_prompt_session):
mock_prompt_session.prompt.return_value = exit_input
goose_prompt_session = GoosePromptSession(mock_prompt_session)
user_input = goose_prompt_session.get_user_input()
assert user_input == UserInput(PromptAction.EXIT)
@pytest.mark.parametrize("error", [EOFError, KeyboardInterrupt])
def test_get_user_input_to_exit_when_error_occurs(error, mock_prompt_session):
mock_prompt_session.prompt.side_effect = error
goose_prompt_session = GoosePromptSession(mock_prompt_session)
user_input = goose_prompt_session.get_user_input()
assert user_input == UserInput(PromptAction.EXIT)
tests/cli/prompt/test_lexer.py
+253
View File
@@ -0,0 +1,253 @@
from goose.cli.prompt.lexer import (
PromptLexer,
command_itself,
completion_for_command,
value_for_command,
)
from prompt_toolkit.document import Document
# Helper function to create a Document and lexer instance
def create_lexer_and_document(commands, text):
lexer = PromptLexer(commands)
document = Document(text)
return lexer, document
# Test cases
def test_lex_document_command():
lexer, document = create_lexer_and_document(["file"], "/file:example.txt")
tokens = lexer.lex_document(document)
expected_tokens = [("class:command", "/file:"), ("class:parameter", "example.txt")]
assert tokens(0) == expected_tokens
def test_lex_document_partial_command():
lexer, document = create_lexer_and_document(["file"], "/fi")
tokens = lexer.lex_document(document)
expected_tokens = [("class:command", "/fi")]
assert tokens(0) == expected_tokens
def test_lex_document_with_text():
lexer, document = create_lexer_and_document(["file"], "Some text /file:example.txt")
tokens = lexer.lex_document(document)
expected_tokens = [
("class:text", "S"),
("class:text", "o"),
("class:text", "m"),
("class:text", "e"),
("class:text", " "),
("class:text", "t"),
("class:text", "e"),
("class:text", "x"),
("class:text", "t"),
("class:text", " "),
("class:command", "/file:"),
("class:parameter", "example.txt"),
]
assert tokens(0) == expected_tokens
def test_lex_document_with_command_in_middle():
lexer, document = create_lexer_and_document(["file"], "Some text /file:example.txt more text")
tokens = lexer.lex_document(document)
expected_tokens = [
("class:text", "S"),
("class:text", "o"),
("class:text", "m"),
("class:text", "e"),
("class:text", " "),
("class:text", "t"),
("class:text", "e"),
("class:text", "x"),
("class:text", "t"),
("class:text", " "),
("class:command", "/file:"),
("class:parameter", "example.txt"),
("class:text", " "),
("class:text", "m"),
("class:text", "o"),
("class:text", "r"),
("class:text", "e"),
("class:text", " "),
("class:text", "t"),
("class:text", "e"),
("class:text", "x"),
("class:text", "t"),
]
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_lex_document_multiple_commands():
lexer, document = create_lexer_and_document(
["command", "anothercommand"],
"/command:example1.txt more text /anothercommand:example2.txt",
)
tokens = lexer.lex_document(document)
expected_tokens = [
("class:command", "/command:"),
("class:parameter", "example1.txt"),
("class:text", " "),
("class:text", "m"),
("class:text", "o"),
("class:text", "r"),
("class:text", "e"),
("class:text", " "),
("class:text", "t"),
("class:text", "e"),
("class:text", "x"),
("class:text", "t"),
("class:text", " "),
("class:command", "/anothercommand:"),
("class:parameter", "example2.txt"),
]
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_lex_document_multiple_same_commands():
lexer, document = create_lexer_and_document(
["command"],
"/command:example1.txt more text /command:example2.txt",
)
tokens = lexer.lex_document(document)
expected_tokens = [
("class:command", "/command:"),
("class:parameter", "example1.txt"),
("class:text", " "),
("class:text", "m"),
("class:text", "o"),
("class:text", "r"),
("class:text", "e"),
("class:text", " "),
("class:text", "t"),
("class:text", "e"),
("class:text", "x"),
("class:text", "t"),
("class:text", " "),
("class:command", "/command:"),
("class:parameter", "example2.txt"),
]
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_lex_document_two_half_commands():
lexer, document = create_lexer_and_document(
["command"],
"/comma /com",
)
tokens = lexer.lex_document(document)
expected_tokens = [
("class:text", "/"),
("class:text", "c"),
("class:text", "o"),
("class:text", "m"),
("class:text", "m"),
("class:text", "a"),
("class:text", " "),
("class:command", "/com"),
]
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_lex_document_command_attached_to_pre_string():
lexer, document = create_lexer_and_document(
["command"],
"some/command:example.txt",
)
expected_tokens = [
("class:text", "s"),
("class:text", "o"),
("class:text", "m"),
("class:text", "e"),
("class:text", "/"),
("class:text", "c"),
("class:text", "o"),
("class:text", "m"),
("class:text", "m"),
("class:text", "a"),
("class:text", "n"),
("class:text", "d"),
("class:text", ":"),
("class:text", "e"),
("class:text", "x"),
("class:text", "a"),
("class:text", "m"),
("class:text", "p"),
("class:text", "l"),
("class:text", "e"),
("class:text", "."),
("class:text", "t"),
("class:text", "x"),
("class:text", "t"),
]
tokens = lexer.lex_document(document)
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_lex_document_partial_command_attached_to_pre_string():
lexer, document = create_lexer_and_document(
["command"],
"some/com",
)
tokens = lexer.lex_document(document)
expected_tokens = [
("class:text", "s"),
("class:text", "o"),
("class:text", "m"),
("class:text", "e"),
("class:text", "/"),
("class:text", "c"),
("class:text", "o"),
("class:text", "m"),
]
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_lex_document_no_command():
lexer, document = create_lexer_and_document([], "Some random text")
tokens = lexer.lex_document(document)
expected_tokens = [("class:text", character) for character in "Some random text"]
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_lex_document_ending_char_of_parameter_is_symbol():
lexer, document = create_lexer_and_document(
["command"],
"/command:example.txt/",
)
expected_tokens = [
("class:command", "/command:"),
("class:parameter", "example.txt/"),
]
tokens = lexer.lex_document(document)
actual_tokens = list(tokens(0))
assert actual_tokens == expected_tokens
def test_command_itself():
pattern = command_itself("file:")
matches = pattern.match("/file:example.txt")
assert matches is not None
assert matches.group(1) == "/file:"
def test_value_for_command():
pattern = value_for_command("file:")
matches = pattern.search("/file:example.txt")
assert matches is not None
assert matches.group(1) == "example.txt"
def test_completion_for_command():
pattern = completion_for_command("file:")
matches = pattern.search("/file:")
assert matches is not None
assert matches.group(1) == "file:"
tests/cli/prompt/test_prompt_validator.py
+37
View File
@@ -0,0 +1,37 @@
from unittest.mock import MagicMock, patch
import pytest
from goose.cli.prompt.prompt_validator import PromptValidator
from prompt_toolkit.validation import ValidationError
@pytest.fixture
def validator():
return PromptValidator()
@patch("prompt_toolkit.document.Document.text")
def test_validate_should_not_raise_error_when_input_is_none(document, validator):
try:
validator.validate(create_mock_document(None))
except Exception as e:
pytest.fail(f"An error was raised: {e}")
@patch("prompt_toolkit.document.Document.text", return_value="user typed something")
def test_validate_should_not_raise_error_when_user_has_input(document, validator):
try:
validator.validate(create_mock_document("user typed something"))
except Exception as e:
pytest.fail(f"An error was raised: {e}")
def test_validate_should_raise_validation_error_when_user_has_empty_input(validator):
with pytest.raises(ValidationError):
validator.validate(create_mock_document(""))
def create_mock_document(text: str) -> MagicMock:
document = MagicMock()
document.text = text
return document
tests/cli/prompt/test_user_input.py
+15
View File
@@ -0,0 +1,15 @@
from goose.cli.prompt.user_input import PromptAction, UserInput
def test_user_input_with_action_continue():
input = UserInput(action=PromptAction.CONTINUE, text="Hello")
assert input.to_continue() is True
assert input.to_exit() is False
assert input.text == "Hello"
def test_user_input_with_action_exit():
input = UserInput(action=PromptAction.EXIT)
assert input.to_continue() is False
assert input.to_exit() is True
assert input.text is None
tests/cli/test_config.py
+81
View File
@@ -0,0 +1,81 @@
from unittest.mock import patch
import pytest
from goose.cli.config import ensure_config, read_config, session_path, write_config
from goose.profile import default_profile
@pytest.fixture
def mock_profile_config_path(tmp_path):
with patch("goose.cli.config.PROFILES_CONFIG_PATH", tmp_path / "profiles.yaml") as mock_path:
yield mock_path
@pytest.fixture
def mock_default_model_configuration():
with patch(
"goose.cli.config.default_model_configuration", return_value=("provider", "processor", "accelerator")
) as mock_default_model_configuration:
yield mock_default_model_configuration
def test_read_write_config(mock_profile_config_path, profile_factory):
profiles = {
"profile1": profile_factory({"provider": "providerA"}),
}
write_config(profiles)
assert read_config() == profiles
def test_ensure_config_create_profiles_file_with_default_profile(
mock_profile_config_path, mock_default_model_configuration
):
assert not mock_profile_config_path.exists()
ensure_config(name="default")
assert mock_profile_config_path.exists()
assert read_config() == {"default": default_profile(*mock_default_model_configuration())}
def test_ensure_config_add_default_profile(mock_profile_config_path, profile_factory, mock_default_model_configuration):
existing_profile = profile_factory({"provider": "providerA"})
write_config({"profile1": existing_profile})
ensure_config(name="default")
assert read_config() == {
"profile1": existing_profile,
"default": default_profile(*mock_default_model_configuration()),
}
@patch("goose.cli.config.Confirm.ask", return_value=True)
def test_ensure_config_overwrite_default_profile(
mock_confirm, mock_profile_config_path, profile_factory, mock_default_model_configuration
):
existing_profile = profile_factory({"provider": "providerA"})
profile_name = "default"
write_config({profile_name: existing_profile})
expected_default_profile = default_profile(*mock_default_model_configuration())
assert ensure_config(name="default") == expected_default_profile
assert read_config() == {"default": expected_default_profile}
@patch("goose.cli.config.Confirm.ask", return_value=False)
def test_ensure_config_keep_original_default_profile(
mock_confirm, mock_profile_config_path, profile_factory, mock_default_model_configuration
):
existing_profile = profile_factory({"provider": "providerA"})
profile_name = "default"
write_config({profile_name: existing_profile})
assert ensure_config(name="default") == existing_profile
assert read_config() == {"default": existing_profile}
def test_session_path(mock_sessions_path):
assert session_path("session1") == mock_sessions_path / "session1.jsonl"
tests/cli/test_main.py
+80
View File
@@ -0,0 +1,80 @@
from datetime import datetime
from time import time
from unittest.mock import MagicMock, patch
import pytest
from click.testing import CliRunner
from exchange import Message
from goose.cli.main import goose_cli
@pytest.fixture
def mock_print():
with patch("goose.cli.main.print") as mock_print:
yield mock_print
@pytest.fixture
def mock_session_files_path(tmp_path):
with patch("goose.cli.main.SESSIONS_PATH", tmp_path) as session_files_path:
yield session_files_path
@pytest.fixture
def mock_session():
with patch("goose.cli.main.Session") as mock_session_class:
mock_session_instance = MagicMock()
mock_session_class.return_value = mock_session_instance
yield mock_session_class, mock_session_instance
def test_session_resume_command_with_session_name(mock_session):
mock_session_class, mock_session_instance = mock_session
runner = CliRunner()
runner.invoke(goose_cli, ["session", "resume", "session1", "--profile", "default"])
mock_session_class.assert_called_once_with(name="session1", profile="default")
mock_session_instance.run.assert_called_once()
def test_session_resume_command_without_session_name_without_session_files(
mock_print, mock_session_files_path, mock_session
):
_, mock_session_instance = mock_session
runner = CliRunner()
runner.invoke(goose_cli, ["session", "resume"])
mock_print.assert_called_with("No sessions found.")
mock_session_instance.run.assert_not_called()
def test_session_resume_command_without_session_name_use_latest_session(
mock_print, mock_session_files_path, mock_session, create_session_file
):
mock_session_class, mock_session_instance = mock_session
for index, session_name in enumerate(["first", "second"]):
create_session_file([Message.user("Hello1")], mock_session_files_path / f"{session_name}.jsonl", time() + index)
runner = CliRunner()
runner.invoke(goose_cli, ["session", "resume", "--profile", "default"])
second_file_path = mock_session_files_path / "second.jsonl"
mock_print.assert_called_once_with(f"Resuming most recent session: second from {second_file_path}")
mock_session_class.assert_called_once_with(name="second", profile="default")
mock_session_instance.run.assert_called_once()
def test_session_list_command(mock_print, mock_session_files_path, create_session_file):
create_session_file([Message.user("Hello")], mock_session_files_path / "abc.jsonl")
runner = CliRunner()
runner.invoke(goose_cli, ["session", "list"])
file_time = datetime.fromtimestamp(mock_session_files_path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S")
mock_print.assert_called_with(f"{file_time} abc")
def test_session_clear_command(mock_session_files_path, create_session_file):
for index, session_name in enumerate(["first", "second"]):
create_session_file([Message.user("Hello1")], mock_session_files_path / f"{session_name}.jsonl", time() + index)
runner = CliRunner()
runner.invoke(goose_cli, ["session", "clear", "--keep", "1"])
session_files = list(mock_session_files_path.glob("*.jsonl"))
assert len(session_files) == 1
assert session_files[0].stem == "second"
tests/cli/test_session.py
+134
View File
@@ -0,0 +1,134 @@
from unittest.mock import MagicMock, patch
import pytest
from exchange import Message
from goose.cli.prompt.goose_prompt_session import GoosePromptSession
from goose.cli.prompt.user_input import PromptAction, UserInput
from goose.cli.session import Session
from prompt_toolkit import PromptSession
SPECIFIED_SESSION_NAME = "mySession"
SESSION_NAME = "test"
@pytest.fixture
def mock_specified_session_name():
with patch.object(PromptSession, "prompt", return_value=SPECIFIED_SESSION_NAME) as specified_session_name:
yield specified_session_name
@pytest.fixture
def create_session_with_mock_configs(mock_sessions_path, exchange_factory, profile_factory):
with patch("goose.cli.session.build_exchange", return_value=exchange_factory()), patch(
"goose.cli.session.load_profile", return_value=profile_factory()
), patch("goose.cli.session.SessionNotifier") as mock_session_notifier, patch(
"goose.cli.session.load_provider", return_value="provider"
):
mock_session_notifier.return_value = MagicMock()
def create_session(session_attributes: dict = {}):
return Session(**session_attributes)
yield create_session
def test_session_does_not_extend_last_user_message_on_init(
create_session_with_mock_configs, mock_sessions_path, create_session_file
):
messages = [Message.user("Hello"), Message.assistant("Hi"), Message.user("Last should be removed")]
create_session_file(messages, mock_sessions_path / f"{SESSION_NAME}.jsonl")
session = create_session_with_mock_configs({"name": SESSION_NAME})
print("Messages after session init:", session.exchange.messages) # Debugging line
assert len(session.exchange.messages) == 2
assert [message.text for message in session.exchange.messages] == ["Hello", "Hi"]
def test_save_session_create_session(mock_sessions_path, create_session_with_mock_configs, mock_specified_session_name):
session = create_session_with_mock_configs()
session.exchange.messages.append(Message.assistant("Hello"))
session.save_session()
session_file = mock_sessions_path / f"{SPECIFIED_SESSION_NAME}.jsonl"
assert session_file.exists()
saved_messages = session.load_session()
assert len(saved_messages) == 1
assert saved_messages[0].text == "Hello"
def test_save_session_resume_session_new_file(
mock_sessions_path, create_session_with_mock_configs, mock_specified_session_name, create_session_file
):
with patch("goose.cli.session.confirm", return_value=False):
existing_messages = [Message.assistant("existing_message")]
existing_session_file = mock_sessions_path / f"{SESSION_NAME}.jsonl"
create_session_file(existing_messages, existing_session_file)
new_session_file = mock_sessions_path / f"{SPECIFIED_SESSION_NAME}.jsonl"
assert not new_session_file.exists()
session = create_session_with_mock_configs({"name": SESSION_NAME})
session.exchange.messages.append(Message.assistant("new_message"))
session.save_session()
assert new_session_file.exists()
assert existing_session_file.exists()
saved_messages = session.load_session()
assert [message.text for message in saved_messages] == ["existing_message", "new_message"]
def test_save_session_resume_session_existing_session_file(
mock_sessions_path, create_session_with_mock_configs, create_session_file
):
with patch("goose.cli.session.confirm", return_value=True):
existing_messages = [Message.assistant("existing_message")]
existing_session_file = mock_sessions_path / f"{SESSION_NAME}.jsonl"
create_session_file(existing_messages, existing_session_file)
session = create_session_with_mock_configs({"name": SESSION_NAME})
session.exchange.messages.append(Message.assistant("new_message"))
session.save_session()
saved_messages = session.load_session()
assert [message.text for message in saved_messages] == ["existing_message", "new_message"]
def test_process_first_message_return_message(create_session_with_mock_configs):
session = create_session_with_mock_configs()
with patch.object(
GoosePromptSession, "get_user_input", return_value=UserInput(action=PromptAction.CONTINUE, text="Hello")
):
message = session.process_first_message()
assert message.text == "Hello"
assert len(session.exchange.messages) == 0
def test_process_first_message_to_exit(create_session_with_mock_configs):
session = create_session_with_mock_configs()
with patch.object(GoosePromptSession, "get_user_input", return_value=UserInput(action=PromptAction.EXIT)):
message = session.process_first_message()
assert message is None
def test_process_first_message_return_last_exchange_message(create_session_with_mock_configs):
session = create_session_with_mock_configs()
session.exchange.messages.append(Message.user("Hi"))
message = session.process_first_message()
assert message.text == "Hi"
assert len(session.exchange.messages) == 0
def test_generate_session_name(create_session_with_mock_configs):
session = create_session_with_mock_configs()
with patch.object(GoosePromptSession, "get_save_session_name", return_value=SPECIFIED_SESSION_NAME):
session.generate_session_name()
assert session.name == SPECIFIED_SESSION_NAME