1 changed files with 63 additions and 0 deletions
@ -0,0 +1,63 @@
|
||||
#!/usr/bin/env python |
||||
|
||||
import os |
||||
import argparse |
||||
from dotenv import load_dotenv |
||||
from openai import OpenAI |
||||
|
||||
def load_openai_key(): |
||||
# Load environment variables in a file called .env |
||||
load_dotenv(override=True) |
||||
api_key = os.getenv('OPENAI_API_KEY') |
||||
|
||||
# Check the key |
||||
if not api_key: |
||||
return "Error: No API key was found!" |
||||
elif not api_key.startswith("sk-proj-"): |
||||
return "Error: An API key was found, but it doesn't start sk-proj-; please check you're using the right key" |
||||
elif api_key.strip() != api_key: |
||||
return "Error: An API key was found, but it looks like it might have space or tab characters at the start or end - please remove them!" |
||||
else: |
||||
return "API key found and looks good so far!" |
||||
|
||||
def ask_llm(client, model, user_prompt): |
||||
system_prompt = """ |
||||
you are a writing assistant with an expertise in children's stories. |
||||
Write a bedtime story inspired by the subject below. |
||||
The story should have a begining, middle, and end. |
||||
The story shoukd be appropriate for children ages 5-8 and have a positive message. |
||||
I should be able to read the entire story in about 3 minutes |
||||
""" |
||||
response = client.chat.completions.create( |
||||
model = model, |
||||
messages = [ {"role": "system", "content": system_prompt}, |
||||
{"role": "user", "content": user_prompt}] |
||||
) |
||||
return response.choices[0].message.content |
||||
|
||||
def main(): |
||||
parser = argparse.ArgumentParser(description="AI Bedtime Storyteller") |
||||
parser.add_argument("provider", choices=["openai", "ollama"], help="AI provider to use") |
||||
parser.add_argument("--model", help="Model to use for Ollama (required if provider is 'ollama')", required="ollama" in parser.parse_known_args()[0].provider) |
||||
parser.add_argument("subject", help="What do you want the story to be about?") |
||||
|
||||
args = parser.parse_args() |
||||
|
||||
if args.provider == "openai": |
||||
load_openai_key() |
||||
client = OpenAI() |
||||
model = "gpt-4o-mini" |
||||
elif args.provider == "ollama": |
||||
client = OpenAI(base_url='http://localhost:11434/v1', api_key='ollama') |
||||
model = args.model |
||||
else: |
||||
return "Error: invalid provider!" |
||||
|
||||
user_prompt = args.subject |
||||
|
||||
result = ask_llm(client, model, user_prompt) |
||||
print("AI Response:", result) |
||||
|
||||
if __name__ == "__main__": |
||||
main() |
||||
|
Loading…
Reference in new issue