Semantic Grep

gistfile1.txt

#!/usr/bin/env python3

# This "semantic grep" matches lines by ChatGPT

# Provide an OpenAI key via environment variable
# Provide a question in plain language as first parameter
# Pipe from cat or whatever
# Optional; provide examples and counter-examples

# cat input.txt | sgrep.py "Is this an animal name?" --example "fox" --counter-example "Rasputin"

import os
import sys
import json
import urllib.request
from urllib.error import URLError, HTTPError

# Function to send requests to OpenAI Chat API
def query_openai(prompt):
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("Error: OPENAI_API_KEY environment variable not set.", file=sys.stderr)
        sys.exit(1)

    data = json.dumps({
        "model": "gpt-4-turbo-preview",  # You can adjust the model
        "messages": prompt,
        "temperature": 0.0,
        "max_tokens": 4
    }).encode("utf-8")

    req = urllib.request.Request("https://api.openai.com/v1/chat/completions",
                                 data=data,
                                 headers={"Content-Type": "application/json",
                                          "Authorization": f"Bearer {api_key}"})
    try:
        with urllib.request.urlopen(req) as response:
            response_body = response.read()
            return json.loads(response_body)
    except HTTPError as e:
        print(f"HTTP Error: {e.code} {e.reason}", file=sys.stderr)
    except URLError as e:
        print(f"URL Error: {e.reason}", file=sys.stderr)

    sys.exit(1)

# Parse command line arguments
def parse_args():
    args = {"examples": []}
    argv = sys.argv[1:]
    while argv:
        arg = argv.pop(0)
        if arg == "--example" or arg == "--counter-example":
            if argv:
                value = argv.pop(0)
                args["examples"].append((arg, value))
            else:
                print(f"Argument {arg} requires a value", file=sys.stderr)
                sys.exit(1)
        else:
            if "sentence" in args:
                print(f"Unexpected argument {arg}", file=sys.stderr)
                sys.exit(1)
            args["sentence"] = arg
    return args

def main():
    args = parse_args()
    sentence = args.get("sentence")
    if not sentence:
        print("Usage: sgrep.py 'sentence' [--example 'positive example'] [--counter-example 'counter example']", file=sys.stderr)
        sys.exit(1)

    setup_messages = [
        {"role": "system", "content": "Write a boolean, that is, either `true` or `false`, showing whether the given line meets this definition: " + sentence},
    ]

    for arg, value in args["examples"]:
        setup_messages.append({"role": "user", "content": value})
        response = "true" if arg == "--example" else "false"
        setup_messages.append({"role": "assistant", "content": response})

    for line in sys.stdin:
        line = line.strip()
        messages = setup_messages.copy()
        messages.append({"role": "user", "content": line})
        result = query_openai(messages)
        if result:
            try:
                normalized = result["choices"][0]["message"]["content"].strip().lower()
                if normalized == "true":
                    print(line)
                elif normalized == "false":
                    continue
                else:
                    print(f"Error: unexpected response for line: {line}", file=sys.stderr)
            except KeyError:
                continue

if __name__ == "__main__":
    main()