refactor: Improve console logs

raestrada · Oct 14, 2024 · 7c06b58 · 7c06b58
1 parent 4db7c6f
commit 7c06b58
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 31 deletions.
diff --git a/storycraftr/agents.py b/storycraftr/agents.py
@@ -13,34 +13,38 @@
 # Function to load all Markdown files from the book's directory and subdirectories
 def load_markdown_files(book_name):
     """Load all Markdown files from the book's directory and subdirectories."""
+    console.print(f"[bold blue]Loading all Markdown files from '{book_name}'...[/bold blue]")  # Progress message
     md_files = glob.glob(f'{book_name}/**/*.md', recursive=True)
+    console.print(f"[bold green]Loaded {len(md_files)} Markdown files.[/bold green]")  # Success message
     return md_files
 
 # Function to delete an existing assistant
 def delete_assistant(name):
+    console.print(f"[bold blue]Checking if assistant '{name}' exists for deletion...[/bold blue]")  # Progress message
     assistants = client.beta.assistants.list()
     for assistant in assistants.data:
         if assistant.name == name:
             console.print(f"Deleting assistant {name}...")
             client.beta.assistants.delete(assistant_id=assistant.id)
-            console.print(f"Assistant {name} deleted.")
+            console.print(f"[bold green]Assistant {name} deleted successfully.[/bold green]")  # Success message
             break
 
 # Function to create or get an assistant
 def create_or_get_assistant(name, book_name):
+    console.print(f"[bold blue]Searching for existing assistant '{name}'...[/bold blue]")  # Progress message
     assistant = None
     assistants = client.beta.assistants.list()
     for assistant in assistants.data:
         if assistant.name == name:
-            console.print(f"Assistant {name} already exists.")
+            console.print(f"[bold yellow]Assistant {name} already exists.[/bold yellow]")  # Success message
             return assistant
 
     # Step 1. Create a vector store for the book
-    console.print(f"Creating vector store for '{book_name}'...")
+    console.print(f"[bold blue]Creating vector store for '{book_name}'...[/bold blue]")  # Progress message
     vector_store = client.beta.vector_stores.create(name=f"{book_name} Docs")
 
     # Step 2. Upload Knowledge (Markdown files)
-    console.print(f"Uploading knowledge from '{book_name}'...")
+    console.print(f"[bold blue]Uploading knowledge from '{book_name}'...[/bold blue]")  # Progress message
     md_files = load_markdown_files(book_name)
 
     file_streams = [open(file_path, "rb") for file_path in md_files]
@@ -50,14 +54,15 @@ def create_or_get_assistant(name, book_name):
     )
 
     while file_batch.status == "queued" or file_batch.status == "in_progress":
-        console.print(f"[bold yellow]{file_batch.status}...[/bold yellow]")
+        console.print(f"[bold yellow]{file_batch.status}...[/bold yellow]")  # Progress message
         time.sleep(1)
 
     # Step 3. Create the Assistant
+    console.print(f"[bold blue]Reading behavior instructions...[/bold blue]")  # Progress message
     with open('behaviors/default.txt', 'r') as file:
         instructions = file.read()
 
-    console.print(f"Creating assistant {name}...")
+    console.print(f"[bold blue]Creating assistant '{name}'...[/bold blue]")  # Progress message
     assistant = client.beta.assistants.create(
         instructions=instructions,
         name=name,
@@ -66,21 +71,16 @@ def create_or_get_assistant(name, book_name):
     )
 
     # Associate the assistant with the vector store
+    console.print(f"[bold blue]Associating assistant '{name}' with vector store...[/bold blue]")  # Progress message
     assistant = client.beta.assistants.update(
         assistant_id=assistant.id,
         tool_resources={"file_search": {"vector_store_ids": [vector_store.id]}},
     )
 
-    console.print(f"[bold green]Assistant '{name}' created successfully.[/bold green]")
+    console.print(f"[bold green]Assistant '{name}' created successfully.[/bold green]")  # Success message
 
     return assistant
 
-import os
-import time
-
-import os
-import time
-
 # Function to create a message in a thread and handle async processing
 def create_message(thread_id, content, assistant, file_path=None):
     """
@@ -95,43 +95,50 @@ def create_message(thread_id, content, assistant, file_path=None):
     Returns:
         str: The text content generated or improved by the assistant.
     """
+    console.print(f"[bold blue]Creating message in thread {thread_id}...[/bold blue]")  # Progress message
+
     # Prepare the base prompt
     if file_path and os.path.exists(file_path):
+        console.print(f"[bold blue]Reading content from {file_path} for improvement...[/bold blue]")  # Progress message
         with open(file_path, 'r', encoding='utf-8') as f:
             file_content = f.read()
             # Append the file content to the prompt asking for improvement
             prompt = f"{content}\n\nHere is the existing content to improve:\n{file_content}"
     else:
+        console.print(f"[bold blue]Using provided prompt to generate new content...[/bold blue]")  # Progress message
         # If no file, use the original prompt for generating new content
         prompt = content
 
     # Send the prompt to OpenAI completions API
