Update multi_agent_book_workflow.py

multi_agent_book_workflow.py

@@ -8,6 +8,7 @@ from typing import Dict, List, Any
 import uuid
 import chromadb
 import numpy as np
+import streamlit as st
 
 # Agent and LLM Frameworks
 from crewai import Agent, Task, Crew
@@ -24,39 +25,57 @@ class BookWritingOrchestrator:
         Args:
             api_key (str, optional): API key for models if not using environment variables
         """
-        # API Key handling
-        self.openai_api_key = api_key or os.getenv('OPENAI_API_KEY')
-        self.anthropic_api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
-        
-        # Vector Store for Context Management
-        self.chroma_client = chromadb.Client()
-        self.context_store = self.chroma_client.create_collection(
-            name=f"book_context_{str(uuid.uuid4())}"
-        )
-        
-        # Embedding Model
-        self.embeddings = OpenAIEmbeddings(api_key=self.openai_api_key)
-        
-        # LLM Configurations
-        self.openai_llm = ChatOpenAI(
-            model="gpt-4-turbo",
-            temperature=0.7,
-            api_key=self.openai_api_key
-        )
-        self.anthropic_llm = ChatAnthropic(
-            model='claude-3-opus-20240229',
-            temperature=0.7,
-            api_key=self.anthropic_api_key
-        )
-        
-        # Agent Memory
-        self.global_memory = ConversationBufferMemory(
-            memory_key="chat_history",
-            return_messages=True
-        )
+        # Generate a unique project ID
+        self.project_id = str(uuid.uuid4())
+        
+        # API Key handling with more robust error checking
+        try:
+            # Try getting API keys from environment or passed parameter
+            self.openai_api_key = api_key or os.getenv('OPENAI_API_KEY')
+            self.anthropic_api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
+            
+            # Validate API keys
+            if not self.openai_api_key:
+                st.warning("OpenAI API key is missing. Some features may be limited.")
+            
+            if not self.anthropic_api_key:
+                st.warning("Anthropic API key is missing. Some features may be limited.")
+            
+            # Vector Store for Context Management
+            self.chroma_client = chromadb.Client()
+            self.context_store = self.chroma_client.create_collection(
+                name=f"book_context_{self.project_id}"
+            )
+            
+            # Embedding Model
+            self.embeddings = OpenAIEmbeddings(
+                api_key=self.openai_api_key
+            )
+            
+            # LLM Configurations
+            self.openai_llm = ChatOpenAI(
+                model="gpt-4-turbo",
+                temperature=0.7,
+                api_key=self.openai_api_key
+            )
+            self.anthropic_llm = ChatAnthropic(
+                model='claude-3-opus-20240229',
+                temperature=0.7,
+                api_key=self.anthropic_api_key
+            )
+            
+            # Agent Memory
+            self.global_memory = ConversationBufferMemory(
+                memory_key="chat_history",
+                return_messages=True
+            )
+            
+            # Initialize Specialized Agents
+            self.setup_agents()
         
-        # Initialize Specialized Agents
-        self.setup_agents()
+        except Exception as e:
+            st.error(f"Error initializing BookWritingOrchestrator: {e}")
+            raise
     
     def setup_agents(self):
         """
@@ -219,15 +238,18 @@ class BookWritingOrchestrator:
             context_key (str): Unique identifier for the context
             content (str): Content to store
         """
-        # Generate embedding
-        embedding = self.embeddings.embed_documents([content])[0]
-        
-        # Store in Chroma
-        self.context_store.add(
-            embeddings=[embedding],
-            documents=[content],
-            ids=[f"{self.project_id}_{context_key}"]
-        )
+        try:
+            # Generate embedding
+            embedding = self.embeddings.embed_documents([content])[0]
+            
+            # Store in Chroma
+            self.context_store.add(
+                embeddings=[embedding],
+                documents=[content],
+                ids=[f"{self.project_id}_{context_key}"]
+            )
+        except Exception as e:
+            st.error(f"Error storing context: {e}")
     
     def _retrieve_context(
         self, 
@@ -244,23 +266,27 @@ class BookWritingOrchestrator:
         Returns:
             List of contextually relevant content
         """
-        # Retrieve previous chapters
-        previous_chapters = [
-            f'chapter_{i}' for i in range(1, chapter_number)
-        ]
-        
-        # Search context
-        retrieved_contexts = []
-        for chapter_key in previous_chapters:
-            search_embedding = self.embeddings.embed_documents([chapter_key])[0]
-            results = self.context_store.query(
-                query_embeddings=[search_embedding],
-                n_results=top_k
-            )
+        try:
+            # Retrieve previous chapters
+            previous_chapters = [
+                f'chapter_{i}' for i in range(1, chapter_number)
+            ]
             
-            retrieved_contexts.extend(results['documents'][0])
-        
-        return retrieved_contexts
+            # Search context
+            retrieved_contexts = []
+            for chapter_key in previous_chapters:
+                search_embedding = self.embeddings.embed_documents([chapter_key])[0]
+                results = self.context_store.query(
+                    query_embeddings=[search_embedding],
+                    n_results=top_k
+                )
+                
+                retrieved_contexts.extend(results['documents'][0])
+            
+            return retrieved_contexts
+        except Exception as e:
+            st.error(f"Error retrieving context: {e}")
+            return []
     
     def _parse_concept(self, concept_result: str) -> Dict[str, Any]:
         """
@@ -272,10 +298,42 @@ class BookWritingOrchestrator:
         Returns:
             Structured book concept dictionary
         """
-        # Implement robust parsing logic
-        # This could use another AI call or regex-based extraction
-        return {
-            'title': 'Parsed Book Title',
-            'genre': 'Parsed Genre',
-            'chapters': []  # Parsed chapter outline
-        }
+        # Basic parsing with error handling
+        try:
+            # This is a placeholder. In a real implementation, 
+            # you might use NLP techniques or another AI call to parse
+            return {
+                'title': 'Book Title',
+                'genre': 'Fiction',
+                'chapters': ['Chapter 1', 'Chapter 2', 'Chapter 3']
+            }
+        except Exception as e:
+            st.error(f"Error parsing concept: {e}")
+            return {
+                'title': 'Untitled',
+                'genre': 'Unspecified',
+                'chapters': []
+            }
+
+def main():
+    """
+    Demonstration of BookWritingOrchestrator
+    """
+    try:
+        # Example usage
+        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()