Update app.py

app.py

@@ -2,17 +2,10 @@ import streamlit as st
 from utils.vector_store import VectorStore
 from utils.document_processor import DocumentProcessor
 from utils.case_manager import CaseManager
-from utils.legal_prompt_generator import LegalPromptGenerator  # Add this
-from components.chat_interface import ChatInterface
-from components.template_generator import render_template_generator
-
-# Initialize components
-case_manager = CaseManager()
-doc_processor = DocumentProcessor()
-vector_store = VectorStore()
-prompt_generator = LegalPromptGenerator()  # Add this
-chat_interface = ChatInterface(case_manager, vector_store, doc_processor)  # Pass components here
-
+from utils.legal_notebook_interface import LegalNotebookInterface
+from datetime import datetime
+import os
+import nltk
 
 # Page configuration
 st.set_page_config(
@@ -22,9 +15,40 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
+# Initialize NLTK data directory
+@st.cache_resource
+def initialize_nltk():
+    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)
+    
+    resources = ['punkt', 'stopwords']
+    for resource in resources:
+        try:
+            nltk.download(resource, download_dir=nltk_data_dir, quiet=True)
+        except Exception as e:
+            st.error(f"Error downloading NLTK resource {resource}: {str(e)}")
+    return True
+
+# Initialize components
+@st.cache_resource
+def init_components():
+    # Initialize NLTK first
+    initialize_nltk()
+    
+    # Set up data directory
+    data_dir = os.path.join(os.getcwd(), "data")
+    os.makedirs(data_dir, exist_ok=True)
+    
+    # Initialize components
+    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)
+    
+    return case_manager, vector_store, doc_processor
 
 def main():
