Update app.py

app.py

@@ -1,8 +1,9 @@
-# app.py
 import streamlit as st
 import os
 from pathlib import Path
 from components.template_generator import render_template_generator
+from components.document_viewer import DocumentViewer
+from utils.case_manager import CaseManager
 from utils.document_processor import DocumentProcessor
 import anthropic
 
@@ -16,31 +17,28 @@ st.set_page_config(
 
 # Initialize Anthropic client
 try:
-    # First try to get from HF secrets
-    api_key = os.environ.get('ANTHROPIC_API_KEY')
+    api_key = os.getenv("ANTHROPIC_API_KEY")
     if not api_key:
-        st.error("Please set the ANTHROPIC_API_KEY in your Space's secrets")
+        st.error("Please set the ANTHROPIC_API_KEY in your environment variables.")
         st.stop()
-    
+
     client = anthropic.Anthropic(api_key=api_key)
 except Exception as e:
-    st.error(f"Error initializing Anthropic client: Please ensure Anthropic_API_KEY is set in HF secrets")
+    st.error(f"Error initializing Anthropic client: {str(e)}")
     st.stop()
 
-# Session state initialization
-if 'processed_docs' not in st.session_state:
+# Initialize session state
+if "processed_docs" not in st.session_state:
     st.session_state.processed_docs = {}
-if 'chat_history' not in st.session_state:
+if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
-def initialize_case_folders():
-    """Initialize folder structure for cases"""
-    base_path = Path("data/cases")
-    base_path.mkdir(parents=True, exist_ok=True)
-    return base_path
+# Initialize CaseManager
+case_manager = CaseManager()
+
 
 def analyze_document(text: str) -> str:
-    """Analyze document using Claude"""
+    """Analyze document using Claude."""
     try:
         message = client.messages.create(
             model="claude-3-sonnet-20240229",
@@ -63,8 +61,9 @@ def analyze_document(text: str) -> str:
         st.error(f"Error during analysis: {str(e)}")
         return ""
 
+
 def chat_with_docs(query: str, context: str) -> str:
-    """Chat with documents using Claude"""
+    """Chat with documents using Claude."""
     try:
         message = client.messages.create(
             model="claude-3-sonnet-20240229",
@@ -73,9 +72,9 @@ def chat_with_docs(query: str, context: str) -> str:
             messages=[{
                 "role": "user",
                 "content": f"""Based on the following context, please answer the question.
-                
+
                 Context: {context}
-                
+
                 Question: {query}"""
             }]
         )
@@ -84,104 +83,97 @@ def chat_with_docs(query: str, context: str) -> str:
         st.error(f"Error generating response: {str(e)}")
         return ""
 
+
 def render_sidebar():
-    """Render sidebar with navigation options"""
+    """Render sidebar with navigation options."""
     with st.sidebar:
         st.title("Legal AI Assistant")
-        
+
         # Navigation
         selected_option = st.radio(
             "Navigation",
-            ["Document Analysis", "Document Generation", "Chat with Documents"]
+            ["📁 Manage Cases", "📝 Document Generation", "🤖 Chat with Documents"]
         )
-        
         return selected_option
 
+
+def manage_cases():
+    """Case management interface."""
+    st.header("Manage Cases")
+    st.write("Create, view, and manage cases and their associated documents.")
+
+    # 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", ["Contract", "Litigation", "Employment", "Other"])
+            create_case = st.form_submit_button("Create Case")
+
+            if create_case and title:
+                case_id = case_manager.create_case(title, description, case_type)
+                st.success(f"Case '{title}' created successfully!")
+
+    # List 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']}")
+
+                # List documents in case
+                documents = case_manager.list_documents(case["id"])
+                for doc in documents:
+                    st.markdown(f"- **{doc['title']}** (Added: {doc['added_at']})")
+
+                # Upload new documents
+                uploaded_file = st.file_uploader("Upload Document", key=f"upload_{case['id']}")
+                if uploaded_file:
+                    doc_data = {"title": uploaded_file.name, "content": uploaded_file.read().decode("utf-8")}
+                    case_manager.add_document(case["id"], doc_data)
+                    st.success(f"Document '{uploaded_file.name}' added to case.")
+
+
 def main():
-    # Initialize folders
-    base_path = initialize_case_folders()
-    
-    # Initialize components
-    doc_processor = DocumentProcessor()
-    
     # Render sidebar and get selected option
     selected_option = render_sidebar()
-    
-    if selected_option == "Document Analysis":
-        st.header("Document Analysis")
-        
-        # File upload
-        uploaded_file = st.file_uploader(
-            "Upload a legal document",
-            type=['pdf', 'docx', 'txt']
-        )
-        
-        if uploaded_file:
-            with st.spinner("Processing document..."):
-                # Process document
-                text, chunks = doc_processor.process_document(uploaded_file)
-                
-                if text:
-                    st.success("Document processed successfully!")
-                    
-                    # Store in session state
-                    st.session_state.processed_docs[uploaded_file.name] = {
-                        'text': text,
-                        'chunks': chunks
-                    }
-                    
-                    # Show document content
-                    with st.expander("Document Content", expanded=False):
-                        st.text(text[:1000] + "..." if len(text) > 1000 else text)
-                    
-                    # Analysis button
-                    if st.button("Analyze Document"):
-                        with st.spinner("Analyzing..."):
-                            analysis = analyze_document(text)
-                            if analysis:
-                                st.markdown("### Analysis Results")
-                                st.markdown(analysis)
-                                
-                                # Download button
-                                st.download_button(
-                                    "Download Analysis",
-                                    analysis,
-                                    file_name=f"analysis_{uploaded_file.name}.txt",
-                                    mime="text/plain"
-                                )
-    
-    elif selected_option == "Document Generation":
+
+    if selected_option == "📁 Manage Cases":
+        manage_cases()
+    elif selected_option == "📝 Document Generation":
         # Render template generator
         render_template_generator()
-    
-    else:  # Chat with Documents
+    elif selected_option == "🤖 Chat with Documents":
         st.header("Chat with Documents")
-        
+
         if not st.session_state.processed_docs:
-            st.warning("Please upload some documents in the Document Analysis section first.")
+            st.warning("Please upload and process some documents in the Manage Cases section first.")
             return
-        
+
         # Chat interface
         query = st.text_input("Ask a question about your documents:")
         if query:
             # Combine all document texts as context
             context = "\n\n".join([doc['text'] for doc in st.session_state.processed_docs.values()])
-            
+
             with st.spinner("Thinking..."):
                 response = chat_with_docs(query, context)
-                
+
                 # Add to chat history
                 st.session_state.chat_history.append({
                     "query": query,
                     "response": response
                 })
-        
+
         # Display chat history
         for chat in st.session_state.chat_history:
             with st.container():
-                st.markdown("**You:** " + chat["query"])
-                st.markdown("**Assistant:** " + chat["response"])
+                st.markdown(f"**You:** {chat['query']}")
+                st.markdown(f"**Assistant:** {chat['response']}")
                 st.markdown("---")
 
+
 if __name__ == "__main__":
-    main()
+    main()