Skip to content

Commit

Permalink
feat: implement gptme-util CLI for utilities (#261)
Browse files Browse the repository at this point in the history
* feat(gptme-util): implement tools subcommands

- Added tools list and info commands
- Added language tag support
- Reused existing code from commands.py
- Added initial tests for utility CLI
- Moved util.py to util/__init__.py for better organization

* fix(gptme-util): improve error handling and tests

- Added empty list handling in chats_list
- Added model validation in tokens_count
- Added comprehensive test cases
- Skip context test until module is available

* refactor: remove duplicate ModelDictMeta definition

- Import _ModelDictMeta from models.py instead of duplicating definition
- Fixes code duplication identified by pylint

* wip

* fix: fixed recursive import

* docs: fixed docs

* docs: added gptme-util to CLI reference

* Apply suggestions from code review

* refactor: moved stuff around

* Apply suggestions from code review

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* build: added pylint command for code duplication checks to Makefile

* fix: more refactor, support summarizing conversations with `gptme-util chats list --summarize`

* test: fixed tests

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
  • Loading branch information
ErikBjare and ellipsis-dev[bot] authored Nov 17, 2024
1 parent e4ce6c8 commit 128676c
Show file tree
Hide file tree
Showing 13 changed files with 532 additions and 98 deletions.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ lint:
! grep -r 'ToolUse("python"' ${SRCDIRS}
@# ruff
poetry run ruff check ${RUFF_ARGS}

poetry run pylint --disable=all --enable=duplicate-code gptme/

format:
poetry run ruff check --fix-only ${RUFF_ARGS}
Expand Down
4 changes: 4 additions & 0 deletions docs/cli.rst
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,7 @@ This is the full CLI reference. For a more concise version, run ``gptme --help``
.. click:: gptme.eval:main
:prog: gptme-eval
:nested: full

.. click:: gptme.util.cli:main
:prog: gptme-util
:nested: full
2 changes: 1 addition & 1 deletion docs/server.rst
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ It can be started by running the following command:
gptme-server
For more CLI usage, see :ref:`the CLI documentation <cli:gptme-server>`.
For more CLI usage, see the :ref:`CLI reference <cli:gptme-server>`.

There are a few different interfaces available:

Expand Down
13 changes: 4 additions & 9 deletions gptme/llm_openai_models.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,10 @@
from typing import TypedDict
from typing_extensions import NotRequired
from typing import TYPE_CHECKING

if TYPE_CHECKING:
from .models import _ModelDictMeta # fmt: skip

class _ModelDictMeta(TypedDict):
context: int
max_output: NotRequired[int]
price_input: NotRequired[float]
price_output: NotRequired[float]


OPENAI_MODELS: dict[str, _ModelDictMeta] = {
OPENAI_MODELS: dict[str, "_ModelDictMeta"] = {
# GPT-4o
"gpt-4o": {
"context": 128_000,
Expand Down
30 changes: 30 additions & 0 deletions gptme/logmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -331,13 +331,26 @@ def _conversation_files() -> list[Path]:

@dataclass(frozen=True)
class ConversationMeta:
"""Metadata about a conversation."""

name: str
path: str
created: float
modified: float
messages: int
branches: int

def format(self, metadata=False) -> str:
"""Format conversation metadata for display."""
output = f"{self.name}"
if metadata:
output += f"\nMessages: {self.messages}"
output += f"\nCreated: {datetime.fromtimestamp(self.created)}"
output += f"\nModified: {datetime.fromtimestamp(self.modified)}"
if self.branches > 1:
output += f"\n({self.branches} branches)"
return output


def get_conversations() -> Generator[ConversationMeta, None, None]:
"""Returns all conversations, excluding ones used for testing, evals, etc."""
Expand Down Expand Up @@ -368,6 +381,23 @@ def get_user_conversations() -> Generator[ConversationMeta, None, None]:
yield conv


def list_conversations(
limit: int = 20,
include_test: bool = False,
) -> list[ConversationMeta]:
"""
List conversations with a limit.
Args:
limit: Maximum number of conversations to return
include_test: Whether to include test conversations
"""
conversation_iter = (
get_conversations() if include_test else get_user_conversations()
)
return list(islice(conversation_iter, limit))


def _gen_read_jsonl(path: PathLike) -> Generator[Message, None, None]:
with open(path) as file:
for line in file.readlines():
Expand Down
25 changes: 24 additions & 1 deletion gptme/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,30 @@ def to_xml(self) -> str:
attrs = f"role='{self.role}'"
return f"<message {attrs}>\n{self.content}\n</message>"

def format(self, oneline: bool = False, highlight: bool = False) -> str:
def format(
self,
oneline: bool = False,
highlight: bool = False,
max_length: int | None = None,
) -> str:
"""Format the message for display.
Args:
oneline: Whether to format the message as a single line
highlight: Whether to highlight code blocks
max_length: Maximum length of the message. If None, no truncation is applied.
If set, will truncate at first newline or max_length, whichever comes first.
"""
if max_length is not None:
first_newline = self.content.find("\n")
max_length = (
min(max_length, first_newline) if first_newline != -1 else max_length
)
content = self.content[:max_length]
if len(content) < len(self.content):
content += "..."
temp_msg = self.replace(content=content)
return format_msgs([temp_msg], oneline=True, highlight=highlight)[0]
return format_msgs([self], oneline=oneline, highlight=highlight)[0]

def print(self, oneline: bool = False, highlight: bool = True) -> None:
Expand Down
3 changes: 2 additions & 1 deletion gptme/prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
from .__version__ import __version__
from .config import get_config, get_project_config
from .message import Message
from .tools import loaded_tools
from .util import document_prompt_function

PromptType = Literal["full", "short"]
Expand Down Expand Up @@ -199,6 +198,8 @@ def prompt_project() -> Generator[Message, None, None]:

def prompt_tools(examples: bool = True) -> Generator[Message, None, None]:
"""Generate the tools overview prompt."""
from .tools import loaded_tools # fmt: skip

assert loaded_tools, "No tools loaded"
prompt = "# Tools Overview"
for tool in loaded_tools:
Expand Down
109 changes: 30 additions & 79 deletions gptme/tools/chats.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,32 +2,16 @@
List, search, and summarize past conversation logs.
"""

import itertools
import logging
import textwrap
from pathlib import Path
from textwrap import indent
from typing import TYPE_CHECKING

from ..message import Message
from .base import ToolSpec, ToolUse

if TYPE_CHECKING:
from ..logmanager import LogManager


logger = logging.getLogger(__name__)


def _format_message_snippet(msg: Message, max_length: int = 100) -> str:
"""Format a message snippet for display."""
first_newline = msg.content.find("\n")
max_length = min(max_length, first_newline) if first_newline != -1 else max_length
content = msg.content[:max_length]
return f"{msg.role.capitalize()}: {content}" + (
"..." if len(content) <= len(msg.content) else ""
)


def _get_matching_messages(log_manager, query: str, system=False) -> list[Message]:
"""Get messages matching the query."""
return [
Expand All @@ -38,36 +22,9 @@ def _get_matching_messages(log_manager, query: str, system=False) -> list[Messag
]


def _summarize_conversation(
log_manager: "LogManager", include_summary: bool
) -> list[str]:
"""Summarize a conversation."""
# noreorder
from ..llm import summarize as llm_summarize # fmt: skip

summary_lines = []
if include_summary:
summary = llm_summarize(log_manager.log.messages)
summary_lines.append(indent(f"Summary: {summary.content}", " "))
else:
non_system_messages = [msg for msg in log_manager.log if msg.role != "system"]
if non_system_messages:
first_msg = non_system_messages[0]
last_msg = non_system_messages[-1]

summary_lines.append(
f" First message: {_format_message_snippet(first_msg)}"
)
if last_msg != first_msg:
summary_lines.append(
f" Last message: {_format_message_snippet(last_msg)}"
)

summary_lines.append(f" Total messages: {len(log_manager.log)}")
return summary_lines


def list_chats(max_results: int = 5, include_summary: bool = False) -> None:
def list_chats(
max_results: int = 5, metadata=False, include_summary: bool = False
) -> None:
"""
List recent chat conversations and optionally summarize them using an LLM.
Expand All @@ -77,24 +34,30 @@ def list_chats(max_results: int = 5, include_summary: bool = False) -> None:
If True, uses an LLM to generate a comprehensive summary.
If False, uses a simple strategy showing snippets of the first and last messages.
"""
# noreorder
from ..logmanager import LogManager, get_user_conversations # fmt: skip
from ..llm import summarize # fmt: skip
from ..logmanager import LogManager, list_conversations # fmt: skip

conversations = list(itertools.islice(get_user_conversations(), max_results))
conversations = list_conversations(max_results)
if not conversations:
print("No conversations found.")
return

print(f"Recent conversations (showing up to {max_results}):")
for i, conv in enumerate(conversations, 1):
print(f"\n{i}. {conv.name}")
print(f" Created: {conv.created}")
if metadata:
print() # Add a newline between conversations
print(f"{i:2}. {textwrap.indent(conv.format(metadata=True), ' ')[4:]}")

log_path = Path(conv.path)
log_manager = LogManager.load(log_path)

summary_lines = _summarize_conversation(log_manager, include_summary)
print("\n".join(summary_lines))
# Use the LLM to generate a summary if requested
if include_summary:
summary = summarize(log_manager.log.messages)
print(
f"\n Summary:\n{textwrap.indent(summary.content, ' > ', predicate=lambda _: True)}"
)
print()


def search_chats(query: str, max_results: int = 5, system=False) -> None:
Expand All @@ -106,11 +69,10 @@ def search_chats(query: str, max_results: int = 5, system=False) -> None:
max_results (int): Maximum number of conversations to display.
system (bool): Whether to include system messages in the search.
"""
# noreorder
from ..logmanager import LogManager, get_user_conversations # fmt: skip
from ..logmanager import LogManager, list_conversations # fmt: skip

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

Expand All @@ -119,37 +81,31 @@ def search_chats(query: str, max_results: int = 5, system=False) -> None:
if matching_messages:
results.append(
{
"conversation": conv.name,
"conversation": conv,
"log_manager": log_manager,
"matching_messages": matching_messages,
}
)

if len(results) >= max_results:
break

# Sort results by the number of matching messages, in descending order
results.sort(key=lambda x: len(x["matching_messages"]), reverse=True)

if not results:
print(f"No results found for query: '{query}'")
return

# Sort results by the number of matching messages, in descending order
results.sort(key=lambda x: len(x["matching_messages"]), reverse=True)

print(f"Search results for query: '{query}'")
print(f"Found matches in {len(results)} conversation(s):")

for i, result in enumerate(results, 1):
print(f"\n{i}. Conversation: {result['conversation']}")
conversation = result["conversation"]
print(f"\n{i}. {conversation.format()}")
print(f" Number of matching messages: {len(result['matching_messages'])}")

summary_lines = _summarize_conversation(
result["log_manager"], include_summary=False
)
print("\n".join(summary_lines))

# Show sample matches
print(" Sample matches:")
for j, msg in enumerate(result["matching_messages"][:3], 1):
print(f" {j}. {_format_message_snippet(msg)}")
print(f" {j}. {msg.format(max_length=100)}")
if len(result["matching_messages"]) > 3:
print(
f" ... and {len(result['matching_messages']) - 3} more matching message(s)"
Expand All @@ -165,23 +121,18 @@ def read_chat(conversation: str, max_results: int = 5, incl_system=False) -> Non
max_results (int): Maximum number of messages to display.
incl_system (bool): Whether to include system messages.
"""
# noreorder
from ..logmanager import LogManager, get_conversations # fmt: skip

conversations = list(get_conversations())
from ..logmanager import LogManager, list_conversations # fmt: skip

for conv in conversations:
for conv in list_conversations():
if conv.name == conversation:
log_path = Path(conv.path)
logmanager = LogManager.load(log_path)
print(f"Reading conversation: {conversation}")
i = 0
for msg in logmanager.log:
if msg.role != "system" or incl_system:
print(f"{i}. {_format_message_snippet(msg)}")
print(f"{i}. {msg.format(max_length=100)}")
i += 1
else:
print(f"{i}. (system message)")
if i >= max_results:
break
break
Expand Down
10 changes: 7 additions & 3 deletions gptme/util.py → gptme/util/__init__.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
"""
Utility package for gptme.
"""

import functools
import io
import logging
Expand All @@ -17,7 +21,7 @@
from rich.console import Console
from rich.syntax import Syntax

from .clipboard import copy, set_copytext
from ..clipboard import copy, set_copytext

EMOJI_WARN = "⚠️"

Expand Down Expand Up @@ -319,8 +323,8 @@ def decorator(func): # pragma: no cover
return func

# noreorder
from .message import len_tokens # fmt: skip
from .tools import init_tools # fmt: skip
from ..message import len_tokens # fmt: skip
from ..tools import init_tools # fmt: skip

init_tools()

Expand Down
Loading

0 comments on commit 128676c

Please sign in to comment.