-    # Header (placed at the start of the app)
+    # Header
     st.markdown(
         """
         <style>
@@ -47,6 +71,13 @@ def main():
             font-size: 14px;
             color: #A9A9A9;
         }
+        .stButton>button {
+            background-color: #2B547E;
+            color: white;
+        }
+        .stTextInput>div>div>input {
+            background-color: #f0f2f6;
+        }
         </style>
         <h1 class="main-header">SuoMoto.AI</h1>
         <p class="main-tagline">Empowering Legal Intelligence: Automate, Analyze, Act.</p>
@@ -54,29 +85,41 @@ def main():
         unsafe_allow_html=True
     )
 
+    # Initialize components with proper error handling
+    try:
+        case_manager, vector_store, doc_processor = init_components()
+    except Exception as e:
+        st.error(f"Error initializing components: {str(e)}")
+        st.stop()
+
     # Sidebar navigation
     st.sidebar.title("SuoMoto.AI ⚖️")
     st.sidebar.markdown(
         """
-        Streamline legal workflows with advanced AI-powered tools for document management, chat, and template generation.
+        Streamline legal workflows with advanced AI-powered tools for document 
+        analysis and management.
         """
     )
+    
     tab = st.sidebar.radio(
         "Navigation",
-        ["📁 Case Manager", "📄 Document Manager", "🤖 Chat", "📝 Generate Templates"]
+        ["📁 Case Manager", "📄 Document Analysis", "🤖 Legal Assistant"]
     )
 
     # Tab 1: Case Manager
     if tab == "📁 Case Manager":
         st.title("📁 Manage Cases")
-        st.markdown("Easily create and manage legal cases with linked documents.")
+        st.markdown("Create and manage legal cases with linked documents.")
 
         # Case creation form
         with st.expander("➕ Create New Case", expanded=True):
             with st.form("create_case_form"):
                 title = st.text_input("Case Title")
                 description = st.text_area("Description", placeholder="Enter case details...")
-                case_type = st.selectbox("Case Type", ["Judgement", "Contract", "MOU", "Will", "Other"])
+                case_type = st.selectbox(
+                    "Case Type", 
+                    ["Judgment", "Contract", "MOU", "Will", "Other"]
+                )
                 create_case = st.form_submit_button("Create Case")
 
                 if create_case and title:
@@ -104,139 +147,81 @@ def main():
 
                     # Upload new document
                     uploaded_file = st.file_uploader(
-                        f"Upload Document to Case: {case['title']}", key=f"upload_{case['id']}"
+                        f"Upload Document to Case: {case['title']}", 
+                        key=f"upload_{case['id']}"
                     )
                     if uploaded_file:
                         with st.spinner("Processing document..."):
                             text, chunks, metadata = doc_processor.process_and_tag_document(uploaded_file)
-                            doc_data = {"title": uploaded_file.name, "text": text, "metadata": metadata}
+                            doc_data = {
+                                "title": uploaded_file.name,
+                                "text": text,
+                                "metadata": metadata
+                            }
                             case_manager.add_document(case["id"], doc_data)
                             st.success(f"Document '{uploaded_file.name}' added to case!")
 
-    # Tab 2: Document Manager
-    elif tab == "📄 Document Manager":
-        st.title("📄 Manage Documents")
-        st.markdown("Search and manage documents across cases.")
-
-        # Search functionality
-        search_query = st.text_input("🔍 Search for a case or document:")
-        if search_query:
-            results = case_manager.search(search_query)
-            st.subheader("Search Results")
-            if results:
-                for result in results:
-                    with st.expander(f"**{result['title']}**"):
-                        st.write(f"**Type**: {result.get('metadata', {}).get('type', 'Unknown')}")
-                        st.write(f"**Jurisdiction**: {result.get('metadata', {}).get('jurisdiction', 'Unknown')}")
-                        st.write(f"**Uploaded On**: {result.get('metadata', {}).get('uploaded_at', 'Unknown')}")
-                        st.write(f"**Preview**: {result.get('text', '')[:300]}...")
-            else:
-                st.warning("No results found for your search.")
-        else:
-            st.subheader("📚 All Documents")
-            cases = case_manager.get_all_cases()
-            if cases:
-                for case in cases:
-                    st.markdown(f"### Case: {case['title']}")
-                    for doc in case_manager.list_documents(case["id"]):
-                        with st.expander(f"Document: {doc.get('title', 'Unknown Title')}"):
-                            st.markdown(
-                                f"""
-                                **Type**: {doc.get('metadata', {}).get('type', 'Unknown')}  
-                                **Jurisdiction**: {doc.get('metadata', {}).get('jurisdiction', 'Unknown')}  
-                                **Uploaded On**: {doc.get('metadata', {}).get('uploaded_at', 'Unknown')}  
-                                """
-                            )
-                            st.markdown(f"**Content Preview**: {doc.get('text', '')[:500]}...")
-            else:
-                st.info("No cases available.")
-
-    # Tab 3: Chat
-    elif tab == "🤖 Chat":
-        st.markdown(
-            """
-            <style>
-            .chat-container {
-                background-color: #f9f9f9;
-                border-radius: 10px;
-                padding: 20px;
-                max-height: 600px;
-                overflow-y: auto;
-            }
-            .user-message {
-                text-align: right;
-                color: #ffffff;
-                background-color: #2B547E;
-                padding: 10px;
-                border-radius: 10px;
-                margin-bottom: 10px;
-                max-width: 70%;
-                margin-left: auto;
-            }
-            .assistant-message {
-                text-align: left;
-                color: #000000;
-                background-color: #e3e3e3;
-                padding: 10px;
-                border-radius: 10px;
-                margin-bottom: 10px;
-                max-width: 70%;
-                margin-right: auto;
-            }
-            </style>
-            """,
-            unsafe_allow_html=True,
-        )
+    # Tab 2: Document Analysis
+    elif tab == "📄 Document Analysis":
+        st.title("📄 Document Analysis")
+        # Initialize notebook interface
+        notebook = LegalNotebookInterface(case_manager, vector_store, doc_processor)
+        # Render the notebook interface
+        notebook.render()
 
-        st.title("🤖 AI-Powered Chat")
-        st.markdown("Chat with cases or documents to get insights.")
+    # Tab 3: Legal Assistant
+    elif tab == "🤖 Legal Assistant":
+        st.title("🤖 AI Legal Assistant")
+        st.markdown("Chat with your AI legal assistant about cases and documents.")
 
-        selected_case = st.selectbox("Select Case", options=[case["id"] for case in case_manager.get_all_cases()])
-        selected_document = st.selectbox(
-            "Select Document", options=case_manager.list_documents(selected_case), format_func=lambda x: x["title"]
+        # Select case and documents
+        cases = case_manager.get_all_cases()
+        if not cases:
+            st.warning("Please create a case and add documents first.")
+            return
+
+        selected_case = st.selectbox(
+            "Select Case",
+            cases,
+            format_func=lambda x: x['title']
         )
 
-        if selected_document:
-            st.write(f"**Chatting about Document:** {selected_document['title']}")
-            chat_history = st.session_state.get("chat_history", [])
+        if selected_case:
+            documents = case_manager.list_documents(selected_case['id'])
+            if not documents:
+                st.warning("Please add documents to the case first.")
+                return
 
-            with st.container():
-                st.markdown('<div class="chat-container">', unsafe_allow_html=True)
-                for message in chat_history:
-                    role_class = "user-message" if message["role"] == "user" else "assistant-message"
-                    st.markdown(f'<div class="{role_class}">{message["content"]}</div>', unsafe_allow_html=True)
-                st.markdown('</div>', unsafe_allow_html=True)
+            # Chat interface
+            if "messages" not in st.session_state:
+                st.session_state.messages = []
 
-            # Chat input
-            query = st.text_input("💬 Ask a question about the document:")
-            if query:
-                chat_history.append({"role": "user", "content": query})
-                st.session_state["chat_history"] = chat_history
-
-                with st.spinner("Generating response..."):
-                    context = selected_document.get('text', '')
-                    if not context:
-                        st.error("The selected document has no text content.")
-                    else:
-                        response = vector_store.chat_with_context(query, context)
-                        chat_history.append({"role": "assistant", "content": response})
-                        st.session_state["chat_history"] = chat_history
-
-    # Tab 4: Generate Templates
-    elif tab == "📝 Generate Templates":
-        st.title("📝 Legal Document Templates")
-        st.markdown("Generate standard legal documents using predefined templates.")
-        render_template_generator()
+            # Display chat history
+            for message in st.session_state.messages:
+                with st.chat_message(message["role"]):
+                    st.markdown(message["content"])
 
+            # Chat input
+            if prompt := st.chat_input("Ask about your legal documents..."):
+                st.session_state.messages.append({"role": "user", "content": prompt})
+                with st.chat_message("user"):
+                    st.markdown(prompt)
+
+                with st.chat_message("assistant"):
+                    response = f"Analysis of your query: {prompt}\n\n"
+                    response += "Based on the documents in this case..."
+                    st.markdown(response)
+                    st.session_state.messages.append({"role": "assistant", "content": response})
+
+    # Footer
     st.markdown(
         """
         <div class="footer">
-        Synaptyx: SuoMoto.AI agent
+        Powered by SuoMoto.AI
         </div>
         """,
         unsafe_allow_html=True
     )
 
-
 if __name__ == "__main__":
-    main()
+    main()