Update utils/legal_notebook_interface.py

utils/legal_notebook_interface.py

@@ -4,6 +4,7 @@ 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
@@ -13,35 +14,19 @@ class LegalNotebookInterface:
             st.session_state.active_documents = []
         if "chat_history" not in st.session_state:
             st.session_state.chat_history = []
-        if "document_context" not in st.session_state:
-            st.session_state.document_context = None
-        if "suggestions" not in st.session_state:
-            st.session_state.suggestions = []
+        if "current_case" not in st.session_state:
+            st.session_state.current_case = None
 
     def render(self):
         """Render the main notebook interface."""
-        # Custom CSS for NotebookLM-like interface
+        # Add custom styling
         st.markdown("""
         <style>
-        .main-container {
-            max-width: 1200px;
-            margin: auto;
-        }
         .document-panel {
             background-color: #f8f9fa;
             border-radius: 10px;
             padding: 20px;
             margin: 10px;
-            height: calc(100vh - 100px);
-            overflow-y: auto;
-        }
-        .chat-panel {
-            background-color: white;
-            border-radius: 10px;
-            padding: 20px;
-            margin: 10px;
-            height: calc(100vh - 100px);
-            overflow-y: auto;
         }
         .message {
             padding: 15px;
@@ -64,13 +49,6 @@ class LegalNotebookInterface:
             border-radius: 20px;
             cursor: pointer;
         }
-        .document-snippet {
-            padding: 10px;
-            background-color: #fff;
-            border-left: 3px solid #ffd700;
-            margin: 10px 0;
-            border-radius: 4px;
-        }
         </style>
         """, unsafe_allow_html=True)
 
@@ -90,153 +68,124 @@ class LegalNotebookInterface:
         st.markdown("### 📚 Documents")
         
         # Document selection
