Skip to content

Commit

Permalink
refactor: refactor tools, codeblock, and tooluse (#113)
Browse files Browse the repository at this point in the history
* refactor: refactor tools, streamline amending __doc__ for tools, refactor type_o_tooluse to ToolUse.from_type

* Apply suggestions from code review

* refactor: move Codeblock to new file codeblock.py

* Apply suggestions from code review

* misc fixes

* fix: more codeblock refactoring

* Apply suggestions from code review

* fixed tests

* more fixes

* fix: removed save command

* Apply suggestions from code review

* removed unused test

* Apply suggestions from code review

* fix: more refactoring, use block_types to register commands for python, shell, etc

* fix: minor improvements to /help and /tools formatting

* fix: misc improvements to docs

* Apply suggestions from code review
  • Loading branch information
ErikBjare authored Sep 9, 2024
1 parent dd6fff7 commit 0cad5ca
Show file tree
Hide file tree
Showing 25 changed files with 429 additions and 422 deletions.
58 changes: 12 additions & 46 deletions docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -6,66 +6,32 @@ Here is the API reference for ``gptme``.
.. toctree::
:maxdepth: 2

gptme package
^^^^^^^^^^^^^
core
----

Some of the main classes in ``gptme``.
Some of the core classes in ``gptme``.

gptme.message
-------------

.. automodule:: gptme.message
.. autoclass:: gptme.message.Message
:members:

gptme.logmanager
----------------

.. automodule:: gptme.logmanager
.. automodule:: gptme.codeblock
:members:

gptme.server
^^^^^^^^^^^^

Endpoint functions for the server.

.. automodule:: gptme.server
.. automodule:: gptme.logmanager
:members:
:undoc-members:

gptme.tools
^^^^^^^^^^^
tools
-----

Tools available to gptme.
Supporting classes and functions for creating and using tools.

.. automodule:: gptme.tools
:members:

gptme.tools.shell
-----------------

.. automodule:: gptme.tools.shell
:members:

Python
server
------

.. automodule:: gptme.tools.python
:members:

gptme.tools.context
-------------------

.. automodule:: gptme.tools.context
:members:

gptme.tools.save
----------------
See `Server <server.html>`_ for more information.

.. automodule:: gptme.tools.save
:members:

gptme.tools.patch
-----------------

.. automodule:: gptme.tools.patch
.. automodule:: gptme.server
:members:
26 changes: 12 additions & 14 deletions gptme/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,12 @@
from rich import print # noqa: F401
from rich.console import Console

from .commands import CMDFIX, action_descriptions, execute_cmd
from .commands import (
CMDFIX,
_gen_help,
action_descriptions,
execute_cmd,
)
from .config import get_workspace_prompt
from .constants import MULTIPROMPT_SEPARATOR, PROMPT_USER
from .dirs import get_logs_dir
Expand All @@ -34,15 +39,9 @@
logger = logging.getLogger(__name__)
print_builtin = __builtins__["print"] # type: ignore

# TODO: these are a bit redundant/incorrect
LLMChoice = Literal["openai", "anthropic", "local"]
ModelChoice = Literal["gpt-3.5-turbo", "gpt-4", "gpt-4-1106-preview"]


script_path = Path(os.path.realpath(__file__))
action_readme = "\n".join(
f" {CMDFIX}{cmd:11s} {desc}." for cmd, desc in action_descriptions.items()
)
commands_help = "\n".join(_gen_help(incl_langtags=False))


docstring = f"""
Expand All @@ -55,7 +54,7 @@
The chat offers some commands that can be used to interact with the system:
\b
{action_readme}"""
{commands_help}"""


@click.command(help=docstring)
Expand Down Expand Up @@ -121,7 +120,7 @@ def main(
prompts: list[str],
prompt_system: str,
name: str,
model: ModelChoice,
model: str,
stream: bool,
verbose: bool,
no_confirm: bool,
Expand Down Expand Up @@ -273,13 +272,12 @@ def chat(
# then exit
elif not interactive:
# noreorder
from .tools import is_supported_codeblock_tool # fmt: skip
from .tools import is_supported_langtag # fmt: skip

# continue if we can run tools on the last message
runnable = False
if codeblock := log.get_last_code_block("assistant", history=1):
lang, _ = codeblock
if is_supported_codeblock_tool(lang):
if codeblock := log.get_last_codeblock("assistant", history=1):
if is_supported_langtag(codeblock.lang):
runnable = True
if not runnable:
logger.info("Non-interactive and exhausted prompts, exiting")
Expand Down
81 changes: 81 additions & 0 deletions gptme/codeblock.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
from collections.abc import Generator
from dataclasses import dataclass
from xml.etree import ElementTree


@dataclass
class Codeblock:
lang: str
content: str
path: str | None = None

# init path in __post_init__ if path is None and lang is pathy
def __post_init__(self):
if self.path is None and self.is_filename:
self.path = self.lang

def to_markdown(self) -> str:
return f"```{self.lang}\n{self.content}\n```"

def to_xml(self) -> str:
return f'<codeblock lang="{self.lang}" path="{self.path}">\n{self.content}\n</codeblock>'

@classmethod
def from_markdown(cls, content: str) -> "Codeblock":
if content.strip().startswith("```"):
content = content[3:]
if content.strip().endswith("```"):
content = content[:-3]
lang = content.splitlines()[0].strip()
return cls(lang, content[len(lang) :])

@classmethod
def from_xml(cls, content: str) -> "Codeblock":
"""
Example:
<codeblock lang="python" path="example.py">
print("Hello, world!")
</codeblock>
"""
root = ElementTree.fromstring(content)
return cls(root.attrib["lang"], root.text or "", root.attrib.get("path"))

@property
def is_filename(self) -> bool:
return "." in self.lang or "/" in self.lang

@classmethod
def iter_from_markdown(cls, markdown: str) -> list["Codeblock"]:
return list(_extract_codeblocks(markdown))


def _extract_codeblocks(markdown: str) -> Generator[Codeblock, None, None]:
# speed check (early exit): check if message contains a code block
backtick_count = markdown.count("```")
if backtick_count < 2:
return

lines = markdown.split("\n")
stack: list[str] = []
current_block = []
current_lang = ""

for line in lines:
stripped_line = line.strip()
if stripped_line.startswith("```"):
if not stack: # Start of a new block
stack.append(stripped_line[3:])
current_lang = stripped_line[3:]
elif stripped_line[3:] and stack[-1] != stripped_line[3:]: # Nested start
current_block.append(line)
stack.append(stripped_line[3:])
else: # End of a block
if len(stack) == 1: # Outermost block
yield Codeblock(current_lang, "\n".join(current_block))
current_block = []
current_lang = ""
else: # Nested end
current_block.append(line)
stack.pop()
elif stack:
current_block.append(line)
93 changes: 45 additions & 48 deletions gptme/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@
import re
import sys
from collections.abc import Generator
from pathlib import Path
from time import sleep
from typing import Literal

from . import llm
from .codeblock import Codeblock
from .constants import CMDFIX
from .logmanager import LogManager
from .message import (
Expand All @@ -18,9 +18,9 @@
)
from .models import get_model
from .tools import (
execute_codeblock,
execute_msg,
execute_python,
execute_shell,
is_supported_langtag,
loaded_tools,
)
from .useredit import edit_text_with_editor
Expand All @@ -35,9 +35,6 @@
"fork",
"summarize",
"context",
"save",
"shell",
"python",
"replay",
"undo",
"impersonate",
Expand All @@ -54,9 +51,6 @@
"rename": "Rename the conversation",
"fork": "Create a copy of the conversation with a new name",
"summarize": "Summarize the conversation",
"save": "Save the last code block to a file",
"shell": "Execute shell code",
"python": "Execute Python code",
"replay": "Re-execute codeblocks in the conversation, wont store output in log",
"impersonate": "Impersonate the assistant",
"tokens": "Show the number of tokens used",
Expand Down Expand Up @@ -89,11 +83,6 @@ def handle_cmd(
name, *args = re.split(r"[\n\s]", cmd)
full_args = cmd.split(" ", 1)[1] if " " in cmd else ""
match name:
# TODO: rewrite to auto-register tools using block_types
case "bash" | "sh" | "shell":
yield from execute_shell(full_args, ask=not no_confirm, args=[])
case "python" | "py":
yield from execute_python(full_args, ask=not no_confirm, args=[])
case "log":
log.undo(1, quiet=True)
log.print(show_hidden="--hidden" in args)
Expand Down Expand Up @@ -124,11 +113,6 @@ def handle_cmd(
# if int, undo n messages
n = int(args[0]) if args and args[0].isdigit() else 1
log.undo(n)
case "save":
# undo
log.undo(1, quiet=True)
filename = args[0] if args else input("Filename: ")
save(log, filename)
case "exit":
sys.exit(0)
case "replay":
Expand Down Expand Up @@ -159,17 +143,23 @@ def handle_cmd(
for tool in loaded_tools:
print(
f"""
- {tool.name} ({tool.desc.rstrip(".")})
tokens (example): {len_tokens(tool.examples)}
""".strip()
# {tool.name}
{tool.desc.rstrip(".")}
tokens (example): {len_tokens(tool.examples)}"""
)
case _:
if log.log[-1].content != f"{CMDFIX}help":
print("Unknown command")
# undo the '/help' command itself
log.undo(1, quiet=True)
log.write()
help()
# the case for python, shell, and other block_types supported by tools
if is_supported_langtag(name):
yield from execute_codeblock(
Codeblock(name, full_args), ask=not no_confirm
)
else:
if log.log[-1].content != f"{CMDFIX}help":
print("Unknown command")
# undo the '/help' command itself
log.undo(1, quiet=True)
log.write()
help()


