From the uDemy course on LLM engineering.
https://www.udemy.com/course/llm-engineering-master-ai-and-large-language-models
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
63 lines
2.3 KiB
63 lines
2.3 KiB
#!/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() |
|
|
|
|