-        all_cases = self.case_manager.get_all_cases()
-        if all_cases:
-            selected_case = st.selectbox(
-                "Select Case",
-                all_cases,
-                format_func=lambda x: x['title']
-            )
+        cases = self.case_manager.get_all_cases()
+        if not cases:
+            st.info("No cases found. 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:
+            st.session_state.current_case = selected_case['id']
+            documents = self.case_manager.list_documents(selected_case['id'])
             
-            if selected_case:
-                documents = self.case_manager.list_documents(selected_case['id'])
-                selected_docs = st.multiselect(
-                    "Select Documents to Analyze",
-                    documents,
-                    format_func=lambda x: x['title']
-                )
-                
-                if selected_docs:
-                    st.session_state.active_documents = selected_docs
-                    
-                    # Show document previews
-                    for doc in selected_docs:
-                        with st.expander(f"📄 {doc['title']}", expanded=False):
-                            st.markdown("#### Quick Preview")
-                            content = doc.get('content', '')[:500]
-                            st.markdown(f"{content}...")
-                            
-                            # Document metadata
-                            st.markdown("#### Metadata")
-                            st.json({
-                                'Type': doc.get('metadata', {}).get('type', 'Unknown'),
-                                'Added': doc.get('added_at', 'Unknown'),
-                                'Pages': doc.get('metadata', {}).get('page_count', 'Unknown')
-                            })
+            if documents:
+                st.write("#### Available Documents")
+                for doc in documents:
+                    with st.expander(f"📄 {doc['title']}", expanded=False):
+                        st.write("**Added:** ", doc.get('added_at', 'Unknown'))
+                        st.write("**Type:** ", doc.get('metadata', {}).get('document_type', 'Unknown'))
+                        if st.button("View Details", key=f"view_{doc['id']}"):
+                            self._show_document_details(doc)
+            else:
+                st.info("No documents in this case. Please add documents first.")
 
     def _render_chat_panel(self):
-        """Render the chat and analysis panel."""
-        st.markdown("### 💬 Legal Analysis")
+        """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
+
         # Display chat history
         for message in st.session_state.chat_history:
-            self._render_message(message)
-
-        # Suggestions
-        if st.session_state.suggestions:
-            st.markdown("#### 💡 Suggested Questions")
-            cols = st.columns(2)
-            for idx, suggestion in enumerate(st.session_state.suggestions):
-                with cols[idx % 2]:
-                    if st.button(suggestion, key=f"suggestion_{idx}"):
-                        self._handle_input(suggestion)
+            message_class = "user-message" if message["role"] == "user" else "assistant-message"
+            st.markdown(f"""
+                <div class="message {message_class}">
+                    {message["content"]}
+                </div>
+            """, unsafe_allow_html=True)
 
         # Chat input
-        with st.container():
-            chat_input = st.text_area("Ask about your documents...", height=100)
-            col1, col2 = st.columns([1, 5])
-            with col1:
-                if st.button("Send", type="primary"):
-                    if chat_input:
-                        self._handle_input(chat_input)
-            with col2:
-                st.markdown("*Press Enter to send*")
-
-    def _render_message(self, message: Dict):
-        """Render a single message with citations and context."""
-        message_class = "user-message" if message['role'] == 'user' else 'assistant-message'
-        
-        st.markdown(f"""
-        <div class="message {message_class}">
-            <p>{message['content']}</p>
-            {self._render_citations(message.get('citations', [])) if message.get('citations') else ''}
-        </div>
-        """, unsafe_allow_html=True)
+        if prompt := st.chat_input("Ask about your documents..."):
+            self._handle_chat_input(prompt)
 
-    def _render_citations(self, citations: List[Dict]) -> str:
-        """Render document citations and snippets."""
-        if not citations:
-            return ""
-            
-        citations_html = "<div class='citations'>"
-        for citation in citations:
-            citations_html += f"""
-            <div class="document-snippet">
-                <small>From {citation['document']}</small>
-                <p>{citation['text']}</p>
-            </div>
-            """
-        citations_html += "</div>"
-        return citations_html
-
-    def _handle_input(self, input_text: str):
+    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': input_text,
-            'timestamp': datetime.datetime.now().isoformat()
+            "role": "user",
+            "content": prompt,
+            "timestamp": datetime.datetime.now().isoformat()
         })
 
-        # Get relevant context from documents
-        context_chunks = []
-        for doc in st.session_state.active_documents:
-            chunks = self.vector_store.similarity_search(
-                input_text,
-                filter_criteria={"doc_id": doc['id']},
-                k=2
-            )
-            context_chunks.extend(chunks)
-
-        # Generate response with citations
-        response = {
-            'role': 'assistant',
-            'content': self._generate_response(input_text, context_chunks),
-            'citations': [
-                {
-                    'document': chunk['metadata']['title'],
-                    'text': chunk['text']
-                }
-                for chunk in context_chunks[:2]
-            ],
-            'timestamp': datetime.datetime.now().isoformat()
-        }
+        # 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
+
+        try:
+            # Get relevant chunks from documents
+            chunks = []
+            for doc in documents:
+                doc_chunks = self.vector_store.similarity_search(
+                    prompt,
+                    filter_criteria={"doc_id": doc['id']},
+                    k=2
+                )
+                if doc_chunks:
+                    chunks.extend(doc_chunks)
 
-        # Add assistant message
-        st.session_state.chat_history.append(response)
+            # Generate response
+            response = {
+                "role": "assistant",
+                "content": self._generate_response(prompt, chunks),
+                "timestamp": datetime.datetime.now().isoformat()
+            }
 
-        # Generate new suggestions
-        self._generate_suggestions(input_text, context_chunks)
+            # Add response to chat history
+            st.session_state.chat_history.append(response)
+
+        except Exception as e:
+            st.error(f"Error analyzing documents: {str(e)}")
+
+    def _generate_response(self, prompt: str, chunks: List[Dict]) -> str:
+        """Generate response based on document chunks."""
+        if not chunks:
+            return "I couldn't find any relevant information in the documents. Please try rephrasing your question."
 
-    def _generate_response(self, query: str, context_chunks: List[Dict]) -> str:
-        """Generate response using the legal analysis pipeline."""
         # Combine context from chunks
-        context = "\n\n".join([chunk['text'] for chunk in context_chunks])
+        context = "\n\n".join([chunk['text'] for chunk in chunks])
         
-        # Get response from your existing chat interface
-        response_content = f"Analysis based on {len(context_chunks)} relevant document sections:\n\n"
-        response_content += "1. Key findings...\n"
-        response_content += "2. Legal implications...\n"
-        response_content += "3. Recommendations...\n"
+        # Generate a simple response (replace with your actual response generation)
+        response = f"Based on the documents, I found {len(chunks)} relevant sections:\n\n"
+        for i, chunk in enumerate(chunks, 1):
+            response += f"{i}. {chunk['text'][:200]}...\n\n"
         
-        return response_content
-
-    def _generate_suggestions(self, query: str, context_chunks: List[Dict]):
-        """Generate follow-up question suggestions."""
-        # Generate relevant follow-up questions based on context
-        st.session_state.suggestions = [
-            "Could you elaborate on the legal implications?",
-            "What are the relevant precedents?",
-            "Are there any potential risks?",
-            "How does this compare to similar cases?"
-        ]
+        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])}")