diff --git a/requirements.txt b/requirements.txt index de677d7..74c0c46 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,9 @@ -openai~=0.27.0 +openai>=0.28.0,<1 +requests>=2.31.0,<3 markdown~=3.4.1 retrying~=1.3.4 setuptools~=65.5.0 pyfiglet~=0.8.post1 streamlit~=1.17.0 gtts~=2.3.1 -tqdm~=4.64.1 \ No newline at end of file +tqdm~=4.64.1 diff --git a/src/__pycache__/__init__.cpython-310.pyc b/src/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 7b646f7..0000000 Binary files a/src/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/src/__pycache__/__init__.cpython-39.pyc b/src/__pycache__/__init__.cpython-39.pyc deleted file mode 100644 index e4d55e7..0000000 Binary files a/src/__pycache__/__init__.cpython-39.pyc and /dev/null differ diff --git a/src/__pycache__/book.cpython-310.pyc b/src/__pycache__/book.cpython-310.pyc deleted file mode 100644 index c61f496..0000000 Binary files a/src/__pycache__/book.cpython-310.pyc and /dev/null differ diff --git a/src/__pycache__/prompts.cpython-310.pyc b/src/__pycache__/prompts.cpython-310.pyc deleted file mode 100644 index 1c5873a..0000000 Binary files a/src/__pycache__/prompts.cpython-310.pyc and /dev/null differ diff --git a/src/app.py b/src/app.py index bfcbea1..4ab7dc9 100644 --- a/src/app.py +++ b/src/app.py @@ -1,10 +1,19 @@ +from typing import Optional + import streamlit as st -import openai from book import Book from utils import * +from ollama_client import check_connection, OLLAMA_BASE_URL, OLLAMA_MODEL + +try: + import openai # type: ignore +except ImportError: # pragma: no cover - optional dependency + openai = None + +BACKEND_OPTIONS = ('Ollama', 'OpenAI') valid = False -content = '' +backend_choice = BACKEND_OPTIONS[0] # Center the title st.title('BookGPT') @@ -12,29 +21,74 @@ def initialize(): - global valid - # Get the API key and check if it is valid - api_key = st.text_input('OpenAI API Key', type='password') - if api_key: - openai.api_key = api_key - - # Check if the API key is valid - try: - openai.Engine.list() - valid = True - st.success('API key is valid!') + global valid, backend_choice + backend_choice = st.radio('LLM Backend', BACKEND_OPTIONS, index=0) - # API key is not valid - except openai.error.AuthenticationError: + if backend_choice == 'OpenAI': + if openai is None: + st.error('The openai package is not installed. Please install it to use the OpenAI backend.') valid = False - st.error('API key is not valid!') + return + + api_key = st.text_input('OpenAI API Key', type='password') + st.text_input('OpenAI Model', value='gpt-3.5-turbo', key='openai_model_input') + + if api_key: + openai.api_key = api_key + try: + openai.Model.list() + valid = True + st.success('API key is valid!') + except openai.error.AuthenticationError: # type: ignore[attr-defined] + valid = False + st.error('API key is not valid!') + except Exception: + valid = False + st.warning('Unable to validate the API key right now. Please try again later.') + else: + valid = False + else: + if check_connection(): + valid = True + st.success(f'Connected to {OLLAMA_MODEL} at {OLLAMA_BASE_URL}.') + else: + valid = False + st.error('Unable to connect to the Ollama server.') def generate_book(chapters, words, category, topic, language): - book = Book(chapters, words, topic, category, language) - - content = book.get_md() - st.markdown(content) + backend = backend_choice.lower() + kwargs = dict( + chapters=chapters, + words_per_chapter=words, + topic=topic, + category=category, + language=language, + llm_backend=backend, + ) + + if backend == 'openai': + kwargs['openai_model'] = st.session_state.get('openai_model_input', 'gpt-3.5-turbo') + + book: Optional[Book] = None + try: + with st.spinner('Generating book...'): + book = Book(**kwargs) + book.get_title() + book.get_structure() + book.finish_base() + book.get_content() + st.markdown(book.to_markdown()) + except Exception as exc: + st.error(f'Failed to generate the book: {exc}') + if book and hasattr(book, 'content'): + try: + st.info('Partial content generated before the error:') + st.markdown(book.to_markdown()) + except Exception: + pass + if book and getattr(book, 'last_saved_path', None): + st.caption(f"Partial book saved to {book.last_saved_path}.") def show_form(): @@ -42,10 +96,10 @@ def show_form(): with st.form('BookGPT'): # Get the number of chapters - chapters = st.number_input('How many chapters should the book have?', min_value=3, max_value=100, value=5) + chapters = st.number_input('How many chapters should the book have?', min_value=1, max_value=100, value=5) # Get the number of words per chapter - words = st.number_input('How many words should each chapter have?', min_value=100, max_value=3500, value=1000, + words = st.number_input('How many words should each chapter have?', min_value=100, max_value=2000, value=1200, step=50) # Get the category of the book @@ -61,9 +115,12 @@ def show_form(): # Submit button submit = st.form_submit_button('Generate') - # Check if the api key was valid + # Check if the selected backend is ready if submit and not valid: - st.error('The API key is not valid!') + if backend_choice == 'OpenAI': + st.error('The OpenAI configuration is not valid. Please check your API key and try again.') + else: + st.error('Unable to reach the Ollama server. Please try again later.') # Check if all fields are filled elif submit and not (chapters and words and category and topic and language): diff --git a/src/book.py b/src/book.py index 7e76670..76e5dd5 100644 --- a/src/book.py +++ b/src/book.py @@ -1,16 +1,58 @@ -import openai from tqdm import tqdm import prompts +import random +from datetime import datetime, timezone, timedelta +import time +from typing import List, Dict, Optional +class GenerationInterrupted(RuntimeError): + """Raised when generation stops mid-chapter but partial content exists.""" + + def __init__(self, message: str, partial_chapter: List[str]): + super().__init__(message) + self.partial_chapter = partial_chapter + +try: + import openai # type: ignore +except ImportError: # pragma: no cover - optional dependency + openai = None + +from ollama_client import chat, OllamaError + class Book: + def __str__(self): + book_structure = "Structure of the book:\n" + for chapter_index, chapter_info in enumerate(self.chapters, start=1): + chapter_title = chapter_info['title'] + chapter_paragraphs = chapter_info['paragraphs'] + book_structure += f"Chapter {chapter_index} ({len(chapter_paragraphs)} paragraphs): {chapter_title}\n" + for paragraph_index, paragraph_info in enumerate(chapter_paragraphs, start=1): + paragraph_title = paragraph_info['title'] + paragraph_words = paragraph_info['words'] + book_structure += f"\tParagraph {paragraph_index} ({paragraph_words} words): {paragraph_title}\n" + + return book_structure + def __init__(self, **kwargs): + excluded_keys = {'tolerance', 'llm_backend', 'openai_model', 'ollama_options'} # Joining the keyword arguments into a single string - self.arguments = '; '.join([f'{key}: {value}' for key, value in kwargs.items() if key != 'tolerance']) + self.arguments = '; '.join([ + f'{key}: {value}' for key, value in kwargs.items() if key not in excluded_keys + ]) # Get 'tolerance' attribute from kwargs self.tolerance = kwargs.get('tolerance', 0.9) + # Configure LLM preferences + self.llm_backend = kwargs.get('llm_backend', 'openai').lower() + self.openai_model = kwargs.get('openai_model', 'gpt-3.5-turbo') + self.ollama_options = kwargs.get('ollama_options') + + # Track whether only partial content is available + self.partial_content = False + self.last_saved_path: Optional[str] = None + # Assign a status variable self.status = 0 @@ -49,9 +91,14 @@ def get_structure(self): self.structure_prompt.append(self.get_message('user', structure_arguments)) self.structure = self.get_response(self.structure_prompt) self.chapters = self.convert_structure(self.structure) - self.paragraph_amounts = self.get_paragraph_amounts(self.chapters) - self.paragraph_words = self.get_paragraph_words(self.chapters) - return self.structure, self.chapters + + # Ensure self.chapters contains the actual chapter information before assigning paragraph amounts and words. + if isinstance(self.chapters, list): + self.paragraph_amounts = self.get_paragraph_amounts(self.chapters) # updated line + self.paragraph_words = self.get_paragraph_words(self.chapters) # updated line + return str(self.structure) + else: + self.output('Error in converting the book structure.') def finish_base(self): if not hasattr(self, 'title'): @@ -77,33 +124,59 @@ def calculate_max_status(self): return self.max_status def get_content(self): - chapters = [] - for i in tqdm(range(len(self.chapters))): - prompt = self.base_prompt.copy() - chapter = self.get_chapter(i, prompt.copy()) - chapters.append(chapter) - self.content = chapters - return self.content + if not hasattr(self, 'chapters'): + raise ValueError('Structure not generated yet.') - def save_book(self): + chapters: List[List[str]] = [] + try: + for i in tqdm(range(len(self.chapters))): + prompt = self.base_prompt.copy() + chapter = self.get_chapter(i, prompt.copy()) + chapters.append(chapter) + except GenerationInterrupted as interrupted: + chapters.append(interrupted.partial_chapter) + self._persist_partial_content(chapters, interrupted) + raise RuntimeError(str(interrupted)) from interrupted.__cause__ + except Exception: + if chapters: + self._persist_partial_content(chapters) + raise + else: + self.content = chapters + self.partial_content = False + return self.content + + def save_book(self, filename: Optional[str] = None) -> str: # Save the book in md format - with open(f'book.md', 'w') as file: - file.write(f'# {self.title}\n\n') - for chapter in self.content: - file.write(f'## {self.chapters[self.content.index(chapter)]["title"]}\n\n') - for paragraph in chapter: - file.write( - f'### {self.chapters[self.content.index(chapter)]["paragraphs"][chapter.index(paragraph)]["title"]}\n\n') - file.write(paragraph + '\n\n') - file.write('\n\n') + # Corrected saving the book with the specified time + desired_time = datetime.now(timezone(timedelta(hours=-5))) # EST timezone + # Use the desired time as a seed for the random number generator + random.seed(desired_time) + # Generate a random 4-digit number + random_number = random.randint(1000009, 9999999) + # Ensure it's 4 digits long + random_number = str(random_number).zfill(random.randint(7, 10)) + suffix = '_partial' if self.partial_content else '' + path = filename or f'book{random_number}{suffix}.md' + with open(path, 'w') as file: + file.write(self.to_markdown()) + self.last_saved_path = path + return path def get_chapter(self, chapter_index, prompt): - if len(self.base_prompt) == 3: + if len(self.base_prompt) <= 9: self.finish_base() paragraphs = [] for i in range(self.paragraph_amounts[chapter_index]): - paragraph = self.get_paragraph(prompt.copy(), chapter_index, i) + try: + paragraph = self.get_paragraph(prompt.copy(), chapter_index, i) + except Exception as exc: # pragma: no cover - network/runtime failure + raise GenerationInterrupted( + f'Failed to generate paragraph {i + 1} of chapter {chapter_index + 1}', + paragraphs, + ) from exc + prompt.append(self.get_message('user', f'!w {chapter_index + 1} {i + 1}')) prompt.append(self.get_message('assistant', paragraph)) self.status += 1 @@ -126,29 +199,34 @@ def get_paragraph(self, prompt, chapter_index, paragraph_index): @staticmethod def get_message(role, content): return {"role": role, "content": content} - + @staticmethod def convert_structure(structure): chapters = structure.split("Chapter") chapters = [x for x in chapters if x != ''] chapter_information = [] - for chapter in chapters: - for line in chapter.split("\n"): - if 'paragraphs' in line.lower(): - chapter_information.append({'title': line.split('): ')[1], 'paragraphs': []}) - elif 'paragraph' in line.lower(): - chapter_information[-1]['paragraphs'].append( - {'title': line.split('): ')[1], 'words': line.split('(')[1].split(')')[0].split(' ')[0]}) - chapter_information[-1]['paragraph_amount'] = len(chapter_information[-1]['paragraphs']) - + chapter_lines = chapter.split("\n") + if len(chapter_lines) > 1: + chapter_title_line = chapter_lines[0] + if 'paragraphs' in chapter_title_line.lower(): + chapter_info = {'title': chapter_title_line.split('): ')[1], 'paragraphs': []} + for line in chapter_lines[1:]: + if 'paragraph' in line.lower(): + words_info = line.split('(')[1].split(')')[0].split(' ') + if len(words_info) >= 2: + paragraph_title = line.split('): ')[1] + paragraph_words = words_info[0] + chapter_info['paragraphs'].append({'title': paragraph_title, 'words': paragraph_words}) + chapter_information.append(chapter_info) return chapter_information + @staticmethod def get_paragraph_amounts(structure): amounts = [] for chapter in structure: - amounts.append(chapter['paragraph_amount']) + amounts.append(len(chapter['paragraphs'])) return amounts @staticmethod @@ -158,13 +236,73 @@ def get_paragraph_words(structure): words.append([int(x['words']) for x in chapter['paragraphs']]) return words - @staticmethod - def get_response(prompt): - return openai.ChatCompletion.create( - model="gpt-3.5-turbo-0301", - messages=prompt - )["choices"][0]["message"]["content"] + def get_response(self, prompt: List[Dict[str, str]], max_retries: int = 5) -> str: + retries = 0 + last_error: Optional[Exception] = None + backend = self.llm_backend + while retries < max_retries: + try: + if backend == 'ollama': + response = chat(prompt, options=self.ollama_options) + elif backend == 'openai': + if openai is None: # pragma: no cover - dependency guard + raise RuntimeError('openai package is not installed') + response = openai.ChatCompletion.create( # type: ignore[attr-defined] + model=self.openai_model, + messages=prompt + )["choices"][0]["message"]["content"] + else: + raise RuntimeError(f"Unsupported LLM backend: {backend}") + + with open("log.txt", "a") as f: + f.write(f"Prompt: {prompt}\nResponse: {response}\n\n") + return response + except OllamaError as exc: + last_error = exc + except Exception as exc: # pragma: no cover - runtime/network failure + last_error = exc + + retries += 1 + print(f"An error occurred: {last_error}. Retrying ({retries}/{max_retries})...") + time.sleep(20) + + if last_error is None: + raise RuntimeError("Unknown error while requesting a response") + raise RuntimeError(f"Failed to get a response after {max_retries} retries.") from last_error + + def to_markdown(self) -> str: + if not hasattr(self, 'content'): + raise ValueError('Content not generated yet.') + + lines = [f'# {getattr(self, "title", "Untitled Book")}'] + if self.partial_content: + lines.append('') + lines.append('> **Note:** Generation stopped early; the content below is partial.') + + for chapter_index, (chapter_meta, paragraphs) in enumerate(zip(self.chapters, self.content), start=1): + lines.append('') + lines.append(f'## Chapter {chapter_index}: {chapter_meta["title"]}') + lines.append('') + for paragraph_index, paragraph in enumerate(paragraphs, start=1): + paragraph_meta = chapter_meta['paragraphs'][paragraph_index - 1] + lines.append(f'### {paragraph_meta["title"]}') + lines.append('') + lines.append(paragraph) + lines.append('') + + return '\n'.join(lines).strip() + '\n' @staticmethod def output(message): print(message) + + def _persist_partial_content(self, chapters: List[List[str]], error: Optional[Exception] = None) -> None: + """Persist partial content to disk for recovery.""" + + self.content = chapters + self.partial_content = True + path = self.save_book() + message = f'Partial book saved to {path}.' + if error is not None: + message = f'{message} Reason: {error}' + self.output(message) diff --git a/src/ollama_client.py b/src/ollama_client.py new file mode 100644 index 0000000..b4bc54a --- /dev/null +++ b/src/ollama_client.py @@ -0,0 +1,90 @@ +"""Lightweight helper for interacting with a remote Ollama instance.""" + +from __future__ import annotations + +import os +from typing import Dict, Iterable, Optional + +import requests + + +class OllamaError(RuntimeError): + """Raised when the Ollama backend returns an error.""" + + +_DEFAULT_HOST = "69.142.141.135" +_DEFAULT_PORT = "11434" +_DEFAULT_MODEL = "gpt-oss:120b-cloud" + + +OLLAMA_HOST = os.getenv("OLLAMA_HOST", _DEFAULT_HOST) +OLLAMA_PORT = os.getenv("OLLAMA_PORT", _DEFAULT_PORT) +OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", _DEFAULT_MODEL) +OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", f"http://{OLLAMA_HOST}:{OLLAMA_PORT}") + + +def _chat_endpoint() -> str: + return f"{OLLAMA_BASE_URL.rstrip('/')}/api/chat" + + +def _tags_endpoint() -> str: + return f"{OLLAMA_BASE_URL.rstrip('/')}/api/tags" + + +def chat( + messages: Iterable[Dict[str, str]], + *, + timeout: Optional[float] = 120, + options: Optional[Dict[str, object]] = None, +) -> str: + """Send a chat completion request to the configured Ollama backend.""" + + payload: Dict[str, object] = { + "model": OLLAMA_MODEL, + "messages": list(messages), + "stream": False, + } + + if options: + payload["options"] = options + + try: + response = requests.post(_chat_endpoint(), json=payload, timeout=timeout) + response.raise_for_status() + except requests.RequestException as exc: # pragma: no cover - network failure + raise OllamaError("Failed to reach the Ollama server") from exc + + try: + data = response.json() + except ValueError as exc: # pragma: no cover - unexpected payload + raise OllamaError("Ollama returned a non-JSON response") from exc + + message = data.get("message") + if isinstance(message, dict): + content = message.get("content") + if isinstance(content, str): + return content + + if "response" in data and isinstance(data["response"], str): + return data["response"] + + raise OllamaError("Unexpected response format from Ollama") + + +def check_connection(timeout: float = 5.0) -> bool: + """Return True when the configured model is reachable on the Ollama server.""" + + try: + response = requests.get(_tags_endpoint(), timeout=timeout) + response.raise_for_status() + body = response.json() + except (requests.RequestException, ValueError): # pragma: no cover - network failure + return False + + models = body.get("models") + if isinstance(models, list): + for entry in models: + if isinstance(entry, dict) and entry.get("name") == OLLAMA_MODEL: + return True + + return False diff --git a/src/prompts.py b/src/prompts.py index 75a7a78..a6693cb 100644 --- a/src/prompts.py +++ b/src/prompts.py @@ -8,4 +8,4 @@ TITLE_INSTRUCTIONS = "You are a title creating AI for books. Create the title based on what title types are used by the best (selling) books. The title is one of the most important things, so it should be really good. It should be in the format: \"title\". The user will give you the topic and other data after you are ready. Type \"Ready\", if you are ready." -STRUCTURE_INSTRUCTIONS = "You are a book structure creating AI. You are using other books of the same type as structure inspiration or you are creating you one structure, if appropriate. The structure should look like the following example: \"Chapter 1 ({the amount of paragraphs}): xxx\n\tParagraph 1 ({amount of recommended words} words): xxx\n\tParagraph 2 ({amount of recommended words} words): xxx\nChapter 2 ({the amount of paragraphs} paragraphs): xxx\n\tParagraph 1 ({amount of recommended words} words): xxx\n\tParagraph 2 ({amount of recommended words} words): xxx\n... Find fitting titles for the chapters and paragraphs, that will give the writer a lot of text to write (Also find a good amount of paragraphs, depending on the words per chapter amount, better more than less). Follow the format, don't write any additional information and make sure, that the words per paragraph add up to the words per chapter amount. The user will give you the topic and other data after you are ready. Type \"Ready\", if you are ready." +STRUCTURE_INSTRUCTIONS = "You are a book structure creating AI. You are using other books of the same type as structure inspiration or you are creating your own structure, if appropriate. The structure should look like the following example: \"Chapter 1 ({the amount of paragraphs}): xxx\n\tParagraph 1 ({amount of recommended words} words): xxx\n\tParagraph 2 ({amount of recommended words} words): xxx\nChapter 2 ({the amount of paragraphs} paragraphs): xxx\n\tParagraph 1 ({amount of recommended words} words): xxx\n\tParagraph 2 ({amount of recommended words} words): xxx\n... Find fitting titles for the chapters and paragraphs, that will give the writer a lot of text to write (Also find a good amount of paragraphs, depending on the words per chapter amount, better more than less). Follow the format, don't write any additional information and make sure, that the words per paragraph add up to the words per chapter amount. The user will give you the topic and other data after you are ready. Type \"Ready\", if you are ready." diff --git a/src/run.py b/src/run.py index 82e08eb..9085616 100644 --- a/src/run.py +++ b/src/run.py @@ -1,17 +1,12 @@ # Imports +import os from pyfiglet import Figlet from book import Book -import json -import openai - - -# Get the OpenAI API key from the config file -def get_api_key(): - # Read the config file - with open('config.json', 'r') as f: - # Return the OpenAI key - return json.load(f)['OpenAI_key'] +try: + import openai # type: ignore +except ImportError: # pragma: no cover - optional dependency + openai = None # Draw the given text in a figlet def draw(text): @@ -48,98 +43,146 @@ def get_option(options): print('Invalid option. Please try again.') -# Main function +def select_backend() -> str: + backend = os.getenv('BOOKGPT_BACKEND', 'ollama').lower() + if backend not in {'openai', 'ollama'}: + backend = 'ollama' + + if backend == 'openai': + if openai is None: + print('openai package is not installed. Falling back to Ollama backend.') + return 'ollama' + + api_key = os.getenv('OPENAI_KEY') + if not api_key: + print('OPENAI_KEY not found. Falling back to Ollama backend.') + return 'ollama' + + openai.api_key = api_key + print('Using OpenAI backend for generation.') + else: + print('Using Ollama backend for generation.') + + return backend + + +def get_default_book_kwargs(backend: str) -> dict: + data = { + 'chapters': int(os.getenv('BOOKGPT_CHAPTERS', 5)), + 'words_per_chapter': int(os.getenv('BOOKGPT_WORDS_PER_CHAPTER', 1200)), + 'category': os.getenv('BOOKGPT_CATEGORY', 'Science Fiction'), + 'topic': os.getenv('BOOKGPT_TOPIC', """William a 27-year-old boy moves to an unfamiliar city and rents a house, where he will begin a new life. He is quiet, socially awkward, and dislikes interacting with people, Nevertheless, he will inevitably encounter various situations requiring social interaction in the future, as well as many moments where friends will be needed, whether for problem-solving or emotional support. Despite his quirky personality, this also makes it easier for him to find genuine friends. These friends, while tolerating his rationality and sharpness, care about him and try their best to help him resolve the problems he encounters. His life is simple. He lives frugally, spending only the necessary money on essential daily necessities. When it comes to interpersonal relationships, he highly values "choice" and "necessity." He believes that no friend or person has any obligation to do anything for him, and he himself has no justification to demand that anyone must do anything for him. If someone tells the protagonist, "You are very important to me," he would feel flustered and overwhelmed. He takes commitments seriously and will always do his utmost to fulfill promises he has made. However, he is usually cautious and tends to avoid making promises altogether. 未自华为备意录"""), + 'tolerance': float(os.getenv('BOOKGPT_TOLERANCE', 0.6)), + 'llm_backend': backend, + } + + if backend == 'openai': + data['openai_model'] = os.getenv('OPENAI_MODEL', 'gpt-3.5-turbo') + + return data + + +def prompt_with_default(prompt_text: str, default: str) -> str: + response = input(f"{prompt_text} [{default}]: ").strip() + return response or default + + +def prompt_int_with_default(prompt_text: str, default: int, minimum: int = 1) -> int: + while True: + response = input(f"{prompt_text} [{default}]: ").strip() + if not response: + return default + try: + value = int(response) + except ValueError: + print('Please enter a valid number.') + continue + + if value < minimum: + print(f'Please enter a value greater than or equal to {minimum}.') + continue + + return value + + +def collect_book_preferences(defaults: dict) -> dict: + print('Provide details for your book (press Enter to keep the default).') + chapters = prompt_int_with_default('How many chapters should the book have?', defaults['chapters'], minimum=1) + words = prompt_int_with_default( + 'How many words should each chapter have?', + defaults['words_per_chapter'], + minimum=100, + ) + category = prompt_with_default('What is the category of the book?', defaults['category']) + topic = prompt_with_default('What is the topic of the book?', defaults['topic']) + + data = { + 'chapters': chapters, + 'words_per_chapter': words, + 'category': category, + 'topic': topic, + 'tolerance': defaults['tolerance'], + 'llm_backend': defaults['llm_backend'], + } + + if defaults.get('openai_model'): + data['openai_model'] = defaults['openai_model'] + + return data + + def main(): - # Set the OpenAI API key - openai.api_key = get_api_key() + backend = select_backend() # Draw the title draw('BookGPT') - # Check if the user wants to generate a new book or not if get_option(['Generate a book', 'Exit']) - 1: return - # Get the number of chapters - print('How many chapters should the book have?') - chapters = int(input('> ')) - - # Get the number of words per chapter - print('How many words should each chapter have?') - # Check if it is below 1200 - words = int(input('> ')) - if words <= 1200: - words = 1200 - print('The number of words per chapter has been set to 1200. (The max number of words per chapter)') - - # Get the category of the book - print('What is the category of the book?') - category = input('> ') - - # Get the topic of the book - print('What is the topic of the book?') - topic = input('> ') - - # What is the tolerance of the book? - print('What is the tolerance of the book? (0.8 means that 80% of the words will be written 100%)') - tolerance = float(input('> ')) - if tolerance < 0 or tolerance > 1: - tolerance = 0.8 - - # Do you want to add any additional parameters? - print('Do you want to add any additional parameters?') - if get_option(['No', 'Yes']) - 1: - print( - 'Please enter the additional parameters in the following format: "parameter1=value1, parameter2=value2, ..."') - additional_parameters = input('> ') - additional_parameters = additional_parameters.split(', ') - for i in range(len(additional_parameters)): - additional_parameters[i] = additional_parameters[i].split('=') - additional_parameters = dict(additional_parameters) - else: - additional_parameters = {} - - # Initialize the Book - book = Book(chapters=chapters, words_per_chapter=words, topic=topic, category=category, tolerance=tolerance, - **additional_parameters) + defaults = get_default_book_kwargs(backend) + book_kwargs = collect_book_preferences(defaults) + book = Book(**book_kwargs) - # Print the title - print(f'Title: {book.get_title()}') + title = book.get_title() + print(f'Title: {title}') - # Ask if he wants to change the title until he is satisfied while True: print('Do you want to generate a new title?') if get_option(['No', 'Yes']) - 1: - print(f'Title: {book.get_title()}') + title = book.get_title() + print(f'Title: {title}') else: break - # Print the structure of the book print('Structure of the book:') - structure, _ = book.get_structure() + structure = book.get_structure() print(structure) - # Ask if he wants to change the structure until he is satisfied while True: print('Do you want to generate a new structure?') if get_option(['No', 'Yes']) - 1: print('Structure of the book:') - structure, _ = book.get_structure() + structure = book.get_structure() print(structure) else: break print('Generating book...') - # Initialize the book generation book.finish_base() + try: + book.get_content() + except Exception as exc: + if getattr(book, 'last_saved_path', None): + print(f'Generation interrupted: {exc}') + print(f'Partial book saved to {book.last_saved_path}.') + else: + print(f'Failed to generate the book: {exc}') + return - content = book.get_content() - - # Save the book - book.save_book() - print('Book saved as book.md.') - + path = book.save_book() + print(f'Book saved to {path}.') # Run the main function if __name__ == "__main__": diff --git a/src/utils/__pycache__/__init__.cpython-310.pyc b/src/utils/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 35f9306..0000000 Binary files a/src/utils/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/src/utils/__pycache__/utils.cpython-310.pyc b/src/utils/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index c322838..0000000 Binary files a/src/utils/__pycache__/utils.cpython-310.pyc and /dev/null differ