LLM powered wordle solver

wordle.py

import os
import sqlite3
import google.generativeai as genai

# from: https://gist.github.com/slushman/34e60d6bc479ac8fc698df8c226e4264
from wordle.five_letter_words import words


GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GEMINI_MODEL = "gemini-1.5-flash"
WORDS_DB_PATH = "wordle/wordle.db"

CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS words (
    word TEXT PRIMARY KEY,
    letter1 CHAR(1) NOT NULL,
    letter2 CHAR(1) NOT NULL, 
    letter3 CHAR(1) NOT NULL,
    letter4 CHAR(1) NOT NULL,
    letter5 CHAR(1) NOT NULL
)
"""


def search(description: str):
    sql = generate_sql(description)
    return get_sql_results(sql)


def get_sql_results(sql: str):
    conn = sqlite3.connect(WORDS_DB_PATH)
    cursor = conn.cursor()
    cursor.execute(sql)
    results = cursor.fetchall()
    conn.close()
    for result in results:
        print("".join(result))


def generate_sql(prompt: str):
    genai.configure(api_key=GEMINI_API_KEY)
    model = genai.GenerativeModel(GEMINI_MODEL)
    prompt = f"You are an assistant to convert natural language in sqlite queries. I have a table with the following columns: {CREATE_TABLE_SQL}. I want to find words that match the following description: '{prompt}'. Please return an sql query that I can put into sqlite3.connect('wordle.db').cursor().execute(sql), i.e. the sql query should be a string without a leading or trailing newline and without `sql` in the string. Use the `word` column to find the words if that's a better way to do it than using the other columns."
    response = model.generate_content(prompt)
    return response.text


def create_db():
    conn = sqlite3.connect(WORDS_DB_PATH)
    cursor = conn.cursor()

    cursor.execute("DROP TABLE IF EXISTS words")

    # Create table with columns for each letter and add indices
    cursor.execute(CREATE_TABLE_SQL)

    # Create indices for each column
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_word ON words(word)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_letter1 ON words(letter1)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_letter2 ON words(letter2)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_letter3 ON words(letter3)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_letter4 ON words(letter4)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_letter5 ON words(letter5)")

    # Insert words from the imported list
    for word in words:
        cursor.execute(
            "INSERT INTO words VALUES (?, ?, ?, ?, ?, ?)",
            (word, word[0], word[1], word[2], word[3], word[4]),
        )

    # Commit changes and close connection
    conn.commit()
    conn.close()