Create utils/legal_notebook_interface.py

utils/legal_notebook_interface.py

@@ -0,0 +1,242 @@
+import streamlit as st
+from typing import List, Dict
+import datetime
+
+class LegalNotebookInterface:
+    def __init__(self, case_manager, vector_store, document_processor):
+        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 = []
+        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 = []
+
+    def render(self):
+        """Render the main notebook interface."""
+        # Custom CSS for NotebookLM-like interface
+        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;
+            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;
+        }
+        .document-snippet {
+            padding: 10px;
+            background-color: #fff;
+            border-left: 3px solid #ffd700;
+            margin: 10px 0;
+            border-radius: 4px;
+        }
+        </style>
+        """, unsafe_allow_html=True)
+
+        # Create two-column layout
+        col1, col2 = st.columns([2, 3])
+
+        # Document Panel (Left)
+        with col1:
+            self._render_document_panel()
+
+        # Chat Panel (Right)
+        with col2:
+            self._render_chat_panel()
+
+    def _render_document_panel(self):
+        """Render the document viewing and management panel."""
+        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']
+            )
+            
+            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')
+                            })
+
+    def _render_chat_panel(self):
+        """Render the chat and analysis panel."""
+        st.markdown("### 💬 Legal Analysis")
+        
+        # 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)
+
+        # 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)
+
+    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):
+        """Process user input and generate response."""
+        # Add user message
+        st.session_state.chat_history.append({
+            'role': 'user',
+            'content': input_text,
+            '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()
+        }
+
+        # Add assistant message
+        st.session_state.chat_history.append(response)
+
+        # Generate new suggestions
+        self._generate_suggestions(input_text, context_chunks)
+
+    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])
+        
+        # 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"
+        
+        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?"
+        ]