Update components/chat_interface.py

components/chat_interface.py

@@ -5,11 +5,15 @@ import os
 from datetime import datetime
 from utils.legal_prompt_generator import LegalPromptGenerator
 
+
 class ChatInterface:
-    def __init__(self, vector_store, document_processor):
+    def __init__(self, case_manager, vector_store, document_processor):
+        """Initialize ChatInterface with all required components."""
+        self.case_manager = case_manager
         self.vector_store = vector_store
         self.document_processor = document_processor
-        
+        self.prompt_generator = LegalPromptGenerator()
+
         try:
             api_key = os.getenv("ANTHROPIC_API_KEY")
             if not api_key:
@@ -27,9 +31,11 @@ class ChatInterface:
             st.session_state.analyzed_documents = []
         if "context_chunks" not in st.session_state:
             st.session_state.context_chunks = []
+        if "current_case" not in st.session_state:
+            st.session_state.current_case = None
 
     def render(self):
-        """Render an improved chat interface with better document context."""
+        """Render the chat interface with document and context management."""
         st.markdown("""
         <style>
         .chat-message {
@@ -61,7 +67,7 @@ class ChatInterface:
         </style>
         """, unsafe_allow_html=True)
 
-        # Display active documents and context
+        # Display active documents in the sidebar
         with st.sidebar:
             st.subheader("📚 Active Documents")
             for doc in st.session_state.analyzed_documents:
@@ -69,7 +75,7 @@ class ChatInterface:
                     st.write(f"Type: {doc.get('metadata', {}).get('type', 'Unknown')}")
                     st.write(f"Added: {doc.get('metadata', {}).get('added_at', 'Unknown')}")
 
-        # Display chat history with improved styling
+        # Display chat history
         for message in st.session_state.messages:
             message_class = "user-message" if message["role"] == "user" else "assistant-message"
             with st.container():
