Skip to content

Commit

Permalink
fix: fix conversation list order in picker, lazily load conversation …
Browse files Browse the repository at this point in the history
…metadata, add get_user_conversations(), add ?limit=<int> to /api/conversations and use it in webui
  • Loading branch information
ErikBjare committed Sep 22, 2024
1 parent e1b881a commit 9c53aa0
Show file tree
Hide file tree
Showing 9 changed files with 185 additions and 63 deletions.
41 changes: 29 additions & 12 deletions gptme/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import urllib.parse
from collections.abc import Generator
from datetime import datetime
from itertools import islice
from pathlib import Path
from typing import Literal

Expand All @@ -28,7 +29,7 @@
from .dirs import get_logs_dir
from .init import init, init_logging
from .llm import reply
from .logmanager import LogManager, _conversations
from .logmanager import Conversation, LogManager, get_user_conversations
from .message import Message
from .models import get_model
from .prompts import get_prompt
Expand Down Expand Up @@ -407,16 +408,24 @@ def get_name(name: str) -> Path:
return logpath


def get_logfile(name: str | Literal["random", "resume"], interactive=True) -> Path:
def get_logfile(
name: str | Literal["random", "resume"], interactive=True, limit=20
) -> Path:
# let user select between starting a new conversation and loading a previous one
# using the library
title = "New conversation or load previous? "
NEW_CONV = "New conversation"
prev_conv_files = list(reversed(_conversations()))
LOAD_MORE = "Load more"
gen_convs = get_user_conversations()
convs: list[Conversation] = []
try:
convs.append(next(gen_convs))
except StopIteration:
pass

if name == "resume":
if prev_conv_files:
return prev_conv_files[0].parent / "conversation.jsonl"
if convs:
return Path(convs[0].path)
else:
raise ValueError("No previous conversations to resume")

Expand All @@ -426,24 +435,32 @@ def get_logfile(name: str | Literal["random", "resume"], interactive=True) -> Pa
# return "-test-" in name or name.startswith("test-")
# prev_conv_files = [f for f in prev_conv_files if not is_test(f.parent.name)]

NEWLINE = "\n"
# load more conversations
convs.extend(islice(gen_convs, limit - 1))

prev_convs = [
f"{f.parent.name:30s} \t{epoch_to_age(f.stat().st_mtime)} \t{len(f.read_text().split(NEWLINE)):5d} msgs"
for f in prev_conv_files
f"{conv.name:30s} \t{epoch_to_age(conv.modified)} \t{conv.messages:5d} msgs"
for conv in convs
]

# don't run pick in tests/non-interactive mode, or if the user specifies a name
if interactive and name in ["random"]:
options = [
NEW_CONV,
] + prev_convs
options = (
[
NEW_CONV,
]
+ prev_convs
+ [LOAD_MORE]
)

index: int
_, index = pick(options, title) # type: ignore
if index == 0:
logdir = get_name(name)
elif index == len(options) - 1:
return get_logfile(name, interactive, limit + 100)
else:
logdir = get_logs_dir() / prev_conv_files[index - 1].parent
logdir = get_logs_dir() / convs[index - 1].name
else:
logdir = get_name(name)

Expand Down
71 changes: 50 additions & 21 deletions gptme/logmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,9 @@
import textwrap
from collections.abc import Generator
from copy import copy
from dataclasses import dataclass
from datetime import datetime
from itertools import zip_longest
from itertools import islice, zip_longest
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any, Literal, TypeAlias
Expand Down Expand Up @@ -288,40 +289,68 @@ def to_dict(self, branches=False) -> dict:
return d


def _conversations() -> list[Path]:
def _conversation_files() -> list[Path]:
# NOTE: only returns the main conversation, not branches (to avoid duplicates)
# returns the most recent first
# returns the conversation files sorted by modified time (newest first)
logsdir = get_logs_dir()
return list(
sorted(logsdir.glob("*/conversation.jsonl"), key=lambda f: -f.stat().st_mtime)
)