+    console.print("[bold blue]Sending prompt to OpenAI API...[/bold blue]")  # Progress message
     response = client.chat.completions.create(
         model="gpt-4o",  # Adjust model as needed
         messages=[{"role": "user", "content": prompt}]
     )
 
     # Retrieve the completion result (generated or improved content)
     generated_content = response.choices[0].message.content
+    console.print(f"[bold green]Generated content received.[/bold green]")  # Success message
 
     # If a file path is provided, save the result to the file
     if file_path:
+        console.print(f"[bold blue]Saving generated content to {file_path}...[/bold blue]")  # Progress message
         with open(file_path, 'w', encoding='utf-8') as f:
             f.write(generated_content)
+        console.print(f"[bold green]Content saved successfully to {file_path}.[/bold green]")  # Success message
 
     # Return the generated content
     return generated_content
 
-
-
 # Function to get a new thread
 def get_thread():
+    console.print(f"[bold blue]Creating new thread...[/bold blue]")  # Progress message
     return client.beta.threads.create()
 
 # Function to update the assistant's knowledge with new files
 def update_agent_files(book_name, assistant):
     """Update the assistant with new Markdown files from the book."""
-    console.print(f"Updating assistant '{assistant.name}' with new files from '{book_name}'...")
+    console.print(f"[bold blue]Updating assistant '{assistant.name}' with new files from '{book_name}'...[/bold blue]")  # Progress message
     md_files = load_markdown_files(book_name)
 
     file_streams = [open(file_path, "rb") for file_path in md_files]
@@ -142,7 +149,7 @@ def update_agent_files(book_name, assistant):
     )
 
     while file_batch.status == "queued" or file_batch.status == "in_progress":
-        console.print(f"[bold yellow]{file_batch.status}...[/bold yellow]")
+        console.print(f"[bold yellow]{file_batch.status}...[/bold yellow]")  # Progress message
         time.sleep(1)
 
-    console.print(f"[bold green]Files updated successfully in assistant '{assistant.name}'.[/bold green]")
+    console.print(f"[bold green]Files updated successfully in assistant '{assistant.name}'.[/bold green]")  # Success message
diff --git a/storycraftr/chapters.py b/storycraftr/chapters.py
@@ -1,18 +1,24 @@
 import os
 from storycraftr.agents import create_or_get_assistant, get_thread, create_message
 from storycraftr.core import get_config
+from rich.console import Console
+
+console = Console()
 
 # Function to save content to a markdown file
 def save_to_markdown(book_name, file_name, header, content):
     """Save the generated content to the specified markdown file."""
     file_path = os.path.join(book_name, 'chapters', file_name)
+    console.print(f"[bold blue]Saving content to {file_path}...[/bold blue]")  # Progress message
     with open(file_path, 'w') as f:
         f.write(f"# {header}\n\n{content}")
+    console.print(f"[bold green]Content saved successfully to {file_path}[/bold green]")  # Success message
     return file_path  # Return the path for potential reuse
 
 # Function to generate a new chapter based on a prompt
 def generate_chapter(book_name, prompt, chapter_number):
     """Generate a new chapter based on a prompt."""
+    console.print(f"[bold blue]Generating chapter {chapter_number}...[/bold blue]")  # Progress message
     assistant = create_or_get_assistant(book_name, book_name)
     thread = get_thread()
 
@@ -22,6 +28,7 @@ def generate_chapter(book_name, prompt, chapter_number):
 
     # Check if the file exists and pass it as an attachment
     if os.path.exists(file_path):
+        console.print(f"[yellow]Existing chapter found at {file_path}. Attaching for further refinement...[/yellow]")  # Progress message
         content = f"Use the attached chapter file as a reference to evolve and improve the content based on this prompt: {prompt}. Write it in {get_config(book_name).primary_language}."
         chapter_content = create_message(
             thread_id=thread.id,
@@ -30,6 +37,7 @@ def generate_chapter(book_name, prompt, chapter_number):
             file_path=file_path
         )
     else:
