# File: multi_agent_book_workflow.py # Description: Core multi-agent book writing orchestration system import os import uuid import numpy as np import streamlit as st from typing import Dict, List, Any, Optional from pydantic import BaseModel, Field # Vector Store and Embedding from langchain.embeddings import OpenAIEmbeddings from langchain.vectorstores import FAISS from langchain.docstore.document import Document # Agent and LLM Frameworks from crewai import Agent, Task, Crew from langchain_openai import ChatOpenAI from langchain_core.language_models.chat_models import BaseChatModel from langchain.schema import AIMessage, HumanMessage from langchain.memory import ConversationBufferMemory import anthropic class CustomChatAnthropic(BaseChatModel, BaseModel): """Custom Chat Model for Anthropic""" anthropic_api_key: str = Field(...) model_name: str = Field(...) temperature: float = Field(default=0.7) client: Any = None def __init__(self, **data): super().__init__(**data) self.client = anthropic.Anthropic(api_key=self.anthropic_api_key) def _generate(self, messages, stop=None, run_manager=None, **kwargs): """Generate a response""" try: message_content = [] for message in messages: if isinstance(message, HumanMessage): message_content.append({"role": "user", "content": message.content}) elif isinstance(message, AIMessage): message_content.append({"role": "assistant", "content": message.content}) response = self.client.messages.create( model=self.model_name, messages=message_content, temperature=self.temperature, max_tokens=1000 ) return AIMessage(content=response.content[0].text) except Exception as e: st.error(f"Error generating response: {str(e)}") raise @property def _llm_type(self) -> str: """Return type of llm.""" return "anthropic" class BookWritingOrchestrator: def __init__(self, api_key: Optional[str] = None): """ Initialize the book writing orchestrator with multi-agent setup """ # Generate a unique project identifier self.project_id = str(uuid.uuid4()) # Set up API keys and debug info self._setup_api_keys(api_key) # Initialize context management self._initialize_context_store() # Setup AI models self._setup_ai_models() # Initialize agents try: self.setup_agents() except Exception as e: st.error(f"Agent setup failed: {e}") raise RuntimeError("Could not initialize book writing agents") from e def _setup_api_keys(self, api_key: Optional[str] = None): """ Validate and set up API keys with comprehensive checks for Hugging Face Spaces """ try: # Debug: Print available environment variables print("Available environment variables:", [k for k in os.environ.keys()]) # OpenAI API Key self.openai_api_key = os.getenv('OPENAI_API_KEY') print("OpenAI API Key found:", bool(self.openai_api_key)) # Anthropic API Key self.anthropic_api_key = os.getenv('ANTHROPIC_API_KEY') print("Anthropic API Key found:", bool(self.anthropic_api_key)) if not self.openai_api_key: raise ValueError("OpenAI API key not found in environment variables") if not self.anthropic_api_key: raise ValueError("Anthropic API key not found in environment variables") except Exception as e: st.error(f"Error accessing API keys: {str(e)}") raise def _initialize_context_store(self): """ Set up FAISS vector store for context management """ 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", metadata={"project_id": self.project_id} ) ], embedding=self.embeddings ) except Exception as e: st.error(f"Context store initialization failed: {e}") raise def _setup_ai_models(self): """ Configure AI models with robust error handling """ try: # OpenAI Configuration self.openai_llm = ChatOpenAI( model="gpt-4-turbo-preview", temperature=0.7, api_key=self.openai_api_key ) # Anthropic Configuration with custom wrapper self.anthropic_llm = CustomChatAnthropic( model_name="claude-3-sonnet-20240229", 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", return_messages=True ) except Exception as e: st.error(f"AI model setup failed: {str(e)}") print(f"Detailed error: {str(e)}") raise def setup_agents(self): """ Create specialized agents for book writing workflow """ try: # Concept Development Agent self.concept_agent = Agent( role='Book Concept Architect', goal='Develop a comprehensive and compelling book concept', backstory='An expert literary strategist who transforms raw ideas into structured book frameworks', verbose=True, llm=self.openai_llm, memory=True ) # Research Agent self.research_agent = Agent( role='Literary Research Specialist', goal='Conduct in-depth research to support the book\'s narrative and authenticity', backstory='A meticulous researcher with expertise in gathering and synthesizing complex information', verbose=True, llm=self.anthropic_llm, memory=True ) # Content Generation Agent self.writing_agent = Agent( role='Master Storyteller', goal='Create engaging, coherent, and stylistically consistent narrative content', backstory='A versatile writer capable of crafting compelling prose across various genres', verbose=True, llm=self.openai_llm, memory=True ) # Editing Agent self.editing_agent = Agent( role='Narrative Refinement Specialist', goal='Ensure narrative consistency, quality, and stylistic excellence', backstory='A seasoned editor with a keen eye for storytelling nuance and structural integrity', verbose=True, llm=self.anthropic_llm, memory=True ) except Exception as e: st.error(f"Error setting up agents: {str(e)}") raise def _fallback_book_concept(self, initial_prompt: str) -> Dict[str, Any]: """Provide fallback book concept""" return { "title": "Untitled Project", "genre": "General Fiction", "target_audience": "General Adult", "core_premise": initial_prompt, "chapter_outline": ["Chapter 1: Introduction"], "narrative_approach": "Standard Third-Person Narrative", "status": "fallback_generated" } def _fallback_chapter_content(self, book_concept: Dict[str, Any], chapter_number: int) -> str: """Provide fallback 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. """ def generate_book_concept(self, initial_prompt: str) -> Dict[str, Any]: """ Generate a comprehensive book concept using multi-agent collaboration """ try: # Concept Development Task concept_task = Task( description=f""" Develop a comprehensive book concept based on this initial idea: {initial_prompt} Provide detailed outputs including: 1. Title: A unique and compelling book title 2. Genre: Primary genre and any relevant subgenres 3. Target Audience: Specific demographic and reader profile 4. Core Premise: The central concept or hook 5. Chapter Outline: Brief outline of proposed chapters 6. Narrative Approach: Point of view and stylistic elements 7. Key Themes: Major themes to be explored Format the output as a structured JSON object. """, expected_output="A detailed JSON-formatted book concept with title, genre, target audience, and other key elements", agent=self.concept_agent ) # Research Validation Task research_task = Task( description=""" Review and enhance the generated book concept with: 1. Market analysis and genre conventions 2. Comparable successful titles 3. Unique selling points 4. Potential areas for deeper exploration Add these insights to the concept structure. """, expected_output="A detailed analysis and enhancement of the book concept with market insights and comparative analysis", agent=self.research_agent ) # Create Crew for Collaborative Processing book_concept_crew = Crew( agents=[self.concept_agent, self.research_agent], tasks=[concept_task, research_task], verbose=True ) # Execute Collaborative Workflow result = book_concept_crew.kickoff() # Parse and structure the result parsed_concept = self._parse_concept(result) # Store Context self._store_context('book_concept', str(parsed_concept)) return parsed_concept 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], chapter_number: int) -> str: """Generate content for a specific chapter""" try: # Get previous context previous_context = self._retrieve_context(chapter_number - 1) if chapter_number > 1 else "" # Writing Task writing_task = Task( description=f""" Write Chapter {chapter_number} for: Title: {book_concept.get('title', 'Untitled')} Genre: {book_concept.get('genre', 'Unspecified')} Previous Context: {previous_context} Requirements: 1. Follow the established narrative style 2. Maintain consistency with previous chapters 3. Advance the plot or themes meaningfully 4. Include appropriate scene-setting and character development Generate a complete, polished chapter. """, expected_output="A complete, polished chapter with narrative continuity and character development", agent=self.writing_agent ) # Editing Task editing_task = Task( description=""" Review and refine the chapter focusing on: 1. Narrative flow and pacing 2. Character consistency 3. Thematic development 4. Language and style polish Provide a final, edited version. """, expected_output="A refined and polished version of the chapter with improved narrative quality", agent=self.editing_agent ) # Create Crew for Chapter Generation chapter_crew = Crew( agents=[self.writing_agent, self.editing_agent], tasks=[writing_task, editing_task], verbose=True ) # Generate Chapter chapter_content = chapter_crew.kickoff() # Store Context self._store_context(f'chapter_{chapter_number}', chapter_content) # Update Memory self.global_memory.chat_memory.add_user_message( f"Chapter {chapter_number} Summary: {chapter_content[:500]}..." ) return chapter_content except Exception as e: st.error(f"Chapter generation failed: {e}") return self._fallback_chapter_content(book_concept, chapter_number) def _parse_concept(self, raw_concept: str) -> Dict[str, Any]: """Parse the raw concept output into a structured format""" try: # Split content into sections sections = raw_concept.strip().split('\n\n') concept = {} # Extract key-value pairs for section in sections: lines = section.strip().split('\n') for line in lines: if ':' in line: key, value = line.split(':', 1) concept[key.strip()] = value.strip() # Ensure required fields required_fields = ['title', 'genre', 'target_audience', 'core_premise'] for field in required_fields: if field not in concept: concept[field] = "Unspecified" return concept except Exception as e: st.error(f"Concept parsing failed: {e}") return self._fallback_book_concept("Failed to parse concept") def _store_context(self, context_key: str, content: str): """Store context in FAISS vector store""" try: document = Document( page_content=content, metadata={ "project_id": self.project_id, "context_key": context_key } ) 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}") def _retrieve_context(self, chapter_number: int) -> str: """Retrieve context for a specific chapter""" try: search_results = self.context_store.similarity_search( f"chapter_{chapter_number}", k=1 ) if search_results: return search_results[0].page_content return "" except Exception as e: st.error(f"Context retrieval failed: {e}") return "" def main(): """ Demonstration of BookWritingOrchestrator """ try: 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()