Update utils/legal_notebook_interface.py

utils/legal_notebook_interface.py

@@ -2,13 +2,14 @@ import streamlit as st
 from typing import List, Dict
 import datetime
 
+
 class LegalNotebookInterface:
     def __init__(self, case_manager, vector_store, document_processor):
         """Initialize the notebook interface with required components."""
         self.case_manager = case_manager
         self.vector_store = vector_store
         self.document_processor = document_processor
-        
+
         # Initialize session state
         if "active_documents" not in st.session_state:
             st.session_state.active_documents = []
@@ -20,37 +21,40 @@ class LegalNotebookInterface:
     def render(self):
         """Render the main notebook interface."""
         # Add custom styling
-        st.markdown("""
-        <style>
-        .document-panel {
-            background-color: #f8f9fa;
-            border-radius: 10px;
-            padding: 20px;
-            margin: 10px;
-        }
-        .message {
-            padding: 15px;
-            margin: 10px 0;
-            border-radius: 8px;
-        }
-        .user-message {
-            background-color: #f0f7ff;
-            border-left: 4px solid #1976d2;
-        }
-        .assistant-message {
-            background-color: #f8f9fa;
-            border-left: 4px solid #4caf50;
-        }
-        .suggestion-chip {
-            display: inline-block;
-            padding: 5px 15px;
-            margin: 5px;
-            background-color: #e3f2fd;
-            border-radius: 20px;
-            cursor: pointer;
-        }
-        </style>
-        """, unsafe_allow_html=True)
+        st.markdown(
+            """
+            <style>
+            .document-panel {
+                background-color: #f8f9fa;
+                border-radius: 10px;
+                padding: 20px;
+                margin: 10px;
+            }
+            .message {
+                padding: 15px;
+                margin: 10px 0;
+                border-radius: 8px;
+            }
+            .user-message {
+                background-color: #f0f7ff;
+                border-left: 4px solid #1976d2;
+            }
+            .assistant-message {
+                background-color: #f8f9fa;
+                border-left: 4px solid #4caf50;
+            }
+            .suggestion-chip {
+                display: inline-block;
+                padding: 5px 15px;
+                margin: 5px;
+                background-color: #e3f2fd;
+                border-radius: 20px;
+                cursor: pointer;
+            }
+            </style>
+            """,
+            unsafe_allow_html=True,
+        )
 
         # Create two-column layout
         col1, col2 = st.columns([2, 3])
@@ -66,7 +70,7 @@ class LegalNotebookInterface:
     def _render_document_panel(self):
         """Render the document viewing and management panel."""
         st.markdown("### 📚 Documents")
-        
+
         # Document selection
         cases = self.case_manager.get_all_cases()
         if not cases:
@@ -78,11 +82,11 @@ class LegalNotebookInterface:
             cases,
             format_func=lambda x: x['title']
         )
-        
+
         if selected_case:
             st.session_state.current_case = selected_case['id']
             documents = self.case_manager.list_documents(selected_case['id'])
-            
+
             if documents:
                 st.write("#### Available Documents")
                 for doc in documents:
@@ -97,7 +101,7 @@ class LegalNotebookInterface:
     def _render_chat_panel(self):
         """Render the chat and analysis interface."""
         st.markdown("### 💬 Document Analysis")
-        
+
         if not st.session_state.current_case:
             st.info("Please select a case from the left panel to start analysis.")
             return
@@ -105,49 +109,51 @@ class LegalNotebookInterface:
         # Display chat history
         for message in st.session_state.chat_history:
             message_class = "user-message" if message["role"] == "user" else "assistant-message"
-            st.markdown(f"""
+            st.markdown(
+                f"""
                 <div class="message {message_class}">
                     {message["content"]}
                 </div>
-            """, unsafe_allow_html=True)
+                """,
+                unsafe_allow_html=True,
+            )
 
         # Chat input
         if prompt := st.chat_input("Ask about your documents..."):
             self._handle_chat_input(prompt)
 
     def _handle_chat_input(self, prompt: str):
-    """Process user input and generate response."""
-    if not st.session_state.current_case:
-        st.error("Please select a case first.")
-        return
-
-    # Add user message
-    st.session_state.chat_history.append({
-        "role": "user",
-        "content": prompt,
-        "timestamp": datetime.datetime.now().isoformat()
-    })
-
-    try:
-        # Get documents from current case
-        documents = self.case_manager.list_documents(st.session_state.current_case)
-        if not documents:
-            st.error("No documents available for analysis.")
+        """Process user input and generate response."""
+        if not st.session_state.current_case:
+            st.error("Please select a case first.")
             return
 
-        # Get relevant chunks
+        # Add user message
+        st.session_state.chat_history.append({
+            "role": "user",
+            "content": prompt,
+            "timestamp": datetime.datetime.now().isoformat()
+        })
+
         try:
+            # Get documents from current case
+            documents = self.case_manager.list_documents(st.session_state.current_case)
+            if not documents:
+                st.error("No documents available for analysis.")
+                return
+
+            # Get relevant chunks
             results = self.vector_store.similarity_search(
                 query=prompt,
                 k=3  # Number of chunks to retrieve
             )
-            
+
             if not results:
                 response = "I couldn't find relevant information in the documents. Please try rephrasing your question."
             else:
                 # Generate response from results
                 response = self._generate_response(prompt, results)
-            
+
             # Add response to chat history
             st.session_state.chat_history.append({
                 "role": "assistant",
@@ -163,39 +169,36 @@ class LegalNotebookInterface:
                 "timestamp": datetime.datetime.now().isoformat()
             })
 
-    except Exception as e:
-        st.error(f"Error analyzing documents: {str(e)}")
-
     def _generate_response(self, prompt: str, results: List[Dict]) -> str:
-    """Generate response based on search results."""
-    if not results:
-        return "I couldn't find relevant information in the documents."
-
-    # Format the response with chunks
-    response = "Based on the documents, here's what I found:\n\n"
-    
-    for i, result in enumerate(results, 1):
-        response += f"**Excerpt {i}**\n"
-        response += f"{result.get('text', '')[:200]}...\n\n"
-        if result.get('metadata', {}).get('title'):
-            response += f"*Source: {result['metadata']['title']}*\n\n"
-
-    response += "\nWould you like me to elaborate on any of these points?"
-    
-    return response
+        """Generate response based on search results."""
+        if not results:
+            return "I couldn't find relevant information in the documents."
+
+        # Format the response with chunks
+        response = "Based on the documents, here's what I found:\n\n"
+
+        for i, result in enumerate(results, 1):
+            response += f"**Excerpt {i}**\n"
+            response += f"{result.get('text', '')[:200]}...\n\n"
+            if result.get('metadata', {}).get('title'):
+                response += f"*Source: {result['metadata']['title']}*\n\n"
+
+        response += "\nWould you like me to elaborate on any of these points?"
+
+        return response
 
     def _show_document_details(self, doc: Dict):
         """Show detailed document information."""
         st.markdown("#### Document Details")
         st.write("**Title:** ", doc['title'])
         st.write("**Added:** ", doc.get('added_at', 'Unknown'))
-        
+
         metadata = doc.get('metadata', {})
         if metadata:
             st.write("**Document Type:** ", metadata.get('document_type', 'Unknown'))
             st.write("**Word Count:** ", metadata.get('word_count', 'Unknown'))
-            
+
             if 'entities' in metadata:
                 st.write("**Named Entities:** ")
                 for entity_type, entities in metadata['entities'].items():
-                    st.write(f"- {entity_type}: {', '.join(entities[:5])}")
+                    st.write(f"- {entity_type}: {', '.join(entities[:5])}")