def get_conversations() -> Generator[dict, None, None]:
for conv_fn in _conversations():
msgs = []
msgs = _read_jsonl(conv_fn)
modified = conv_fn.stat().st_mtime
first_timestamp = msgs[0].timestamp.timestamp() if msgs else modified
yield {
"name": f"{conv_fn.parent.name}",
"path": str(conv_fn),
"created": first_timestamp,
"modified": modified,
"messages": len(msgs),
"branches": 1 + len(list(conv_fn.parent.glob("branches/*.jsonl"))),
}
@dataclass
class Conversation:
name: str
path: str
created: float
modified: float
messages: int
branches: int


def _read_jsonl(path: PathLike) -> list[Message]:
msgs = []
def get_conversations() -> Generator[Conversation, None, None]:
"""Returns all conversations, excluding ones used for testing, evals, etc."""
for conv_fn in _conversation_files():
msgs = _read_jsonl(conv_fn, limit=1)
# TODO: can we avoid reading the entire file? maybe wont even be used, due to user convo filtering
len_msgs = conv_fn.read_text().count("}\n{")
assert len(msgs) <= 1
modified = conv_fn.stat().st_mtime
first_timestamp = msgs[0].timestamp.timestamp() if msgs else modified
yield Conversation(
name=f"{conv_fn.parent.name}",
path=str(conv_fn),
created=first_timestamp,
modified=modified,
messages=len_msgs,
branches=1 + len(list(conv_fn.parent.glob("branches/*.jsonl"))),
)


def get_user_conversations() -> Generator[Conversation, None, None]:
"""Returns all user conversations, excluding ones used for testing, evals, etc."""
for conv in get_conversations():
if any(conv.name.startswith(prefix) for prefix in ["tmp", "test-"]) or any(
substr in conv.name for substr in ["gptme-evals-"]
):
continue
yield conv


def _gen_read_jsonl(path: PathLike) -> Generator[Message, None, None]:
with open(path) as file:
for line in file.readlines():
json_data = json.loads(line)
if "timestamp" in json_data:
json_data["timestamp"] = datetime.fromisoformat(json_data["timestamp"])
msgs.append(Message(**json_data))
return msgs
yield Message(**json_data)


def _read_jsonl(path: PathLike, limit=None) -> list[Message]:
gen = _gen_read_jsonl(path)
if limit:
gen = islice(gen, limit) # type: ignore
return list(gen)


def _write_jsonl(path: PathLike, msgs: list[Message]) -> None:
Expand Down
8 changes: 8 additions & 0 deletions gptme/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@

logger = logging.getLogger(__name__)

# max tokens allowed in a single system message
# if you hit this limit, you and/or I f-ed up, and should make the message shorter
# maybe we should make it possible to store long outputs in files, and link/summarize it/preview it in the message
max_system_len = 20000


@dataclass(frozen=True, eq=False)
class Message:
Expand Down Expand Up @@ -51,6 +56,9 @@ class Message:

def __post_init__(self):
assert isinstance(self.timestamp, datetime)
if self.role == "system":
if (length := len_tokens(self)) >= max_system_len:
logger.warning(f"System message too long: {length} tokens")

