File size: 3,187 Bytes
58af945 ce2a268 58af945 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
#!/usr/bin/env python
from pydantic import BaseModel
from typing import List
from crewai.flow.flow import Flow, listen, start
from src.educational_books.crews.book_outline.book_outline import BookOutlineCrew
from src.educational_books.crews.book_content.book_content import BookContentCrew
from src.educational_books.types import Section, SectionOutline
import asyncio
class BookState(BaseModel):
title: str = "Mastering Data Structures and Algorithms"
book: List[Section] = []
book_outline: List[SectionOutline] = []
topic: str = "Sorting Algorithms"
book_outline_md: str = ""
book_md: str = ""
section_finished: int = 0
class BookFlow(Flow[BookState]):
initial_state = BookState
def write_to_file(self, filename: str, content: str, mode: str = "a"):
"""Helper method to write content to the file"""
with open(filename, mode) as f:
f.write(content)
@start()
def initialize(self, topic: str = None):
if topic:
self.state.topic = topic
@listen(initialize)
def generate_book_outline(self):
print("Generating book outline")
crew = BookOutlineCrew()
outline = crew.crew().kickoff(inputs={"topic": self.state.topic})
self.state.title = outline["title"]
for section in outline["sections"]:
self.state.book_outline_md += f"# {section.title}\n"
self.state.book_outline_md += f"{section.description}\n\n"
self.state.book_outline_md += f"## Covered Skills\n"
for skill in section.covered_skills:
self.state.book_outline_md += f"- {skill}\n"
self.state.book_outline_md += "\n"
self.state.book_outline_md += f"## Learning Objectives\n"
for objective in section.objectives:
self.state.book_outline_md += f"- {objective}\n"
self.state.book_outline_md += "\n"
self.state.book_outline = outline["sections"]
@listen(generate_book_outline)
async def generate_book(self):
print("Generating book")
crew = BookContentCrew()
tasks = []
book_outline = [(i, section.title, section.description, section.covered_skills) for i, section in enumerate(self.state.book_outline)]
async def generate_section(section):
content = crew.crew().kickoff(inputs={
"topic": self.state.topic,
"section_title": section.title,
"section_description": section.description,
"covered_skills": section.covered_skills,
"book_outline": book_outline
})
title = content["title"]
content = content["content"]
section = Section(title=title, content=content)
section_content = f"# {title}\n\n{content}\n\n"
self.state.book_md += section_content
self.state.section_finished += 1
return section
for section in self.state.book_outline:
print(f"Generating content for {section.title}")
task = asyncio.create_task(generate_section(section))
tasks.append(task)
sections = await asyncio.gather(*tasks)
self.state.book = sections
print(f"Book generated with {len(self.state.book)} sections")
def kickoff():
book_name = BookFlow()
book_name.kickoff()
def plot():
book_name = BookFlow()
book_name.plot()
if __name__ == "__main__":
kickoff()
|