+        console.print("[yellow]No existing chapter found. Generating new content...[/yellow]")  # Progress message
         content = f"Write a detailed chapter for the following book premise: {prompt}. Write it in {get_config(book_name).primary_language}."
         chapter_content = create_message(
             thread_id=thread.id,
@@ -39,13 +47,15 @@ def generate_chapter(book_name, prompt, chapter_number):
 
     # Save the updated chapter content to markdown
     save_to_markdown(book_name, chapter_file, f"Chapter {chapter_number}", chapter_content)
+    console.print(f"[bold green]✔ Chapter {chapter_number} generated successfully[/bold green]")  # Success message
     return chapter_content
 
 def generate_cover(book_name, prompt):
     """
     Generate a professional book cover in markdown format using the book's metadata 
     and a prompt for additional guidance.
     """
+    console.print("[bold blue]Generating book cover...[/bold blue]")  # Progress message
     # Obtener los datos del archivo de configuración
     config = get_config(book_name)
     language = config.primary_language
@@ -76,13 +86,15 @@ def generate_cover(book_name, prompt):
 
     # Guardar el contenido en el archivo markdown
     save_to_markdown(book_name, "cover.md", "Cover", cover_content)
+    console.print("[bold green]✔ Cover generated successfully[/bold green]")  # Success message
 
     return cover_content
 
 
 # Function to generate the back cover page
 def generate_back_cover(book_name, prompt):
     """Generate the back cover page for the book."""
+    console.print("[bold blue]Generating back cover...[/bold blue]")  # Progress message
     assistant = create_or_get_assistant(book_name, book_name)
     thread = get_thread()
 
@@ -95,11 +107,13 @@ def generate_back_cover(book_name, prompt):
 
     # Save to markdown
     save_to_markdown(book_name, "back_cover.md", "Back Cover", back_cover_content)
+    console.print("[bold green]✔ Back cover generated successfully[/bold green]")  # Success message
     return back_cover_content
 
 # Function to generate the epilogue of the book
 def generate_epilogue(book_name, prompt):
     """Generate the epilogue for the book."""
+    console.print("[bold blue]Generating epilogue...[/bold blue]")  # Progress message
     assistant = create_or_get_assistant(book_name, book_name)
     thread = get_thread()
 
@@ -108,6 +122,7 @@ def generate_epilogue(book_name, prompt):
 
     # Check if the file exists and pass it as an attachment
     if os.path.exists(file_path):
+        console.print(f"[yellow]Existing epilogue found at {file_path}. Attaching for further refinement...[/yellow]")  # Progress message
         content = f"Use the attached epilogue file as a reference to evolve and improve the content based on this prompt: {prompt}. Write it in {get_config(book_name).primary_language}."
         epilogue_content = create_message(
             thread_id=thread.id,
@@ -116,6 +131,7 @@ def generate_epilogue(book_name, prompt):
             file_path=file_path
         )
     else:
+        console.print("[yellow]No existing epilogue found. Generating new content...[/yellow]")  # Progress message
         content = f"Generate the epilogue for the book based on this prompt: {prompt}. Write it in {get_config(book_name).primary_language}."
         epilogue_content = create_message(
             thread_id=thread.id,
@@ -125,4 +141,5 @@ def generate_epilogue(book_name, prompt):
 
     # Save the updated epilogue content to markdown
     save_to_markdown(book_name, "epilogue.md", "Epilogue", epilogue_content)
+    console.print("[bold green]✔ Epilogue generated successfully[/bold green]")  # Success message
     return epilogue_content
diff --git a/storycraftr/cli.py b/storycraftr/cli.py
@@ -155,7 +155,6 @@ def geography(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_geography(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Geography generated successfully.", style="bold green")
 
 @worldbuilding.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -166,7 +165,6 @@ def history(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_history(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] History generated successfully.", style="bold green")
 
 @worldbuilding.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -177,7 +175,6 @@ def culture(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_culture(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Culture generated successfully.", style="bold green")
 
 @worldbuilding.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -188,7 +185,6 @@ def magic_system(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_magic_system(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Magic system generated successfully.", style="bold green")
 
 @worldbuilding.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -199,7 +195,6 @@ def technology(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_technology(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Technology generated successfully.", style="bold green")
 
 # CLI for outline
 @cli.group()
@@ -216,7 +211,6 @@ def general_outline(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_general_outline(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] General outline generated successfully.", style="bold green")
 
 @outline.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -227,7 +221,6 @@ def character_summary(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_character_summary(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Character summary generated successfully.", style="bold green")
 
 @outline.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -238,7 +231,6 @@ def plot_points(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_plot_points(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Plot points generated successfully.", style="bold green")
 
 @outline.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -249,7 +241,6 @@ def chapter_synopsis(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_chapter_synopsis(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Chapter synopsis generated successfully.", style="bold green")
 
 # CLI for chapters
 @cli.group()
@@ -267,7 +258,6 @@ def chapter(chapter_number, prompt):
         book_name = os.getcwd()
     load_config(book_name)
     generate_chapter(book_name, chapter_number, prompt)
-    console.print(f"[bold green]✔[/bold green] Chapter {chapter_number} generated successfully.", style="bold green")
 
 @chapters.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -278,7 +268,6 @@ def cover(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_cover(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Cover generated successfully.", style="bold green")
 
 @chapters.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -289,7 +278,6 @@ def back_cover(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_back_cover(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Back cover generated successfully.", style="bold green")
 
 @chapters.command()
 @click.option('--book-name', type=click.Path(), help='Path to the book directory')
@@ -300,7 +288,6 @@ def epilogue(prompt, book_name=None):
         book_name = os.getcwd()
     load_config(book_name)
     generate_epilogue(book_name, prompt)
-    console.print(f"[bold green]✔[/bold green] Epilogue generated successfully.", style="bold green")
 
 
 if __name__ == "__main__":