def __repr__(self):
content = textwrap.shorten(self.content, 20, placeholder="...")
Expand Down
3 changes: 2 additions & 1 deletion gptme/server/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
Server for gptme.
"""

from .api import create_app, main
from .api import create_app
from .cli import main

__all__ = ["main", "create_app"]
14 changes: 5 additions & 9 deletions gptme/server/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,15 @@
from contextlib import redirect_stdout
from datetime import datetime
from importlib import resources
from itertools import islice

import flask
from flask import current_app
from flask import current_app, request

from ..commands import execute_cmd
from ..dirs import get_logs_dir
from ..llm import reply
from ..logmanager import LogManager, get_conversations
from ..logmanager import LogManager, get_user_conversations
from ..message import Message
from ..models import get_model
from ..tools import execute_msg
Expand All @@ -32,7 +33,8 @@ def api_root():

@api.route("/api/conversations")
def api_conversations():
conversations = list(get_conversations())
limit = int(request.args.get("limit", 100))
conversations = list(islice(get_user_conversations(), limit))
return flask.jsonify(conversations)


Expand Down Expand Up @@ -149,9 +151,3 @@ def create_app() -> flask.Flask:
app = flask.Flask(__name__, static_folder=static_path)
app.register_blueprint(api)
return app


def main() -> None:
"""Run the Flask app."""
app = create_app()
app.run(debug=True)
12 changes: 6 additions & 6 deletions gptme/server/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,18 +3,20 @@
import click

from ..init import init, init_logging
from .api import create_app

logger = logging.getLogger(__name__)


@click.command("gptme-server")
@click.option("-v", "--verbose", is_flag=True, help="Verbose output.")
@click.option("--debug", is_flag=True, help="Debug mode")
@click.option("-v", "--verbose", is_flag=True, help="Verbose output")
@click.option(
"--model",
default=None,
help="Model to use by default, can be overridden in each request.",
)
def main(verbose: bool, model: str | None): # pragma: no cover
def main(debug: bool, verbose: bool, model: str | None): # pragma: no cover
"""
Starts a server and web UI for gptme.
Expand All @@ -34,7 +36,5 @@ def main(verbose: bool, model: str | None): # pragma: no cover
exit(1)
click.echo("Initialization complete, starting server")

# noreorder
from gptme.server.api import main as server_main # fmt: skip

server_main()
app = create_app()
app.run(debug=debug)
29 changes: 17 additions & 12 deletions gptme/tools/chats.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,14 @@
import logging
from pathlib import Path
from textwrap import indent
from typing import TYPE_CHECKING

from ..message import Message
from .base import ToolSpec

if TYPE_CHECKING:
from ..logmanager import LogManager

logger = logging.getLogger(__name__)


Expand All @@ -33,7 +37,9 @@ def _get_matching_messages(log_manager, query: str, system=False) -> list[Messag
]


def _summarize_conversation(log_manager, include_summary: bool) -> list[str]:
def _summarize_conversation(
log_manager: "LogManager", include_summary: bool
) -> list[str]:
"""Summarize a conversation."""
# noreorder
from ..llm import summarize as llm_summarize # fmt: skip
Expand Down Expand Up @@ -80,11 +86,10 @@ def list_chats(max_results: int = 5, include_summary: bool = False) -> None:

print(f"Recent conversations (showing up to {max_results}):")
for i, conv in enumerate(conversations, 1):
print(f"\n{i}. {conv['name']}")
if "created_at" in conv:
print(f" Created: {conv['created_at']}")
print(f"\n{i}. {conv.name}")
print(f" Created: {conv.created}")

log_path = Path(conv["path"])
log_path = Path(conv.path)
log_manager = LogManager.load(log_path)

summary_lines = _summarize_conversation(log_manager, include_summary)
Expand All @@ -101,19 +106,19 @@ def search_chats(query: str, max_results: int = 5, system=False) -> None:
system (bool): Whether to include system messages in the search.
"""
# noreorder
from ..logmanager import LogManager, get_conversations # fmt: skip
from ..logmanager import LogManager, get_user_conversations # fmt: skip

results = []
for conv in get_conversations():
log_path = Path(conv["path"])
results: list[dict] = []
for conv in get_user_conversations():
log_path = Path(conv.path)
log_manager = LogManager.load(log_path)

matching_messages = _get_matching_messages(log_manager, query, system)

if matching_messages:
results.append(
{
"conversation": conv["name"],
"conversation": conv.name,
"log_manager": log_manager,
"matching_messages": matching_messages,
}
Expand Down Expand Up @@ -165,8 +170,8 @@ def read_chat(conversation: str, max_results: int = 5, incl_system=False) -> Non
conversations = list(get_conversations())

for conv in conversations:
if conv["name"] == conversation:
log_path = Path(conv["path"])
if conv.name == conversation:
log_path = Path(conv.path)
logmanager = LogManager.load(log_path)
print(f"Reading conversation: {conversation}")
i = 0
Expand Down
Loading

0 comments on commit 9c53aa0

Please sign in to comment.