def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover
Expand All @@ -193,22 +183,6 @@ def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover
print("Applied edited messages, write /log to see the result")


def save(log: LogManager, filename: str):
# save the most recent code block to a file
codeblock = log.get_last_code_block()
if not codeblock:
print("No code block found")
return
_, content = codeblock
if Path(filename).exists():
confirm = ask_execute("File already exists, overwrite?", default=False)
if not confirm:
return
with open(filename, "w") as f:
f.write(content)
print(f"Saved code block to {filename}")


def rename(log: LogManager, new_name: str, ask: bool = True):
if new_name in ["", "auto"]:
new_name = llm.generate_name(log.prepare_messages())
Expand All @@ -225,8 +199,31 @@ def rename(log: LogManager, new_name: str, ask: bool = True):
print(f"Renamed conversation to {log.logfile.parent}")


def help():
longest_cmd = max(len(cmd) for cmd in COMMANDS)
print("Available commands:")
def _gen_help(incl_langtags: bool = True) -> Generator[str, None, None]:
yield "Available commands:"
max_cmdlen = max(len(cmd) for cmd in COMMANDS)
for cmd, desc in action_descriptions.items():
print(f" /{cmd.ljust(longest_cmd)} {desc}")
yield f" /{cmd.ljust(max_cmdlen)} {desc}"

if incl_langtags:
yield ""
yield "To execute code with supported tools, use the following syntax:"
yield " /<langtag> <code>"
yield ""
yield "Example:"
yield " /sh echo hello"
yield " /python print('hello')"
yield ""
yield "Supported langtags:"
for tool in loaded_tools:
if tool.block_types:
yield f" - {tool.block_types[0]}" + (
f" (alias: {', '.join(tool.block_types[1:])})"
if len(tool.block_types) > 1
else ""
)


def help():
for line in _gen_help():
print(line)
Loading

0 comments on commit 0cad5ca

Please sign in to comment.