Fixed errors

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
+tqdm~=4.64.1

src/app.py

@@ -1,51 +1,105 @@
+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')
 st.markdown('---')
 
 
 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():
     # Create form for user input
     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):