Update multi_agent_book_workflow.py

multi_agent_book_workflow.py

@@ -20,26 +20,27 @@ from langchain_openai import ChatOpenAI
 from langchain_anthropic import ChatAnthropic
 from langchain.memory import ConversationBufferMemory
 
+
 class BookWritingOrchestrator:
     def __init__(self, api_key=None):
         """
         Initialize the book writing orchestrator with multi-agent setup
-        
+
         Args:
             api_key (str, optional): API key for models if not using environment variables
         """
         # Generate a unique project identifier
         self.project_id = str(uuid.uuid4())
-        
+
         # Comprehensive API key management
         self._setup_api_keys(api_key)
-        
+
         # Initialize context management
         self._initialize_context_store()
-        
+
         # Setup AI models
         self._setup_ai_models()
-        
+
         # Initialize agents with robust error handling
         try:
             self.setup_agents()
@@ -55,7 +56,7 @@ class BookWritingOrchestrator:
         self.openai_api_key = api_key or os.getenv('OPENAI_API_KEY')
         if not self.openai_api_key:
             st.warning("OpenAI API key is missing. Some features may be limited.")
-        
+
         # Anthropic API Key
         self.anthropic_api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
         if not self.anthropic_api_key:
@@ -68,15 +69,15 @@ class BookWritingOrchestrator:
         try:
             # Initialize embeddings
             self.embeddings = OpenAIEmbeddings(api_key=self.openai_api_key)
-            
+
             # Create initial empty FAISS vector store
             self.context_store = FAISS.from_documents(
                 documents=[
                     Document(
-                        page_content="Initial project context", 
+                        page_content="Initial project context",
                         metadata={"project_id": self.project_id}
                     )
-                ], 
+                ],
                 embedding=self.embeddings
             )
         except Exception as e:
@@ -84,9 +85,9 @@ class BookWritingOrchestrator:
             raise
 
     def _setup_ai_models(self):
-    """
-    Configure AI models with robust error handling
-    """
+        """
+        Configure AI models with robust error handling
+        """
         try:
             # LLM Configurations
             self.openai_llm = ChatOpenAI(
@@ -94,14 +95,14 @@ class BookWritingOrchestrator:
                 temperature=0.7,
                 api_key=self.openai_api_key
             )
-            
+
             # Updated Anthropic model configuration
             self.anthropic_llm = ChatAnthropic(
                 model_name="claude-3-sonnet-20240229",  # Updated model name
                 temperature=0.7,
                 anthropic_api_key=self.anthropic_api_key
             )
-            
+
             # Global memory for cross-chapter context
             self.global_memory = ConversationBufferMemory(
                 memory_key="book_writing_context",
@@ -114,10 +115,10 @@ class BookWritingOrchestrator:
     def _fallback_book_concept(self, initial_prompt: str) -> Dict[str, Any]:
         """
         Provide a fallback response when book concept generation fails
-        
+
         Args:
             initial_prompt (str): Original user prompt
-            
+
         Returns:
             Dict: Basic book concept structure
         """
@@ -134,19 +135,19 @@ class BookWritingOrchestrator:
     def _fallback_chapter_content(self, book_concept: Dict[str, Any], chapter_number: int) -> str:
         """
         Provide fallback chapter content when generation fails
-        
+
         Args:
             book_concept (Dict): Book concept data
             chapter_number (int): Chapter number
-            
+
         Returns:
             str: Basic chapter content
         """
         return f"""
         Chapter {chapter_number}
-        
+
         [Placeholder content for {book_concept.get('title', 'Untitled')}]
-        
+
         This is auto-generated fallback content due to an error in chapter generation.
         Please try regenerating this chapter or contact support if the issue persists.
         """
@@ -165,7 +166,7 @@ class BookWritingOrchestrator:
             llm=self.openai_llm,
             memory=True
         )
-        
+
         # Research Agent
         self.research_agent = Agent(
             role='Literary Research Specialist',
@@ -175,7 +176,7 @@ class BookWritingOrchestrator:
             llm=self.anthropic_llm,
             memory=True
         )
-        
+
         # Content Generation Agent
         self.writing_agent = Agent(
             role='Master Storyteller',
@@ -185,7 +186,7 @@ class BookWritingOrchestrator:
             llm=self.openai_llm,
             memory=True
         )
