Update app.py

app.py

@@ -1,11 +1,13 @@
 import streamlit as st
 from utils.vector_store import VectorStore
 from utils.document_processor import DocumentProcessor
+from utils.case_manager import CaseManager
 from components.template_generator import render_template_generator
 
 # Initialize components
-vector_store = None  # Placeholder for lazy initialization
+case_manager = CaseManager()
 doc_processor = DocumentProcessor()
+vector_store = VectorStore()
 
 # Page configuration
 st.set_page_config(
@@ -15,86 +17,104 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-def get_vector_store():
-    """Lazy initialization of VectorStore to avoid circular imports."""
-    global vector_store
-    if vector_store is None:
-        vector_store = VectorStore()
-    return vector_store
-
 def main():
-    """Main function to handle app navigation and functionality."""
     # Sidebar navigation
     tab = st.sidebar.radio(
         "Navigation",
-        ["📁 Manage Documents", "📝 Generate Templates", "🔍 Search Documents"]
+        ["📁 Case Manager", "📄 Document Manager", "🤖 Chat", "📝 Generate Templates"]
     )
 
-    # Tab 1: Manage Documents
-    if tab == "📁 Manage Documents":
-        st.title("📁 Manage Documents")
-        uploaded_file = st.file_uploader("Upload Document", type=["pdf", "docx", "txt"])
+    # Tab 1: Case Manager
+    if tab == "📁 Case Manager":
+        st.title("📁 Case Manager")
+        
+        # Case creation form
+        with st.expander("Create New Case"):
+            with st.form("create_case_form"):
+                title = st.text_input("Case Title")
+                description = st.text_area("Description")
+                case_type = st.selectbox("Case Type", ["Judgement", "Contract", "MOU", "Will", "Other"])
+                create_case = st.form_submit_button("Create Case")
 
-        if uploaded_file:
-            try:
-                with st.spinner("Processing document..."):
-                    # Process document and extract text/chunks
-                    text, chunks = doc_processor.process_document(uploaded_file)
-                    st.success("Document processed successfully!")
+                if create_case and title:
+                    case_id = case_manager.create_case(title, description, case_type)
+                    st.success(f"Case '{title}' created successfully!")
 
-                    # Add to vector store
-                    vector_store_instance = get_vector_store()
-                    vector_store_instance.add_texts(
-                        texts=[chunk["text"] for chunk in chunks],
-                        metadatas=[{
-                            "text": chunk["text"],
-                            "chunk_id": chunk["chunk_id"],
-                            "filename": uploaded_file.name
-                        } for chunk in chunks]
-                    )
-                    st.success("Document added to vector store!")
+        # List all cases
+        st.subheader("All Cases")
+        cases = case_manager.get_all_cases()
+        if cases:
+            for case in cases:
+                with st.expander(f"Case: {case['title']}"):
+                    st.write(f"**Description**: {case['description']}")
+                    st.write(f"**Type**: {case['case_type']}")
+                    st.write(f"**Created At**: {case['created_at']}")
 
-                    # Display document preview
-                    with st.expander("Document Preview", expanded=False):
-                        st.text_area(
-                            "Content",
-                            value=text[:1000] + "..." if len(text) > 1000 else text,
-                            height=300,
-                            disabled=True
-                        )
+                    # List documents in case
+                    documents = case_manager.list_documents(case["id"])
+                    for doc in documents:
+                        st.markdown(f"- **{doc['title']}** (Added: {doc['added_at']})")
 
-            except Exception as e:
-                st.error(f"Error processing document: {e}")
+                    # Upload new document
+                    uploaded_file = st.file_uploader("Upload Document", key=f"upload_{case['id']}")
+                    if uploaded_file:
+                        with st.spinner("Processing document..."):
+                            text, metadata = doc_processor.process_and_tag_document(uploaded_file)
+                            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!")
 
-        # List processed documents
-        st.subheader("Processed Documents")
-        vector_store_instance = get_vector_store()
-        processed_docs = vector_store_instance.metadata
-        if processed_docs:
-            for idx, doc in enumerate(processed_docs):
-                st.markdown(f"{idx+1}. **{doc.get('filename', 'Unknown')}** - Chunk ID: {doc['chunk_id']}")
+    # Tab 2: Document Manager
+    elif tab == "📄 Document Manager":
+        st.title("📄 Document Manager")
+        
+        # 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")
+            for result in results:
+                st.write(f"- **{result['title']}** ({result['type']})")
         else:
-            st.info("No documents uploaded yet.")
+            st.subheader("All Documents")
+            cases = case_manager.get_all_cases()
+            for case in cases:
+                st.write(f"**Case: {case['title']}**")
+                for doc in case_manager.list_documents(case["id"]):
+                    with st.expander(f"Document: {doc['title']}"):
+                        st.write(f"**Metadata**: {doc.get('metadata', 'No metadata available')}")
+                        st.write(f"**Content Preview**: {doc['text'][:500]}...")
+
+    # Tab 3: Chat
+    elif tab == "🤖 Chat":
+        st.title("🤖 Chat")
+        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))
+
+        if selected_document:
+            st.write(f"**Chatting about Document:** {selected_document['title']}")
 
-    # Tab 2: Generate Templates
+            # Chat input and history
+            query = st.text_input("Ask a question about the document:")
+            if query:
+                with st.spinner("Generating response..."):
+                    context = selected_document['text']
+                    response = vector_store.chat_with_context(query, context)
+                    st.write(f"**Response**: {response}")
+
+                # Allow adding new documents mid-chat
+                uploaded_file = st.file_uploader("Upload Document During Chat", key="chat_upload")
+                if uploaded_file:
+                    with st.spinner("Processing document..."):
+                        text, metadata = doc_processor.process_and_tag_document(uploaded_file)
+                        doc_data = {"title": uploaded_file.name, "text": text, "metadata": metadata}
+                        case_manager.add_document(selected_case, doc_data)
+                        st.success(f"Document '{uploaded_file.name}' added to case!")
+
+    # Tab 4: Generate Templates
     elif tab == "📝 Generate Templates":
         st.title("📝 Generate Templates")
         render_template_generator()
 
-    # Tab 3: Search Documents
-    elif tab == "🔍 Search Documents":
-        st.title("🔍 Search Documents")
-        query = st.text_input("Enter your query:")
-        if query:
-            with st.spinner("Searching for relevant chunks..."):
-                vector_store_instance = get_vector_store()
-                results = vector_store_instance.similarity_search(query)
-                if results:
-                    st.success("Found relevant chunks:")
-                    for result in results:
-                        st.markdown(f"- **Chunk:** {result['text'][:200]}... (Relevance: {result['distance']:.2f})")
-                else:
-                    st.warning("No relevant chunks found.")
-
 if __name__ == "__main__":
     main()