Update app.py

app.py

@@ -1,13 +1,10 @@
 # app.py
 import streamlit as st
-from components.sidebar import render_sidebar
-from components.document_viewer import DocumentViewer
-from components.chat_interface import ChatInterface
+import os
+from pathlib import Path
 from components.template_generator import render_template_generator
-from utils.case_manager import CaseManager
 from utils.document_processor import DocumentProcessor
-from utils.vector_store import VectorStore
-import os
+import anthropic
 
 # Page config
 st.set_page_config(
@@ -17,86 +14,174 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# Load custom CSS
-with open('static/styles.css') as f:
-    st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
+# Initialize Anthropic client
+try:
+    # First try to get from HF secrets
+    api_key = os.environ.get('Anthropic_API_KEY')
+    if not api_key:
+        st.error("Please set the Anthropic_API_KEY in your Space's secrets")
+        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.stop()
+
+# Session state initialization
+if 'processed_docs' not in st.session_state:
+    st.session_state.processed_docs = {}
+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
+
+def analyze_document(text: str) -> str:
+    """Analyze document using Claude"""
+    try:
+        message = client.messages.create(
+            model="claude-3-sonnet-20240229",
+            max_tokens=2000,
+            temperature=0.7,
+            messages=[{
+                "role": "user",
+                "content": f"""Analyze this legal document and provide:
+                1. Key points and findings
+                2. Legal implications
+                3. Important dates and deadlines
+                4. Potential risks
+                5. Recommended actions
+
+                Document text: {text}"""
+            }]
+        )
+        return message.content[0].text
+    except Exception as e:
+        st.error(f"Error during analysis: {str(e)}")
+        return ""
+
+def chat_with_docs(query: str, context: str) -> str:
+    """Chat with documents using Claude"""
+    try:
+        message = client.messages.create(
+            model="claude-3-sonnet-20240229",
+            max_tokens=2000,
+            temperature=0.7,
+            messages=[{
+                "role": "user",
+                "content": f"""Based on the following context, please answer the question.
+                
+                Context: {context}
+                
+                Question: {query}"""
+            }]
+        )
+        return message.content[0].text
+    except Exception as e:
+        st.error(f"Error generating response: {str(e)}")
+        return ""
+
+def render_sidebar():
+    """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"]
+        )
+        
+        return selected_option
 
 def main():
+    # Initialize folders
+    base_path = initialize_case_folders()
+    
     # Initialize components
-    case_manager = CaseManager()
     doc_processor = DocumentProcessor()
     
-    # Render sidebar and get selected case
-    selected_case = render_sidebar(case_manager)
+    # Render sidebar and get selected option
+    selected_option = render_sidebar()
     
-    if selected_case:
-        # Initialize vector store for selected case
-        vector_store = VectorStore(selected_case)
+    if selected_option == "Document Analysis":
+        st.header("Document Analysis")
         
-        # Get case details
-        case = case_manager.get_case(selected_case)
+        # File upload
+        uploaded_file = st.file_uploader(
+            "Upload a legal document",
+            type=['pdf', 'docx', 'txt']
+        )
         
-        # Main area
-        st.title(case['title'])
-        st.markdown(f"**Type:** {case['case_type']} | **Created:** {case['created_at'][:10]}")
+        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":
+        # Render template generator
+        render_template_generator()
+    
+    else:  # Chat with Documents
+        st.header("Chat with Documents")
         
-        # Tabs for different functions
-        tab1, tab2, tab3 = st.tabs(["📚 Documents", "💭 Chat", "📝 Generate"])
+        if not st.session_state.processed_docs:
+            st.warning("Please upload some documents in the Document Analysis section first.")
+            return
         
-        with tab1:
-            # Document upload
-            uploaded_files = st.file_uploader(
-                "Upload Documents",
-                accept_multiple_files=True,
-                type=['pdf', 'docx', 'txt']
-            )
+        # 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()])
             
-            if uploaded_files:
-                for file in uploaded_files:
-                    with st.spinner(f"Processing {file.name}..."):
-                        # Process document
-                        text, chunks = doc_processor.process_document(file)
-                        
-                        # Add to vector store
-                        vector_store.add_texts(
-                            [chunk["text"] for chunk in chunks],
-                            [chunk["metadata"] for chunk in chunks]
-                        )
-                        
-                        # Add to case
-                        doc_data = {
-                            "title": file.name,
-                            "content": text,
-                            "chunks": chunks
-                        }
-                        case_manager.add_document(selected_case, doc_data)
-                        
-                        st.success(f"✅ {file.name} processed successfully!")
-            
-            # Document viewer
-            doc_viewer = DocumentViewer(case_manager)
-            doc_viewer.render(selected_case)
-        
-        with tab2:
-            # Chat interface
-            chat = ChatInterface(vector_store)
-            chat.render()
-        
-        with tab3:
-            # Document generation
-            render_template_generator()
-    
-    else:
-        # No case selected
-        st.markdown("## Welcome to Legal AI Assistant")
-        st.markdown("""
-        This tool helps you:
-        - Manage legal cases and documents
-        - Chat with your documents
-        - Generate legal documents from templates
+            with st.spinner("Thinking..."):
+                response = chat_with_docs(query, context)
+                
+                # Add to chat history
+                st.session_state.chat_history.append({
+                    "query": query,
+                    "response": response
+                })
         
-        Get started by creating a new case in the sidebar!
-        """)
+        # 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("---")
 
 if __name__ == "__main__":
     main()