Update utils/legal_notebook_interface.py

utils/legal_notebook_interface.py

@@ -116,63 +116,73 @@ class LegalNotebookInterface:
             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()
-        })
-
+    """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.")
             return
 
+        # Get relevant chunks
         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)
-
-            # Generate response
-            response = {
+            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",
-                "content": self._generate_response(prompt, chunks),
+                "content": response,
                 "timestamp": datetime.datetime.now().isoformat()
-            }
-
-            # 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."
-
-        # Combine context from chunks
-        context = "\n\n".join([chunk['text'] for chunk in chunks])
-        
-        # 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
+            st.error(f"Error searching documents: {str(e)}")
+            st.session_state.chat_history.append({
+                "role": "assistant",
+                "content": "I encountered an error while searching the documents. Please try again.",
+                "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
 
     def _show_document_details(self, doc: Dict):
         """Show detailed document information."""