copy_edit.py

#!/usr/bin/env python3
"""
Requirements:
pip install click langchain openai

Environment Variables:
OPENAI_API_KEY=...
"""
import sys
import click
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
class CustomHandler(StreamingStdOutCallbackHandler):
    def on_llm_start(self, serialized, prompts, **_) -> None:
        click.echo("\033[33m", nl=False)
        pass

    def on_llm_new_token(self, token: str, **_) -> None:
        click.echo(token, nl=False)

    def on_llm_end(self, response, **kwargs) -> None:
        click.echo("\033[0m")

from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage


def run(data, model):
    llm=ChatOpenAI(
            client=None,
            model=model,
            temperature=0.1,
            verbose=True,
            callbacks=[CustomHandler()],
            streaming=True,
        )
    messages = [
        SystemMessage(content="""You are Copy Editor 3000, an expert copy editing software tool.
You do your best to preserve the style and tone of the original texts.
Fix all spelling, grammar, and punctuation mistakes.
If needed add connectors to make the text flow better.
Ask questions something is unclear."""
        ),
        HumanMessage(content=f"original:\n{data}\n\n"),
    ]
    llm(messages)


@click.command()
@click.option('--model', '-m', default='gpt-3.5-turbo-16k-0613')
@click.argument('src',
              type=click.File('r'),
              default=sys.stdin)
def my_command(model, src):
    data = None
    with src:
        for line in src:
            if line.strip() == "":
                continue
            run(line, model)

if __name__ == '__main__':
    my_command()

echo 'I like to fead kats.' | ./copy_edit.py
I like to feed cats.