Update app.py

app.py

@@ -2,11 +2,13 @@ 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
 from utils.legal_notebook_interface import LegalNotebookInterface
 from datetime import datetime
 import os
 import nltk
 import spacy
+import anthropic
 
 # Page configuration
 st.set_page_config(
@@ -73,6 +75,115 @@ def init_components():
         st.error(f"Error initializing components: {str(e)}")
         raise
 
+class ChatInterface:
+    def __init__(self, case_manager, vector_store, document_processor):
+        self.case_manager = case_manager
+        self.vector_store = vector_store
+        self.document_processor = document_processor
+        self.prompt_generator = LegalPromptGenerator()
+
+        # Initialize Anthropics client
+        try:
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+            if not api_key:
+                st.error("Please set the ANTHROPIC_API_KEY environment variable.")
+                st.stop()
+            self.client = anthropic.Anthropic(api_key=api_key)
+        except Exception as e:
+            st.error(f"Error initializing Anthropic client: {str(e)}")
+            st.stop()
+
+        # Initialize session state
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
+        if "current_case" not in st.session_state:
+            st.session_state.current_case = None
+
+    def render(self):
+        """Render the chat interface."""
+        st.title("🤖 Legal Assistant")
+        st.sidebar.title("📁 Select Case")
+
+        # Sidebar case selection
+        cases = self.case_manager.get_all_cases()
+        if not cases:
+            st.sidebar.info("No cases found. Create a case first.")
+            return
+
+        selected_case = st.sidebar.selectbox(
+            "Choose a case",
+            cases,
+            format_func=lambda x: x['title']
+        )
+        st.session_state.current_case = selected_case['id']
+
+        # Chat interface
+        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..."):
+            self._handle_chat_input(prompt)
+
+    def _handle_chat_input(self, prompt: str):
+        """Handle user input and generate a response."""
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        with st.spinner("Analyzing documents..."):
+            try:
+                # Get case documents
+                documents = self.case_manager.list_documents(st.session_state.current_case)
+                if not documents:
+                    st.error("No documents available for analysis.")
+                    return
+
+                # Search for relevant chunks
+                context_chunks = self.vector_store.similarity_search(
+                    query=prompt,
+                    k=5
+                )
+
+                # Generate response using the LLM
+                response, references = self.generate_response(prompt, context_chunks)
+
+                # Append assistant response to chat
+                st.session_state.messages.append({
+                    "role": "assistant",
+                    "content": f"{response}\n\n{references}"
+                })
+            except Exception as e:
+                st.error(f"Error processing input: {str(e)}")
+
+    def generate_response(self, prompt: str, context_chunks: list) -> tuple[str, str]:
+        """Generate a response using the LLM and `LegalPromptGenerator`."""
+        try:
+            # Generate structured messages
+            messages = self.prompt_generator._generate_messages(prompt, context_chunks)
+
+            # Call the LLM API
+            message = self.client.messages.create(
+                model="claude-3",
+                max_tokens=2000,
+                temperature=0.7,
+                messages=messages
+            )
+
+            # Format references
+            references = self._format_references(context_chunks)
+            return message.content[0].text, references
+
+        except Exception as e:
+            st.error(f"Error generating response: {str(e)}")
+            return "An error occurred while processing your request.", ""
+
+    def _format_references(self, chunks: list) -> str:
+        """Format references for the assistant response."""
+        references = []
+        for i, chunk in enumerate(chunks, 1):
+            references.append(f"**Reference {i}:** {chunk['metadata'].get('title', 'Untitled')}\n- Section: {chunk['text'][:200]}...")
+        return "\n".join(references)
+
 def main():
     # Add custom styles
     st.markdown("""
@@ -102,202 +213,25 @@ def main():
     st.markdown("<h1 class='main-header'>SuoMoto.AI</h1>", unsafe_allow_html=True)
     st.markdown("<p class='main-tagline'>Empowering Legal Intelligence: Automate, Analyze, Act.</p>", unsafe_allow_html=True)
 
-    # Initialize components with error handling
+    # Initialize components
     try:
         case_manager, vector_store, doc_processor = init_components()
     except Exception as e:
         st.error("Failed to initialize application components. Please try again.")
         st.stop()
 
-    # Sidebar navigation
-    st.sidebar.title("SuoMoto.AI ⚖️")
-    st.sidebar.markdown("""
-        Streamline legal workflows with advanced AI-powered tools for document 
-        analysis and management.
-    """)
-    
+    # Navigation
     tab = st.sidebar.radio(
         "Navigation",
         ["📁 Case Manager", "📄 Document Analysis", "🤖 Legal Assistant"]
     )
 
-    # Tab 1: Case Manager
-    if tab == "📁 Case Manager":
-        st.title("📁 Manage Cases")
-        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", 
-                    ["Judgment", "Contract", "MOU", "Will", "Other"]
-                )
-                create_case = st.form_submit_button("Create Case")
-
-                if create_case and title:
-                    try:
-                        case_id = case_manager.create_case(title, description, case_type)
-                        st.success(f"Case '{title}' created successfully!")
-                    except Exception as e:
-                        st.error(f"Error creating case: {str(e)}")
-
-        # List all cases
-        st.subheader("📜 All Cases")
-        try:
-            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"])
-                        if documents:
-                            st.markdown("### Documents")
-                            for doc in documents:
-                                st.markdown(f"- **{doc['title']}** (Added: {doc['added_at']})")
-                        else:
-                            st.info("No documents uploaded yet.")
-
-                        # Upload new document
-                        uploaded_file = st.file_uploader(
-                            f"Upload Document to Case: {case['title']}", 
-                            key=f"upload_{case['id']}"
-                        )
-                        # In the Case Manager tab section, where documents are uploaded
-                        if uploaded_file:
-                            if uploaded_file:
-                                with st.spinner("Processing document..."):
-                                    try:
-                                        # Process document
-                                        text, chunks, metadata = doc_processor.process_and_tag_document(uploaded_file)
-                                        
-                                        # Add each chunk to vector store
-                                        for chunk in chunks:
-                                            try:
-                                                vector_store.add_document(
-                                                    doc_id=f"{metadata['doc_id']}_chunk_{chunk['chunk_id']}",
-                                                    text=chunk['text'],
-                                                    metadata={
-                                                        'doc_id': metadata['doc_id'],
-                                                        'chunk_id': chunk['chunk_id'],
-                                                        'title': uploaded_file.name,
-                                                        'type': metadata.get('document_type', 'unknown')
-                                                    }
-                                                )
-                                            except Exception as e:
-                                                st.warning(f"Error adding chunk to vector store: {str(e)}")
-                                        
-                                        # Add document to case
-                                        doc_data = {
-                                            "id": metadata['doc_id'],
-                                            "title": uploaded_file.name,
-                                            "text": text,
-                                            "metadata": metadata,
-                                            "chunks": chunks
-                                        }
-                                        case_manager.add_document(case["id"], doc_data)
-                                        st.success(f"Document '{uploaded_file.name}' processed and added successfully!")
-                                        
-                                    except Exception as e:
-                                        st.error(f"Error processing document: {str(e)}")
-            else:
-                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)}")
-
-    # Tab 2: Document Analysis
-   # In the Document Analysis tab section
-    elif tab == "📄 Document Analysis":
-        st.title("📄 Document Analysis")
-        try:
-            cases = case_manager.get_all_cases()
-            if not cases:
-                st.info("Please create a case and add documents first.")
-            else:
-                notebook = LegalNotebookInterface(case_manager, vector_store, doc_processor)
-                notebook.render()
-        except Exception as e:
-            st.error(f"Error initializing document analysis: {str(e)}")
-        
-    # 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.")
-
-        try:
-            # 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_case:
-                documents = case_manager.list_documents(selected_case['id'])
-                if not documents:
-                    st.warning("Please add documents to the case first.")
-                    return
-
-                # Initialize chat history
-                if "messages" not in st.session_state:
-                    st.session_state.messages = []
-
-                # 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"):
-                        with st.spinner("Analyzing documents..."):
-                            try:
-                                # Get relevant document chunks
-                                chunks = vector_store.similarity_search(prompt, k=3)
-                                context = "\n".join([chunk['text'] for chunk in chunks])
-                                
-                                # Generate response
-                                response = f"Analysis based on your documents:\n\n"
-                                response += "1. Key findings...\n"
-                                response += "2. Legal implications...\n"
-                                response += "3. Recommendations...\n\n"
-                                response += f"Context length: {len(context)} characters"
-                                
-                                st.markdown(response)
-                                st.session_state.messages.append({
-                                    "role": "assistant",
-                                    "content": response
-                                })
-                            except Exception as e:
-                                st.error(f"Error processing query: {str(e)}")
-        
-        except Exception as e:
-            st.error(f"Error in Legal Assistant: {str(e)}")
-
-    # Footer
-    st.markdown(
-        """
-        <div class="footer">
-        Powered by SuoMoto.AI
-        </div>
-        """,
-        unsafe_allow_html=True
-    )
+    if tab == "🤖 Legal Assistant":
+        chat_interface = ChatInterface(case_manager, vector_store, doc_processor)
+        chat_interface.render()
+    else:
+        # Implement other tabs
+        st.info("Other tabs are under construction.")
 
 if __name__ == "__main__":
-    main()
+    main()