S11.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv(dotenv_path='.env')\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import LLMChain\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.docstore.document import Document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseStructureChain:\n",
    "    PROMPT = \"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.llm = ChatOpenAI()\n",
    "        self.chain = LLMChain.from_string(\n",
    "            llm=self.llm,\n",
    "            template=self.PROMPT,\n",
    "        )\n",
    "\n",
    "        self.chain.verbose = True\n",
    "\n",
    "class BaseEventChain:\n",
    "    PROMPT = \"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.llm = ChatOpenAI(model=\"gpt-3.5-turbo-16k\")\n",
    "        self.chain = LLMChain.from_string(\n",
    "            llm=self.llm,\n",
    "            template=self.PROMPT,\n",
    "        )\n",
    "\n",
    "        self.chain.verbose = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Automating book writing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Characters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MainCharacterChain(BaseStructureChain):\n",
    "    PROMPT = \"\"\"\n",
    "    You are provided with the resume of a person.\n",
    "    Describe the person's profile in a few sentences and include that person's name.\n",
    "\n",
    "    Resume: {text}\n",
    "\n",
    "    Profile:\n",
    "\"\"\"\n",
    "\n",
    "    def load_resume(self, file_path: str) -> Document:\n",
    "        loader = PyPDFLoader(file_path)\n",
    "        return loader.load_and_split()\n",
    "\n",
    "    def run(self, file_path: str):\n",
    "        docs = self.load_resume(file_path)\n",
    "        resume = \"\\n\\n\".join([doc.page_content for doc in docs])\n",
    "\n",
    "        return self.chain.predict(\n",
    "            text=resume\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.docstore.document import Document\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chains import LLMChain\n",
    "\n",
    "class TitleChain(BaseStructureChain):\n",
    "    PROMPT = \"\"\"\n",
    "    Your job is to generate the title for a novel about the following subject and main character.\n",
    "    Return a title and only a title!\n",
    "    The title should be consistent with the genre of the novel.\n",
    "    The title should be consistent with the style of the author.\n",
    "\n",
    "    Subject: {subject}\n",
    "    Genre: {genre}\n",
    "    Author: {author}\n",
    "\n",
    "    Main character's profile: {profile}\n",
    "\n",
    "    Title:\n",
    "\"\"\"\n",
    "    def __init__(self) -> None:\n",
    "        self.llm = ChatOpenAI()\n",
    "        self.chain = LLMChain.from_string(\n",
    "            llm=self.llm,\n",
    "            template=self.PROMPT,\n",
    "        )\n",
    "\n",
    "        self.chain.verbose = True\n",
    "\n",
    "    def load_resume(self, file_path: str) -> Document:\n",
    "        loader = PyPDFLoader(file_path)\n",
    "        return loader.load_and_split()\n",
    "\n",
    "    def run(self, subject: str, genre: str, author: str, profile: str):\n",
    "        return self.chain.predict(\n",
    "            subject=subject,\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            profile=profile,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlotChain(BaseStructureChain):\n",
    "    PROMPT = \"\"\"\n",
    "    Your job is to generate the plot for a novel. Return a plot and only a plot.\n",
    "    Describe the full plot of the story and don't hesitate to create new characters to make it compelling.\n",
    "    You are provided with the subject, title, and main character's profile.\n",
    "    Make sure that the main character is at the center of the story.\n",
    "    The plot should be consistent with the genre of the novel.\n",
    "    The plot should be consistent with the style of the author.\n",
    "\n",
    "    Consider the following attributes to write an exciting story:\n",
    "    {features}\n",
    "\n",
    "   \n",
    "    Subject: {subject}\n",
    "    Genre: {genre}\n",
    "    Author: {author}\n",
    "\n",
    "    Title: {title}\n",
    "    Main character's profile: {profile}\n",
    "\n",
    "    Plot:\n",
    "    \"\"\"\n",
    "\n",
    "    HELPER_PROMPT = \"\"\"\n",
    "    Generate a list of attributes that characterized an exciting story.\n",
    "\n",
    "    List of attributes:\"\"\"\n",
    "\n",
    "    def run(self, subject: str, genre: str, author: str, profile: str, title: str):\n",
    "        features = ChatOpenAI().predict(self.HELPER_PROMPT)\n",
    "\n",
    "        return self.chain.predict(\n",
    "            subject=subject,\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            profile=profile,\n",
    "            title=title,\n",
    "            features=features,\n",
    "        )\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chapter list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChapterListChain(BaseStructureChain):\n",
    "    PROMPT = \"\"\"\n",
    "    Your job is to generate a list of chapters.\n",
    "    ONLY the list and nothing more!\n",
    "\n",
    "    You are provided with a title, a plot, and a main character for a novel.\n",
    "    Gnerate a list of chapters describing the plot of that novel.\n",
    "    Make sure the chapters are consistent with the genre of the novel.\n",
    "    Make sure the chapters are consistent with the style of the author.\n",
    "\n",
    "    Follow this template:\n",
    "\n",
    "    Prologue: [description of the prologue]\n",
    "    Chapter 1: [description of the first chapter]\n",
    "    ...\n",
    "    Epilogue: [description of the epilogue]\n",
    "\n",
    "    Make sure the chapter is followed by the character `:` and its description. For example `Chapter 1: [description of the first chapter]`\n",
    "\n",
    "    Subject: {subject}\n",
    "    Genre: {genre}\n",
    "    Author: {author}\n",
    "    \n",
    "    Title: {title}\n",
    "    Main character's profile: {profile}\n",
    "\n",
    "    Plot: {plot}\n",
    "\n",
    "    Return the chapter list and only the chapter list:\n",
    "    \"\"\"\n",
    "\n",
    "    def run(self, subject: str, genre: str, author: str, profile: str, title: str, plot: str) -> dict:\n",
    "        response = self.chain.predict(\n",
    "            subject=subject,\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            profile=profile,\n",
    "            title=title,\n",
    "            plot=plot,\n",
    "        )\n",
    "\n",
    "        return self.parse(response)\n",
    "\n",
    "    def parse(self, response: str) -> dict:\n",
    "        chapter_list = response.strip().split(\"\\n\")\n",
    "        chapter_list = [chapter for chapter in chapter_list if \":\" in chapter]\n",
    "\n",
    "        return dict([\n",
    "            chapter.strip().split(\":\") for chapter in chapter_list\n",
    "        ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Chapter's plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChapterPlotChain(BaseEventChain):\n",
    "    HELPER_PROMPT = \"\"\"\n",
    "    Generate a list of attributes that characterized an exciting story.\n",
    "    List of attributes:\"\"\"\n",
    "\n",
    "    PROMPT = \"\"\"\n",
    "    You are a writer and your job is to generate the prot for one and only one chapter of a novel.\n",
    "\n",
    "    You are provided with a title, a plot, and a main character for a novel.\n",
    "    Additionally, you are provided with the plots of the previous chapters and the outline of the novel.\n",
    "\n",
    "    Each chapter should have its own arc, but it should also be consistent with the other chapters, and overall plot of the novel.\n",
    "\n",
    "    The summary should be consistent with the genre of the novel.\n",
    "    The summary should be consistent with the style of the author.\n",
    "\n",
    "    Consider the following attributes to write an exciting story:\n",
    "    {features}\n",
    "\n",
    "    Subject: {subject}\n",
    "    Genre: {genre}\n",
    "    Author: {author}\n",
    "\n",
    "    Title: {title}\n",
    "    Main character's profile: {profile}\n",
    "\n",
    "    Novel's plot: {plot}\n",
    "\n",
    "    Outline:\n",
    "    {outline}\n",
    "\n",
    "    Chapter plots:\n",
    "    {summaries}\n",
    "\n",
    "    Return the plot for the chapter and only the plot:\n",
    "    Plot for chapter {chapter}:\n",
    "    \"\"\"\n",
    "\n",
    "    def run(self, subject: str, genre: str, author: str, profile: str, title: str, plot: str, summaries_dict: dict, chapter_dict: dict, chapter:str) -> str:\n",
    "        features = ChatOpenAI().predict(self.HELPER_PROMPT)\n",
    "\n",
    "        outline = \"\\n\".join([\n",
    "            f\"{chapter}: {description}\" for chapter, description in chapter_dict.items()\n",
    "        ])\n",
    "\n",
    "        summaries = \"\\n\\n\".join([\n",
    "            f\"Plot of {chapter}: {summary}\" for chapter, summary in summaries_dict.items()\n",
    "        ])\n",
    "\n",
    "        return self.chain.predict(\n",
    "            subject=subject,\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            profile=profile,\n",
    "            title=title,\n",
    "            plot=plot,\n",
    "            outline=outline,\n",
    "            summaries=summaries,\n",
    "            chapter=chapter,\n",
    "            features=features,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Event"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EventChain(BaseEventChain):\n",
    "    PROMPT = \"\"\"\n",
    "    You are a writer and your job is to come up with a detailed list of events happens in the current chapter of a novel.\n",
    "    Those events describe the plot of the chapter, and the actions of the different characters in chronological order.\n",
    "    You are provided with the title, the main plot of the novel, the main character, and a summary of that chapter.\n",
    "    Additionally, you are provided with the list of the events that were outlined in the previous chapters.\n",
    "    The event list should be consistent with the genre of the novel.\n",
    "    The event list should be consistent with the style of the author.\n",
    "\n",
    "    The each element of that list should be returhned on different lines.\n",
    "    Follow this template:\n",
    "\n",
    "    Event 1\n",
    "    Event 2\n",
    "    ...\n",
    "    Final event\n",
    "\n",
    "    Subject: {subject}\n",
    "    Genre: {genre}\n",
    "    Author: {author}\n",
    "\n",
    "    Title: {title}\n",
    "    Main character's profile: {profile}\n",
    "\n",
    "    Novel's plot: {plot}\n",
    "\n",
    "    Events you outlined for previous chapters: {previous_events}\n",
    "\n",
    "    Summary of the current chapter: {summary}\n",
    "\n",
    "    Return the events and only the events!\n",
    "    Event list forr that chapter:\n",
    "\n",
    "\"\"\"\n",
    "    def run(self, subject: str, genre: str, author: str, profile: str, title: str, plot: str, event_dict: dict, summary: str) -> list[str]:\n",
    "        previous_events = \"\"\n",
    "        for chapter, events in event_dict.items():\n",
    "            previous_events += \"\\n\" + chapter\n",
    "\n",
    "            for event in events:\n",
    "                previous_events += \"\\n\" + event\n",
    "\n",
    "        response = self.chain.predict(\n",
    "            subject=subject,\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            profile=profile,\n",
    "            title=title,\n",
    "            plot=plot,\n",
    "            previous_events=previous_events,\n",
    "            summary=summary,\n",
    "        )\n",
    "\n",
    "        event_list = response.strip().split(\"\\n\")\n",
    "\n",
    "        return [event for event in event_list if event.strip()]\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing the book"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WriteChain(BaseEventChain):\n",
    "    PROMPT = \"\"\"\n",
    "    You are a novel writer. Thie is described by a list of events.\n",
    "    You have already written the novel up to the last event.\n",
    "    Your job is to generate the paragraphs of the novel about the new event.\n",
    "    You are provided with the title, a novel's plot, a description of the main character, and a summary of the current chapter.\n",
    "    Make sure the paragraphs are consistent with the plot of the chapter.\n",
    "    Additionally you are provided with the list of the events you have already written about.\n",
    "    The paragraphs should be consistent with the genre of the novel.\n",
    "    The paragraphs should be consistent with the style of the author.\n",
    "\n",
    "    Genre: {genre}\n",
    "    Author: {author}\n",
    "\n",
    "    Title: {title}\n",
    "    Main character's profile: {profile}\n",
    "\n",
    "    Novel's plot: {plot}\n",
    "    Previous events: {previous_events}\n",
    "\n",
    "    Current chapter summary: {summary}\n",
    "\n",
    "    Previous paragraphs: {previous_paragraphs}\n",
    "\n",
    "    New event you need to write about: {current_event}\n",
    "\n",
    "    Paragarphs of the novel describing that event:\n",
    "    \"\"\"\n",
    "\n",
    "    def run(self, genre: str, author: str, title: str, profile: str, plot: str, previous_events: list[str], summary: str, previous_paragraphs: str, current_event: str) -> str:\n",
    "        return self.chain.predict(\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            title=title,\n",
    "            profile=profile,\n",
    "            plot=plot,\n",
    "            previous_events=\"\\n\".join(previous_events),\n",
    "            summary=summary,\n",
    "            previous_paragraphs=previous_paragraphs,\n",
    "            current_event=current_event,\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_structure(subject: str, genre: str, author: str, profile) -> tuple[str, str, dict]:\n",
    "    title_chain = TitleChain()\n",
    "    plot_chain = PlotChain()\n",
    "    chapters_chain = ChapterListChain()\n",
    "\n",
    "    title = title_chain.run(subject, genre, author, profile)\n",
    "    plot = plot_chain.run(subject, genre, author, profile, title)\n",
    "    chapters = chapters_chain.run(subject, genre, author, profile, title, plot)\n",
    "\n",
    "    return title, plot, chapters\n",
    "\n",
    "def get_events(subject: str, genre: str, author: str, profile: str, title: str, plot: str, chapter_dict: dict) -> tuple[dict, dict]:\n",
    "    chapter_plot_chain = ChapterPlotChain()\n",
    "    event_chain = EventChain()\n",
    "\n",
    "    summaries_dict = {}\n",
    "    event_dict = {}\n",
    "    for chapter, _ in chapter_dict.items():\n",
    "        summaries_dict[chapter] = chapter_plot_chain.run(\n",
    "            subject=subject,\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            profile=profile,\n",
    "            title=title,\n",
    "            plot=plot,\n",
    "            summaries_dict=summaries_dict,\n",
    "            chapter_dict=chapter_dict,\n",
    "            chapter=chapter,\n",
    "        ) \n",
    "\n",
    "        event_dict[chapter] = event_chain.run(\n",
    "            subject=subject,\n",
    "            genre=genre,\n",
    "            author=author,\n",
    "            profile=profile,\n",
    "            title=title,\n",
    "            plot=plot,\n",
    "            event_dict=event_dict,\n",
    "            summary=summaries_dict[chapter],\n",
    "        )\n",
    "\n",
    "    return summaries_dict, event_dict\n",
    "\n",
    "def write_book(genre: str, author: str, title: str, profile: str, plot: str, summaries_dict: dict, event_dict: str) -> str:\n",
    "    writer_chain = WriteChain()\n",
    "    previous_events = []\n",
    "    book = {}\n",
    "    paragraphs = \"\"\n",
    "\n",
    "    for chapter, event_list in event_dict.items():\n",
    "        book[chapter] = []\n",
    "\n",
    "        for event in event_list:\n",
    "            paragraphs = writer_chain.run(\n",
    "                genre=genre,\n",
    "                author=author,\n",
    "                title=title,\n",
    "                profile=profile,\n",
    "                plot=plot,\n",
    "                previous_events=previous_events,\n",
    "                summary=summaries_dict[chapter],\n",
    "                previous_paragraphs=paragraphs,\n",
    "                current_event=event,\n",
    "            )\n",
    "\n",
    "            previous_events.append(event)\n",
    "            book[chapter].append(paragraphs)\n",
    "\n",
    "    return book"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main_character_chain = MainCharacterChain()\n",
    "\n",
    "subject = \"Machine War\"\n",
    "author = \"Ernest Hemingway\"\n",
    "genre = \"Science Fiction\"\n",
    "\n",
    "profile = main_character_chain.run(\"./docs/Profile.pdf\")\n",
    "\n",
    "title, plot, chapter_dict = get_structure(subject, genre, author, profile)\n",
    "summaries_dict, event_dict = get_events(subject, genre, author, profile, title, plot, chapter_dict)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "book = write_book(genre, author, title, profile, plot, summaries_dict, event_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write to file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "import docx\n",
    "\n",
    "class DocWriter:\n",
    "    def __init__(self) -> None:\n",
    "        self.doc = docx.Document()\n",
    "\n",
    "    def write_doc(self, book, chapter_dict, title):\n",
    "        self.doc.add_heading(title, 0)\n",
    "\n",
    "        for chapter, paragraph_list in book.items():\n",
    "            description = chapter_dict[chapter]\n",
    "\n",
    "            chapter_name = f\"Chapter {chapter.strip()}: {description.strip()}\"\n",
    "\n",
    "            self.doc.add_heading(chapter_name, 1)\n",
    "\n",
    "            text = \"\\n\\n\".join(paragraph_list)\n",
    "            self.doc.add_paragraph(text)\n",
    "\n",
    "        self.doc.save('./docs/book.docx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_writer = DocWriter()\n",
    "\n",
    "doc_writer.write_doc(book, chapter_dict, title)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}