-        
+
         # Editing Agent
         self.editing_agent = Agent(
             role='Narrative Refinement Specialist',
@@ -199,10 +200,10 @@ class BookWritingOrchestrator:
     def generate_book_concept(self, initial_prompt: str) -> Dict[str, Any]:
         """
         Generate a comprehensive book concept using multi-agent collaboration
-        
+
         Args:
             initial_prompt (str): User's initial book idea
-        
+
         Returns:
             Dict containing book concept details
         """
@@ -211,7 +212,7 @@ class BookWritingOrchestrator:
             description=f"""
             Develop a comprehensive book concept based on this initial idea:
             {initial_prompt}
-            
+
             Provide detailed outputs including:
             1. Unique Book Title
             2. Genre and Subgenre
@@ -223,62 +224,62 @@ class BookWritingOrchestrator:
             agent=self.concept_agent,
             expected_output="A detailed JSON-like structure of the book concept"
         )
-        
+
         # Research Validation Task
         research_task = Task(
             description="Validate and enrich the book concept with additional research insights",
             agent=self.research_agent,
             expected_output="Research-backed annotations and potential depth areas"
         )
-        
+
         # Create Crew for Collaborative Processing
         book_concept_crew = Crew(
             agents=[self.concept_agent, self.research_agent],
             tasks=[concept_task, research_task],
             verbose=True  # Changed from 2 to True
         )
-        
+
         # Execute Collaborative Workflow
         try:
             result = book_concept_crew.kickoff()
-            
+
             # Store Context in FAISS Vector Store
             self._store_context('book_concept', result)
-            
+
             # Parse and return concept
             return self._parse_concept(result)
-        
+
         except Exception as e:
             st.error(f"Book concept generation failed: {e}")
             return self._fallback_book_concept(initial_prompt)
 
     def generate_chapter_content(
-        self, 
-        book_concept: Dict[str, Any], 
+        self,
+        book_concept: Dict[str, Any],
         chapter_number: int
     ) -> str:
         """
         Generate content for a specific chapter using multi-agent workflow
-        
+
         Args:
             book_concept (Dict): Comprehensive book concept
             chapter_number (int): Chapter to generate
-        
+
         Returns:
             str: Generated chapter content
         """
         # Retrieve Previous Context
         previous_context = self._retrieve_context(chapter_number)
-        
+
         # Content Generation Task
         writing_task = Task(
             description=f"""
             Write Chapter {chapter_number} for the book
-            
+
             Book Concept: {book_concept}
-            
+
             Previous Context: {previous_context}
-            
+
             Generate a draft that:
             1. Maintains narrative continuity
             2. Advances the story
@@ -288,43 +289,43 @@ class BookWritingOrchestrator:
             agent=self.writing_agent,
             expected_output="A complete chapter draft with narrative coherence"
         )
-        
+
         # Editing and Refinement Task
         editing_task = Task(
             description="Review and refine the generated chapter draft",
             agent=self.editing_agent,
             expected_output="Polished chapter content with suggested improvements"
         )
-        
+
         # Create Crew for Chapter Generation
         chapter_crew = Crew(
             agents=[self.writing_agent, self.editing_agent],
             tasks=[writing_task, editing_task],
             verbose=True  # Changed from 2 to True
         )
-        
+
         try:
             # Execute Collaborative Workflow
             chapter_content = chapter_crew.kickoff()
-            
+
             # Store Chapter Context in FAISS
             self._store_context(f'chapter_{chapter_number}', chapter_content)
-            
+
             # Update global memory with chapter context
             self.global_memory.chat_memory.add_user_message(chapter_content)
-            
+
             return chapter_content
-        
+
         except Exception as e:
             st.error(f"Chapter generation failed: {e}")
             return self._fallback_chapter_content(book_concept, chapter_number)
 
     # Rest of the implementation remains the same as in the previous version
-    
+
     def _store_context(self, context_key: str, content: str):
         """
         Store context in FAISS vector store
-        
+
         Args:
             context_key (str): Unique identifier for the context
             content (str): Content to store
@@ -332,22 +333,23 @@ class BookWritingOrchestrator:
         try:
             # Create a document with metadata
             document = Document(
-                page_content=content, 
+                page_content=content,
                 metadata={
                     "project_id": self.project_id,
                     "context_key": context_key
                 }
             )
-            
+
             # Add document to FAISS vector store
             new_store = FAISS.from_documents([document], self.embeddings)
             self.context_store.merge_from(new_store)
-        
+
         except Exception as e:
             st.error(f"Context storage failed: {e}")
 
     # Remaining methods stay the same as in the previous implementation
 
+
 def main():
     """
     Demonstration of BookWritingOrchestrator
@@ -355,18 +357,19 @@ def main():
     try:
         # Create orchestrator
         orchestrator = BookWritingOrchestrator()
-        
+
         # Generate book concept
         initial_prompt = "A science fiction story about space exploration"
         book_concept = orchestrator.generate_book_concept(initial_prompt)
         print("Book Concept:", book_concept)
-        
+
         # Generate first chapter
         first_chapter = orchestrator.generate_chapter_content(book_concept, 1)
         print("\nFirst Chapter:\n", first_chapter)
-    
+
     except Exception as e:
         print(f"An error occurred: {e}")
 
+
 if __name__ == "__main__":
     main()