anki-csv2ankicards/deck_creation.py

import openai
import os
import json
import genanki
from logging_config import setup_logging
from constants import API_KEY_ENV, CHAT_MODEL


setup_logging()


API_KEY = os.environ.get(API_KEY_ENV)
if not API_KEY:
    raise ValueError("Please set the OPENAI_API_KEY environment variable.")

openai.api_key = API_KEY

# Given prompt template
PROMPT_TEMPLATE = """
Please come up with a title for the deck and a set of 10 index cards for memorization, 
including a title, front, and back for each card. The index cards should completely 
capture the main points and themes of the text. In addition, they should contain any 
numbers or data that humans might find difficult to remember. The goal of the index 
card set is that one who memorizes it can provide a summary of the text to someone 
else, conveying the main points and themes.

You will provide the deck title, and the titles, questions, and answers for each card 
in a structured format as follows:
```
Deck Title: Title of the Deck
Cards:
- Title: Card Title 1
  Front: What is the capital of New York?
  Back: Albany
- Title: Card Title 2
  Front: Where in the world is Carmen San Diego?
  Back: Nobody knows
```

{content}
"""


def prompt_for_card_content(text_content):
    # Prepare the prompt
    prompt = PROMPT_TEMPLATE.format(content=text_content)

    # Get completion from the OpenAI ChatGPT API
    response = openai.ChatCompletion.create(
      model=CHAT_MODEL,
      messages=[
          {"role": "user", "content": prompt}
      ],
      temperature=0,
    )

    # Extract content from response and save to a new file
    return response.choices[0]['message']['content']


def response_to_json(response_text):
    lines = [line.strip() for line in response_text.split("\n") if line.strip()]

    deck_title = None
    cards = []
    current_card = {}

    for line in lines:
        if "Deck Title:" in line and not deck_title:
            deck_title = line.split("Deck Title:", 1)[1].strip()
        elif "Title:" in line:
            if current_card:  # If there's a card being processed, add it to cards
                cards.append(current_card)
                current_card = {}
            current_card["Title"] = line.split("Title:", 1)[1].strip()
        elif "Front:" in line:
            current_card["Question"] = line.split("Front:", 1)[1].strip()
        elif "Back:" in line:
            current_card["Answer"] = line.split("Back:", 1)[1].strip()

    if current_card:  # Add the last card if it exists
        cards.append(current_card)

    return {
        "DeckTitle": deck_title,
        "Cards": cards
    }


# Create a new model for our cards. This is necessary for genanki.
MY_MODEL = genanki.Model(
    1607372319,
    "Simple Model",
    fields=[
        {"name": "Title"},
        {"name": "Question"},
        {"name": "Answer"},
    ],
    templates=[
        {
            "name": "{{Title}}",
            "qfmt": "{{Question}}",
            "afmt": "{{FrontSide}}<hr id='answer'>{{Answer}}",
        },
    ])

def json_file_to_package(json_path):
    with open(json_path, 'r', encoding='utf-8') as f:
        json_data = json.load(f)
        package = to_package(json_data)

    return package

def to_package(deck_json):
    deck_title = deck_json["DeckTitle"]
    deck = genanki.Deck(1607372319, deck_title)

    for card_json in deck_json["Cards"]:
        title = card_json["Title"]
        question = card_json["Question"]
        answer = card_json["Answer"]

        note = genanki.Note(
                model=MY_MODEL,
                fields=[title, question, answer]
            )

        deck.add_note(note)

    return genanki.Package(deck)