@@ -80,134 +86,90 @@ class ChatInterface:
                 </div>
                 """, unsafe_allow_html=True)
 
-        # Chat input with improved context handling
+        # Chat input
         if prompt := st.chat_input("Ask about your documents..."):
             self._handle_chat_input(prompt)
 
     def _handle_chat_input(self, prompt: str):
-        """Handle chat input with improved context management."""
-        # Add user message
+        """Process user input and generate a response."""
         st.session_state.messages.append({"role": "user", "content": prompt})
-        
-        # Get relevant context chunks
-        context_chunks = self.vector_store.similarity_search(
-            prompt,
-            k=5,
-            filter_criteria={"metadata.type": [doc["metadata"]["type"] for doc in st.session_state.analyzed_documents]}
-        )
-        
-        # Generate response
-        with st.spinner("Analyzing documents and generating response..."):
-            response_content, references = self.generate_response(prompt, context_chunks)
-            
-            # Add assistant message with references
-            st.session_state.messages.append({
-                "role": "assistant",
-                "content": response_content,
-                "references": references
-            })
-            
-            # Store context chunks for future reference
-            st.session_state.context_chunks = context_chunks
+
+        with st.spinner("Analyzing documents and generating a response..."):
+            try:
+                # Retrieve relevant document chunks
+                context_chunks = self.vector_store.similarity_search(
+                    query=prompt,
+                    k=5,
+                    filter_criteria={"metadata.type": [doc["metadata"]["type"] for doc in st.session_state.analyzed_documents]}
+                )
+
+                # Generate the response
+                response, references = self.generate_response(prompt, context_chunks)
+
+                # Add assistant response
+                st.session_state.messages.append({
+                    "role": "assistant",
+                    "content": response,
+                    "references": references
+                })
+
+                # Update context for future queries
+                st.session_state.context_chunks = context_chunks
+            except Exception as e:
+                st.error(f"Error generating response: {str(e)}")
 
     def generate_response(self, prompt: str, context_chunks: List[Dict]) -> tuple[str, str]:
-        """Generate response using Claude with improved context handling."""
+        """Generate a response using the LLM and LegalPromptGenerator."""
         try:
-            # Prepare context from chunks
-            context = "\n".join([
-                f"Document: {chunk['metadata']['title']}\n"
-                f"Section: {chunk['text']}\n"
-                f"Type: {chunk['metadata']['type']}\n"
-                f"Jurisdiction: {chunk['metadata']['jurisdiction']}\n"
-                for chunk in context_chunks
-            ])
-
-            # Generate system message using ontology
-            system_message = self._generate_system_message(prompt, context_chunks)
-
-            # Call Claude API
-            message = self.client.messages.create(
-                model="claude-3-sonnet-20240229",
+            # Generate structured messages
+            messages = self._generate_messages(prompt, context_chunks)
+
+            # Call the LLM
+            response = self.client.messages.create(
+                model="claude-3",
                 max_tokens=2000,
                 temperature=0.7,
-                messages=[
-                    {"role": "system", "content": system_message},
-                    {"role": "user", "content": f"Question: {prompt}\n\nContext:\n{context}"}
-                ]
+                messages=messages
             )
-            
+
             # Format references
             references_html = self._format_references(context_chunks)
-            
-            return message.content[0].text, references_html
-            
-        except Exception as e:
-            st.error(f"Error generating response: {str(e)}")
-            return "I apologize, but I encountered an error generating the response.", ""
+            return response.content[0].text, references_html
 
-    def __init__(self, case_manager, vector_store, document_processor):
-        """Initialize ChatInterface with enhanced components."""
-        self.case_manager = case_manager
-        self.vector_store = vector_store
-        self.document_processor = document_processor
-        self.prompt_generator = LegalPromptGenerator()
-        
-        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 "analyzed_documents" not in st.session_state:
-            st.session_state.analyzed_documents = []
-        if "context_chunks" not in st.session_state:
-            st.session_state.context_chunks = []
-        if "current_case" not in st.session_state:
-            st.session_state.current_case = None
+            st.error(f"Error generating response: {str(e)}")
+            return "An error occurred while processing your query.", ""
 
     def _generate_messages(self, prompt: str, context_chunks: List[Dict]) -> List[Dict]:
-        """Generate messages for the Claude API with enhanced prompts."""
+        """Generate structured messages for LLM input."""
         # Get case metadata if available
         case_metadata = None
         if st.session_state.current_case:
             case_metadata = self.case_manager.get_case(st.session_state.current_case)
-        
-        # Generate enhanced system message
+
+        # Generate system message
         system_message = self.prompt_generator.generate_system_message(
             context_chunks=context_chunks,
             query=prompt,
             case_metadata=case_metadata
         )
-        
-        # Prepare context from chunks
+
+        # Generate user message
         context = "\n".join([
-            f"Document: {chunk['metadata']['title']}\n"
+            f"Document: {chunk['metadata'].get('title', 'Untitled')}\n"
             f"Section: {chunk['text']}\n"
-            f"Type: {chunk['metadata']['type']}\n"
-            f"Jurisdiction: {chunk['metadata']['jurisdiction']}\n"
             for chunk in context_chunks
         ])
-        
-        # Generate user message
         user_message = self.prompt_generator.generate_user_message(prompt, context)
-        
-        # Check if this is a follow-up question
+
+        # Handle follow-up questions
         if st.session_state.messages:
             previous_query = next(
-                (m["content"] for m in reversed(st.session_state.messages)
-                if m["role"] == "user"),
+                (m["content"] for m in reversed(st.session_state.messages) if m["role"] == "user"),
                 None
             )
             previous_response = next(
-                (m["content"] for m in reversed(st.session_state.messages)
-                if m["role"] == "assistant"),
+                (m["content"] for m in reversed(st.session_state.messages) if m["role"] == "assistant"),
                 None
             )
             if previous_query and previous_response:
@@ -217,19 +179,19 @@ class ChatInterface:
                     previous_response=previous_response,
                     context_chunks=context_chunks
                 )
-        
+
         return [
             {"role": "system", "content": system_message},
             {"role": "user", "content": user_message}
         ]
 
     def _format_references(self, chunks: List[Dict]) -> str:
-        """Format reference citations in HTML."""
+        """Format references as HTML for display."""
         references = []
         for i, chunk in enumerate(chunks, 1):
             references.append(f"""
             <div class="document-chunk">
-                <strong>Reference {i}:</strong> {chunk['metadata']['title']}
+                <strong>Reference {i}:</strong> {chunk['metadata'].get('title', 'Untitled')}
                 <br/>
                 <em>Section:</em> {chunk['text'][:200]}...
             </div>
@@ -237,7 +199,7 @@ class ChatInterface:
         return "\n".join(references)
 
     def add_analyzed_document(self, doc: Dict):
-        """Add a document with improved metadata tracking."""
+        """Add a document to session state with metadata tracking."""
         doc['metadata']['added_at'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         if doc not in st.session_state.analyzed_documents:
-            st.session_state.analyzed_documents.append(doc)
+            st.session_state.analyzed_documents.append(doc)