Skip to content

Commit

Permalink
refactor: work on programmatic interface, refactored LogManager into …
Browse files Browse the repository at this point in the history 
…mutable manager and immutable Log dataclass, added wip treeofthought script
  • Loading branch information
@ErikBjare
ErikBjare committed Oct 15, 2024
1 parent 1724534 commit d421cc8
Show file tree
Hide file tree
Showing 9 changed files with 262 additions and 146 deletions.
46 changes: 26 additions & 20 deletions gptme/chat.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
from .init import init
from .interrupt import clear_interruptible, set_interruptible
from .llm import reply
from .logmanager import LogManager
from .logmanager import Log, LogManager, prepare_messages
from .message import Message
from .models import get_model
from .tools import ToolUse, execute_msg, has_tool
Expand Down Expand Up @@ -59,9 +59,7 @@ def chat(
stream = False

console.log(f"Using logdir {path_with_tilde(logdir)}")
log = LogManager.load(
logdir, initial_msgs=initial_msgs, show_hidden=show_hidden, create=True
)
manager = LogManager.load(logdir, initial_msgs=initial_msgs, create=True)

# change to workspace directory
# use if exists, create if @log, or use given path
Expand All @@ -82,13 +80,13 @@ def chat(
# check if message is already in log, such as upon resume
if (
workspace_prompt
and workspace_prompt not in [m.content for m in log]
and "user" not in [m.role for m in log]
and workspace_prompt not in [m.content for m in manager.log]
and "user" not in [m.role for m in manager.log]
):
log.append(Message("system", workspace_prompt, hide=True, quiet=True))
manager.append(Message("system", workspace_prompt, hide=True, quiet=True))

# print log
log.print()
manager.log.print(show_hidden=show_hidden)
console.print("--- ^^^ past messages ^^^ ---")

# main loop
Expand All @@ -99,34 +97,39 @@ def chat(
msg = prompt_msgs.pop(0)
if not msg.content.startswith("/"):
msg = _include_paths(msg)
log.append(msg)
manager.append(msg)
# if prompt is a user-command, execute it
if execute_cmd(msg, log):
if execute_cmd(msg, manager):
continue

# Generate and execute response for this prompt
while True:
set_interruptible()
try:
response_msgs = list(step(log, no_confirm, stream=stream))
response_msgs = list(step(manager, no_confirm, stream=stream))
except KeyboardInterrupt:
console.log("Interrupted. Stopping current execution.")
log.append(Message("system", "Interrupted"))
manager.append(Message("system", "Interrupted"))
break
finally:
clear_interruptible()

for response_msg in response_msgs:
log.append(response_msg)
manager.append(response_msg)
# run any user-commands, if msg is from user
if response_msg.role == "user" and execute_cmd(
response_msg, log
response_msg, manager
):
break

# Check if there are any runnable tools left
last_content = next(
(m.content for m in reversed(log) if m.role == "assistant"), ""
(
m.content
for m in reversed(manager.log)
if m.role == "assistant"
),
"",
)
if not any(
tooluse.is_runnable
Expand All @@ -148,19 +151,22 @@ def chat(

# ask for input if no prompt, generate reply, and run tools
clear_interruptible() # Ensure we're not interruptible during user input
for msg in step(log, no_confirm, stream=stream): # pragma: no cover
log.append(msg)
for msg in step(manager, no_confirm, stream=stream): # pragma: no cover
manager.append(msg)
# run any user-commands, if msg is from user
if msg.role == "user" and execute_cmd(msg, log):
if msg.role == "user" and execute_cmd(msg, manager):
break


def step(
log: LogManager,
log: Log | LogManager,
no_confirm: bool,
stream: bool = True,
) -> Generator[Message, None, None]:
"""Runs a single pass of the chat."""
if isinstance(log, LogManager):
log = log.log

# If last message was a response, ask for input.
# If last message was from the user (such as from crash/edited log),
# then skip asking for input and generate response
Expand All @@ -184,7 +190,7 @@ def step(
set_interruptible()
try:
# performs reduction/context trimming, if necessary
msgs = log.prepare_messages()
msgs = prepare_messages(log.messages)

for m in msgs:
logger.debug(f"Prepared message: {m}")
Expand Down
4 changes: 2 additions & 2 deletions gptme/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from .dirs import get_logs_dir
from .init import init_logging
from .interrupt import handle_keyboard_interrupt
from .logmanager import Conversation, get_user_conversations
from .logmanager import ConversationMeta, get_user_conversations
from .message import Message
from .prompts import get_prompt
from .tools import all_tools, init_tools
Expand Down Expand Up @@ -273,7 +273,7 @@ def pick_log(limit=20) -> Path: # pragma: no cover
NEW_CONV = "New conversation"
LOAD_MORE = "Load more"
gen_convs = get_user_conversations()
convs: list[Conversation] = []
convs: list[ConversationMeta] = []

# load conversations
convs.extend(islice(gen_convs, limit))
Expand Down
66 changes: 32 additions & 34 deletions gptme/commands.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from typing import Literal

from . import llm
from .logmanager import LogManager
from .logmanager import LogManager, prepare_messages
from .message import (
Message,
len_tokens,
Expand Down Expand Up @@ -68,7 +68,7 @@ def execute_cmd(msg: Message, log: LogManager) -> bool:


def handle_cmd(
cmd: str, log: LogManager, no_confirm: bool
cmd: str, manager: LogManager, no_confirm: bool
) -> Generator[Message, None, None]:
"""Handles a command."""
cmd = cmd.lstrip("/")
Expand All @@ -77,44 +77,44 @@ def handle_cmd(
full_args = cmd.split(" ", 1)[1] if " " in cmd else ""
match name:
case "log":
log.undo(1, quiet=True)
log.print(show_hidden="--hidden" in args)
manager.undo(1, quiet=True)
manager.log.print(show_hidden="--hidden" in args)
case "rename":
log.undo(1, quiet=True)
log.write()
manager.undo(1, quiet=True)
manager.write()
# rename the conversation
print("Renaming conversation (enter empty name to auto-generate)")
new_name = args[0] if args else input("New name: ")
rename(log, new_name, ask=not no_confirm)
rename(manager, new_name, ask=not no_confirm)
case "fork":
# fork the conversation
new_name = args[0] if args else input("New name: ")
log.fork(new_name)
manager.fork(new_name)
case "summarize":
msgs = log.prepare_messages()
msgs = prepare_messages(manager.log.messages)
msgs = [m for m in msgs if not m.hide]
summary = llm.summarize(msgs)
print(f"Summary: {summary}")
case "edit":
# edit previous messages
# first undo the '/edit' command itself
log.undo(1, quiet=True)
yield from edit(log)
manager.undo(1, quiet=True)
yield from edit(manager)
case "undo":
# undo the '/undo' command itself
log.undo(1, quiet=True)
manager.undo(1, quiet=True)
# if int, undo n messages
n = int(args[0]) if args and args[0].isdigit() else 1
log.undo(n)
manager.undo(n)
case "exit":
log.undo(1, quiet=True)
log.write()
manager.undo(1, quiet=True)
manager.write()
sys.exit(0)
case "replay":
log.undo(1, quiet=True)
log.write()
manager.undo(1, quiet=True)
manager.write()
print("Replaying conversation...")
for msg in log.log:
for msg in manager.log:
if msg.role == "assistant":
for reply_msg in execute_msg(msg, ask=True):
print_msg(reply_msg, oneline=False)
Expand All @@ -124,16 +124,16 @@ def handle_cmd(
yield msg
yield from execute_msg(msg, ask=not no_confirm)
case "tokens":
log.undo(1, quiet=True)
n_tokens = len_tokens(log.log)
manager.undo(1, quiet=True)
n_tokens = len_tokens(manager.log.messages)
print(f"Tokens used: {n_tokens}")
model = get_model()
if model:
print(f"Model: {model.model}")
if model.price_input:
print(f"Cost (input): ${n_tokens * model.price_input / 1_000_000}")
case "tools":
log.undo(1, quiet=True)
manager.undo(1, quiet=True)
print("Available tools:")
for tool in loaded_tools:
print(
Expand All @@ -148,18 +148,18 @@ def handle_cmd(
if tooluse.is_runnable:
yield from tooluse.execute(ask=not no_confirm)
else:
if log.log[-1].content.strip() == "/help":
if manager.log[-1].content.strip() == "/help":
# undo the '/help' command itself
log.undo(1, quiet=True)
log.write()
manager.undo(1, quiet=True)
manager.write()
help()
else:
print("Unknown command")


def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover
def edit(manager: LogManager) -> Generator[Message, None, None]: # pragma: no cover
# generate editable toml of all messages
t = msgs_to_toml(reversed(log.log)) # type: ignore
t = msgs_to_toml(reversed(manager.log)) # type: ignore
res = None
while not res:
t = edit_text_with_editor(t, "toml")
Expand All @@ -172,26 +172,24 @@ def edit(log: LogManager) -> Generator[Message, None, None]: # pragma: no cover
except KeyboardInterrupt:
yield Message("system", "Interrupted")
return
log.edit(list(reversed(res)))
# now we need to redraw the log so the user isn't seeing stale messages in their buffer
# log.print()
manager.edit(list(reversed(res)))
print("Applied edited messages, write /log to see the result")


def rename(log: LogManager, new_name: str, ask: bool = True):
def rename(manager: LogManager, new_name: str, ask: bool = True):
if new_name in ["", "auto"]:
new_name = llm.generate_name(log.prepare_messages())
new_name = llm.generate_name(prepare_messages(manager.log.messages))
assert " " not in new_name
print(f"Generated name: {new_name}")
if ask:
confirm = ask_execute("Confirm?")
if not confirm:
print("Aborting")
return
log.rename(new_name, keep_date=True)
manager.rename(new_name, keep_date=True)
else:
log.rename(new_name, keep_date=False)
print(f"Renamed conversation to {log.logfile.parent}")
manager.rename(new_name, keep_date=False)
print(f"Renamed conversation to {manager.logfile.parent}")


def _gen_help(incl_langtags: bool = True) -> Generator[str, None, None]:
Expand Down
Loading

0 comments on commit d421cc8

Please sign in to comment.