multi_agent_book_workflow.py

# 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()