Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor filesystem with pathlib + improve CLI UX #19

Closed
wants to merge 3 commits into from
Closed

Refactor filesystem with pathlib + improve CLI UX #19

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
08be4aa
Refactor filesystem with pathlib
sderev Jun 15, 2023
9cb5224
Improve CLI UX for no-arg scenarios
sderev Jun 15, 2023
974f9d2
Add a helper message for the `logs` command
sderev Jun 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 27 additions & 24 deletions llm/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import json
import openai
import os
from pathlib import Path
import sqlite_utils
import sys
import warnings
Expand All @@ -17,6 +18,8 @@

DEFAULT_MODEL = "gpt-3.5-turbo"

LOG_DB = Path().home() / ".llm" / "log.db"


@click.group(
cls=DefaultGroup,
Expand Down Expand Up @@ -53,8 +56,15 @@ def cli():
def openai_(prompt, system, gpt4, model, stream, no_log, code, _continue, chat_id):
"Execute a prompt against on OpenAI model"
if prompt is None:
if sys.stdin.isatty():
# No data being piped in and no arguments provided,
# so show help message and exit.
click.echo(cli.get_help(click.get_current_context()))
return

# Read from stdin instead
prompt = sys.stdin.read()

openai.api_key = get_openai_api_key()
if gpt4:
model = "gpt-4"
Expand Down Expand Up @@ -114,13 +124,12 @@ def openai_(prompt, system, gpt4, model, stream, no_log, code, _continue, chat_i

@cli.command()
def init_db():
"Ensure ~/.llm/log.db SQLite database exists"
path = get_log_db_path()
if os.path.exists(path):
"Ensure the log SQLite database exists"
if LOG_DB.exists():
return
# Ensure directory exists
os.makedirs(os.path.dirname(path), exist_ok=True)
db = sqlite_utils.Database(path)
LOG_DB.parent.mkdir(parents=True, exist_ok=True)
db = sqlite_utils.Database(LOG_DB)
db.vacuum()


Expand All @@ -139,10 +148,10 @@ def init_db():
)
@click.option("-t", "--truncate", is_flag=True, help="Truncate long strings in output")
def logs(count, path, truncate):
path = path or get_log_db_path()
if not os.path.exists(path):
raise click.ClickException("No log database found at {}".format(path))
db = sqlite_utils.Database(path)
"Show log entries"
if not LOG_DB.exists():
raise click.ClickException(f"No log database found at {LOG_DB}. Run `llm init-db` to create {LOG_DB}")
db = sqlite_utils.Database(LOG_DB)
rows = list(
db["log"].rows_where(order_by="-rowid", select="rowid, *", limit=count or None)
)
Expand All @@ -163,27 +172,22 @@ def get_openai_api_key():
# Expand this to home directory / ~.openai-api-key.txt
if "OPENAI_API_KEY" in os.environ:
return os.environ["OPENAI_API_KEY"]
path = os.path.expanduser("~/.openai-api-key.txt")
API_key_path = Path.home() / ".openai-api-key.txt"
# If the file exists, read it
if os.path.exists(path):
with open(path) as fp:
return fp.read().strip()
if API_key_path.exists():
with open(API_key_path) as file:
return file.read().strip()
raise click.ClickException(
"No OpenAI API key found. Set OPENAI_API_KEY environment variable or create ~/.openai-api-key.txt"
)


def get_log_db_path():
return os.path.expanduser("~/.llm/log.db")


def log(no_log, provider, system, prompt, response, model, chat_id=None):
if no_log:
return
log_path = get_log_db_path()
if not os.path.exists(log_path):
if not LOG_DB.exists():
return
db = sqlite_utils.Database(log_path)
db = sqlite_utils.Database(LOG_DB)
db["log"].insert(
{
"provider": provider,
Expand All @@ -201,12 +205,11 @@ def log(no_log, provider, system, prompt, response, model, chat_id=None):
def get_history(chat_id):
if chat_id is None:
return None, []
log_path = get_log_db_path()
if not os.path.exists(log_path):
if not LOG_DB.exists():
raise click.ClickException(
"This feature requires logging. Run `llm init-db` to create ~/.llm/log.db"
f"This feature requires logging. Run `llm init-db` to create {LOG_DB}"
)
db = sqlite_utils.Database(log_path)
db = sqlite_utils.Database(LOG_DB)
# Check if the chat_id column exists in the DB. If not create it. This is a
# migration path for people who have been using llm before chat_id was
# added.
Expand Down