Update multi_agent_book_workflow.py

multi_agent_book_workflow.py

@@ -1,12 +1,12 @@
 # File: multi_agent_book_workflow.py
-# Location: /multi_agent_book_workflow.py (in root directory)
+# Location: /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
+from typing import Dict, List, Any, Optional
 
 # Vector Store and Embedding
 from langchain.embeddings import OpenAIEmbeddings
@@ -16,11 +16,46 @@ from langchain.docstore.document import Document
 # Agent and LLM Frameworks
 from crewai import Agent, Task, Crew
 from langchain_openai import ChatOpenAI
-from langchain.chat_models import ChatAnthropic
+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):
+    """Custom Chat Model for Anthropic"""
+    
+    def __init__(self, model_name: str, temperature: float, anthropic_api_key: str):
+        """Initialize the custom Anthropic chat model"""
+        super().__init__()
+        self.client = anthropic.Anthropic(api_key=anthropic_api_key)
+        self.model_name = model_name
+        self.temperature = temperature
+
+    def _generate(self, messages, stop=None, run_manager=None, **kwargs):
+        """Generate a response"""
+        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)
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "anthropic"
 
 class BookWritingOrchestrator:
-    def __init__(self, api_key=None):
+    def __init__(self, api_key: Optional[str] = None):
         """
         Initialize the book writing orchestrator with multi-agent setup
         
@@ -46,7 +81,7 @@ class BookWritingOrchestrator:
             st.error(f"Agent setup failed: {e}")
             raise RuntimeError("Could not initialize book writing agents") from e
 
-    def _setup_api_keys(self, api_key=None):
+    def _setup_api_keys(self, api_key: Optional[str] = None):
         """
         Validate and set up API keys with comprehensive checks
         """
@@ -94,11 +129,11 @@ class BookWritingOrchestrator:
                 api_key=self.openai_api_key
             )
             
-            # Anthropic Configuration
-            self.anthropic_llm = ChatAnthropic(
-                model="claude-3-sonnet-20240229",
-                anthropic_api_key=self.anthropic_api_key,
-                temperature=0.7
+            # 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
@@ -156,46 +191,6 @@ class BookWritingOrchestrator:
             memory=True
         )
 
-    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
-        """
-        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 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.
-        """
-
     def generate_book_concept(self, initial_prompt: str) -> Dict[str, Any]:
         """
         Generate a comprehensive book concept using multi-agent collaboration
@@ -213,22 +208,31 @@ class BookWritingOrchestrator:
             {initial_prompt}
             
             Provide detailed outputs including:
-            1. Unique Book Title
-            2. Genre and Subgenre
-            3. Target Audience
-            4. Core Narrative Premise
-            5. Potential Chapter Outline
-            6. Distinctive Narrative Approach
+            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.
             """,
-            agent=self.concept_agent,
-            expected_output="A detailed JSON-like structure of the book concept"
+            agent=self.concept_agent
         )
         
         # 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"
+            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.
+            """,
+            agent=self.research_agent
         )
         
         # Create Crew for Collaborative Processing
@@ -238,106 +242,68 @@ class BookWritingOrchestrator:
             verbose=True
         )
         
-        # Execute Collaborative Workflow
         try:
+            # Execute Collaborative Workflow
             result = book_concept_crew.kickoff()
             
+            # Parse and structure the result
+            parsed_concept = self._parse_concept(result)
+            
             # Store Context in FAISS Vector Store
-            self._store_context('book_concept', result)
+            self._store_context('book_concept', str(parsed_concept))
+            
+            return parsed_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 _store_context(self, context_key: str, content: str):
-        """
-        Store context in FAISS vector store
+    def generate_chapter_content(self, book_concept: Dict[str, Any], chapter_number: int) -> str:
         """
-        try:
-            # Create a document with metadata
-            document = Document(
-                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}")
-
-    def _parse_concept(self, raw_concept: str) -> Dict[str, Any]:
-        """
-        Parse the raw concept output into a structured format
+        Generate content for a specific chapter
         
         Args:
-            raw_concept (str): Raw concept from agents
-            
-        Returns:
-            Dict: Structured book concept
-        """
-        try:
-            # Basic parsing - in practice you might want more sophisticated parsing
-            lines = raw_concept.strip().split('\n')
-            concept = {}
-            
-            for line in lines:
-                if ':' in line:
-                    key, value = line.split(':', 1)
-                    concept[key.strip()] = value.strip()
+            book_concept (Dict): Book concept data
+            chapter_number (int): Chapter number to generate
             
-            return concept
-        except Exception as e:
-            st.error(f"Concept parsing failed: {e}")
-            return self._fallback_book_concept("Failed to parse concept")
-
-    def generate_chapter_content(
-        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
         """
         # Get previous context if available
         previous_context = self._retrieve_context(chapter_number - 1) if chapter_number > 1 else ""
         
-        # Content Generation Task
+        # Writing Task
         writing_task = Task(
             description=f"""
-            Write Chapter {chapter_number} for the book
+            Write Chapter {chapter_number} for:
+            Title: {book_concept.get('title', 'Untitled')}
+            Genre: {book_concept.get('genre', 'Unspecified')}
             
-            Book Concept: {book_concept}
             Previous Context: {previous_context}
             
-            Generate a draft that:
-            1. Maintains narrative continuity
-            2. Advances the story
-            3. Matches the established tone and style
-            4. Incorporates research insights
+            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.
             """,
-            agent=self.writing_agent,
-            expected_output="A complete chapter draft with narrative coherence"
+            agent=self.writing_agent
         )
         
         # Editing Task
         editing_task = Task(
-            description="Review and refine the generated chapter draft",
-            agent=self.editing_agent,
-            expected_output="Polished chapter content with suggested improvements"
+            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.
+            """,
+            agent=self.editing_agent
         )
         
         # Create Crew for Chapter Generation
@@ -354,29 +320,74 @@ class BookWritingOrchestrator:
             # Store Chapter Context
             self._store_context(f'chapter_{chapter_number}', chapter_content)
             
-            # Update global memory
+            # Update global memory with chapter summary
             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 the 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
         
         Args:
-            chapter_number (int): Chapter number to retrieve context for
+            chapter_number (int): Chapter number
             
         Returns:
             str: Retrieved context or empty string
         """
         try:
-            # Search for relevant context in vector store
             search_results = self.context_store.similarity_search(
                 f"chapter_{chapter_number}",
                 k=1
@@ -391,12 +402,34 @@ class BookWritingOrchestrator:
             st.error(f"Context retrieval failed: {e}")
             return ""
 
+    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 main():
     """
     Demonstration of BookWritingOrchestrator
     """
     try:
-        # Create orchestrator
         orchestrator = BookWritingOrchestrator()
         
         # Generate book concept