tools.py

from ebook_agent.text_gen import TextGenerator
from ebook_agent.image_gen import ImageGenerator
from ebook_agent.ebook_compiler import EbookCompiler
import os

class EbookAgentTools:
    def __init__(self, hf_token=None):
        token = hf_token or os.getenv("HF_TOKEN")
        self.text_gen = TextGenerator(token=token)
        self.image_gen = ImageGenerator(token=token)
        self.compiler = EbookCompiler()

    def create_book_outline(self, topic: str, genre: str) -> str:
        """Create a book outline."""
        return self.text_gen.outline_book(topic, genre)

    def write_book_chapter(self, title: str, context: str) -> str:
        """Write a specific chapter."""
        return self.text_gen.write_chapter(title, context)

    def create_book_cover(self, title: str, genre: str, description: str) -> str:
        """Generate a cover image."""
        return self.image_gen.generate_cover(title, genre, description)

    def build_final_ebook(self, title: str, author: str, chapters_json: str, cover_path: Optional[str] = None) -> str:
        """
        Compile the final ebook.
        chapters_json should be a list of dictionaries with 'title' and 'content' keys.
        """
        import json
        try:
            chapters = json.loads(chapters_json)
            return self.compiler.compile_epub(title, author, chapters, cover_path)
        except Exception as e:
            return f"Error compiling ebook: {str(e)}"