Update app.py

app.py

@@ -5,6 +5,7 @@ from utils.case_manager import CaseManager
 from utils.legal_notebook_interface import LegalNotebookInterface
 from datetime import datetime
 import os
+from pathlib import Path
 import nltk
 import spacy
 
@@ -16,53 +17,96 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
+def init_huggingface_directories():
+    """Initialize directory structure for Hugging Face Spaces."""
+    # Get the HF space data directory
+    if os.environ.get('SPACE_ID'):
+        # We're in a Hugging Face Space
+        base_dir = Path("/data")
+    else:
+        # Local development
+        base_dir = Path(os.getcwd()) / "data"
+    
+    # Create required directories
+    directories = {
+        'data': base_dir,
+        'cases': base_dir / "cases",
+        'vectors': base_dir / "vectors",
+        'nltk_data': base_dir / "nltk_data",
+        'temp': base_dir / "temp",
+        'logs': base_dir / "logs"
+    }
+    
+    # Create directories if they don't exist
+    for dir_name, dir_path in directories.items():
+        dir_path.mkdir(parents=True, exist_ok=True)
+        st.sidebar.write(f"✓ {dir_name} directory: {dir_path}")
+        
+    return directories
+
 # Initialize NLTK data directory
 @st.cache_resource
-def initialize_nltk():
+def initialize_nltk(nltk_data_dir):
     """Initialize NLTK with all required resources."""
-    nltk_data_dir = os.path.join(os.getcwd(), "data", "nltk_data")
-    os.makedirs(nltk_data_dir, exist_ok=True)
-    nltk.data.path.append(nltk_data_dir)
+    nltk.data.path.append(str(nltk_data_dir))
+
+    # List of all required NLTK resources
+    resources = [
+        'punkt',
+        'averaged_perceptron_tagger',
+        'maxent_ne_chunker',
+        'words',
+        'stopwords'
+    ]
 
-    resources = ['punkt', 'averaged_perceptron_tagger', 'maxent_ne_chunker', 'words', 'stopwords']
+    # Download each resource
     for resource in resources:
         try:
-            nltk.download(resource, download_dir=nltk_data_dir, quiet=True)
+            nltk.download(resource, download_dir=str(nltk_data_dir), quiet=True)
+            st.sidebar.write(f"✓ NLTK {resource}")
         except Exception as e:
             st.error(f"Error downloading NLTK resource {resource}: {str(e)}")
 
     return True
 
-# Initialize components with better error handling
+# Initialize components
 @st.cache_resource
 def init_components():
-    """Initialize all components with proper error handling and feedback."""
+    """Initialize all components with proper HF Spaces directory structure."""
     try:
-        # Set up data directory
-        data_dir = os.path.join(os.getcwd(), "data")
-        os.makedirs(data_dir, exist_ok=True)
-
+        st.sidebar.markdown("### 🚀 Initializing Components")
+        
+        # Initialize directory structure
+        directories = init_huggingface_directories()
+        
         # Initialize NLTK with status
-        with st.spinner("Initializing NLP components..."):
-            initialize_nltk()
-            
-            # Initialize spaCy with proper error handling
+        with st.sidebar.expander("📚 NLTK Resources", expanded=False):
+            initialize_nltk(directories['nltk_data'])
+
+        # Initialize spaCy with proper error handling
+        with st.sidebar.expander("🔧 SpaCy Model", expanded=False):
             try:
                 nlp = spacy.load("en_core_web_sm")
+                st.write("✓ SpaCy model loaded")
             except OSError:
-                with st.spinner("Downloading required language model..."):
-                    os.system("python -m spacy download en_core_web_sm")
-                    nlp = spacy.load("en_core_web_sm")
+                st.write("Downloading SpaCy model...")
+                os.system("python -m spacy download en_core_web_sm")
+                nlp = spacy.load("en_core_web_sm")
+                st.write("✓ SpaCy model downloaded and loaded")
 
-        # Initialize components with status updates
-        case_manager = CaseManager(base_path=os.path.join(data_dir, "cases"))
-        vector_store = VectorStore(storage_path=os.path.join(data_dir, "vectors"))
-        doc_processor = DocumentProcessor(base_path=data_dir)
+        # Initialize components with updated paths
+        case_manager = CaseManager(base_path=str(directories['cases']))
+        vector_store = VectorStore(storage_path=str(directories['vectors']))
+        doc_processor = DocumentProcessor(base_path=str(directories['data']))
 
+        # Log initialization status
+        st.sidebar.success("✅ All components initialized successfully")
+        
         return case_manager, vector_store, doc_processor
 
     except Exception as e:
         st.error(f"Error initializing components: {str(e)}")
+        st.exception(e)
         raise
 
 def main():
@@ -98,6 +142,26 @@ def main():
                 background-color: #fff3cd;
                 color: #664d03;
             }
+            .chat-message {
+                padding: 1.5rem;
+                border-radius: 0.5rem;
+                margin-bottom: 1rem;
+            }
+            .user-message {
+                background-color: #f0f2f6;
+            }
+            .assistant-message {
+                background-color: #ffffff;
+                border: 1px solid #e0e0e0;
+            }
+            .source-reference {
+                margin-top: 0.5rem;
+                padding: 0.5rem;
+                background-color: #f8f9fa;
+                border: 1px solid #e9ecef;
+                border-radius: 0.25rem;
+                font-size: 0.875rem;
+            }
             </style>
             """, unsafe_allow_html=True)
 
@@ -106,8 +170,7 @@ def main():
 
         # Initialize components with error handling
         try:
-            with st.spinner("Initializing application components..."):
-                case_manager, vector_store, doc_processor = init_components()
+            case_manager, vector_store, doc_processor = init_components()
         except Exception as e:
             st.error("Failed to initialize application components. Please try again.")
             st.exception(e)
@@ -184,8 +247,22 @@ def main():
                                     text, chunks, metadata = doc_processor.process_and_tag_document(uploaded_file)
                                     progress_bar.progress(50)
                                     
+                                    # Store in vector store
+                                    status_text.text("Creating document embeddings...")
+                                    for chunk in chunks:
+                                        vector_store.add_document(
+                                            doc_id=metadata['doc_id'],
+                                            text=chunk['text'],
+                                            metadata={
+                                                **metadata,
+                                                'chunk_index': chunk['index'],
+                                                'total_chunks': len(chunks)
+                                            }
+                                        )
+                                    progress_bar.progress(75)
+                                    
                                     # Prepare and add document
-                                    status_text.text("Adding to case...")
+                                    status_text.text("Saving to case...")
                                     doc_data = {
                                         "id": metadata['doc_id'],
                                         "title": uploaded_file.name,
@@ -193,7 +270,6 @@ def main():
                                         "metadata": metadata,
                                         "chunks": chunks
                                     }
-                                    progress_bar.progress(75)
                                     
                                     case_manager.add_document(case["id"], doc_data)
                                     progress_bar.progress(100)
@@ -208,6 +284,7 @@ def main():
                                     
                                 except Exception as e:
                                     st.error(f"Error processing document: {str(e)}")
+                                    st.exception(e)
                                 finally:
                                     # Clean up progress indicators
                                     progress_bar.empty()
@@ -216,6 +293,7 @@ def main():
                     st.info("No cases created yet. Use the form above to create your first case.")
             except Exception as e:
                 st.error(f"Error loading cases: {str(e)}")
+                st.exception(e)
 
         elif tab == "📄 Document Analysis":
             st.title("📄 Document Analysis")