groq

#!/usr/bin/env python

import os
import sys

from groq import Groq

def main():
    client = Groq(
        api_key=os.environ.get("GROQ_API_KEY"),
    )

    model="mixtral-8x7b-32768"

    input = None

    #get input from args or stdin if no args
    if len(sys.argv) > 1:
        if sys.argv[1] == "-h" or sys.argv[1] == "--help":
            print("""Usage: grok [input]
            -l: use llama2-70b-4096 model
            -m: use mixtral-8x7b-32768 model""")
            sys.exit(0)
        for arg in sys.argv[1:]:
            if arg == "-l":
                model = "llama2-70b-4096"
                continue
            if arg == "-m":
                model = "mixtral-8x7b-32768"
                continue
            input = arg
    else:
        input = sys.stdin.read()

    if input is None:
        input = sys.stdin.read()

    if input is None or input == "":
        print("""Usage: grok [input]
        -l: use llama2-70b-4096 model
        -m: use mixtral-8x7b-32768 model""")
        sys.exit(1)

    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": f"{input}",
            }
        ],
        model=model,
    )
    print(chat_completion.choices[0].message.content)

if __name__ == "__main__":
    main()