Update components/chat_interface.py

components/chat_interface.py

@@ -2,98 +2,181 @@ import streamlit as st
 from typing import List, Dict
 import anthropic
 import os
+from datetime import datetime
 
 class ChatInterface:
-    def __init__(self, vector_store):
+    def __init__(self, vector_store, document_processor):
         self.vector_store = vector_store
-
-        # Initialize Anthropic client using environment variable
+        self.document_processor = document_processor
+        
         try:
             api_key = os.getenv("ANTHROPIC_API_KEY")
             if not api_key:
-                st.error("Please set the ANTHROPIC_API_KEY in your environment variables.")
+                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 chat history and analyzed documents
+        # 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 = []
 
     def render(self):
-        """Render chat interface with analyzed documents and history."""
-        st.subheader("Chat with Documents")
+        """Render an improved chat interface with better document context."""
+        st.markdown("""
+        <style>
+        .chat-message {
+            padding: 1.5rem;
+            border-radius: 0.5rem;
+            margin-bottom: 1rem;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }
+        .user-message {
+            background-color: #f0f7ff;
+            border-left: 4px solid #2B547E;
+        }
+        .assistant-message {
+            background-color: #ffffff;
+            border-left: 4px solid #4CAF50;
+        }
+        .reference-box {
+            background-color: #f5f5f5;
+            padding: 0.8rem;
+            border-radius: 0.3rem;
+            font-size: 0.9em;
+            margin-top: 0.5rem;
+        }
+        .document-chunk {
+            border-left: 3px solid #2196F3;
+            padding-left: 1rem;
+            margin: 0.5rem 0;
+        }
+        </style>
+        """, unsafe_allow_html=True)
 
-        # Display analyzed documents
-        with st.expander("Analyzed Documents", expanded=True):
-            if st.session_state.analyzed_documents:
-                for doc in st.session_state.analyzed_documents:
-                    st.markdown(f"- **{doc['name']}**")
-            else:
-                st.info("No documents analyzed yet.")
+        # Display active documents and context
+        with st.sidebar:
+            st.subheader("📚 Active Documents")
+            for doc in st.session_state.analyzed_documents:
+                with st.expander(f"📄 {doc['name']}", expanded=False):
+                    st.write(f"Type: {doc.get('metadata', {}).get('type', 'Unknown')}")
+                    st.write(f"Added: {doc.get('metadata', {}).get('added_at', 'Unknown')}")
 
-        # Display chat history
+        # Display chat history with improved styling
         for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                st.markdown(message["content"])
+            message_class = "user-message" if message["role"] == "user" else "assistant-message"
+            with st.container():
+                st.markdown(f"""
+                <div class="chat-message {message_class}">
+                    {message["content"]}
+                    {'<div class="reference-box">' + message.get("references", "") + '</div>' if message.get("references") else ""}
+                </div>
+                """, unsafe_allow_html=True)
+
+        # Chat input with improved context handling
+        if prompt := st.chat_input("Ask about your documents..."):
+            self._handle_chat_input(prompt)
 
-        # Chat input
-        if prompt := st.chat_input("Ask a question about your documents"):
-            # Add user message to history
-            st.session_state.messages.append({"role": "user", "content": prompt})
+    def _handle_chat_input(self, prompt: str):
+        """Handle chat input with improved context management."""
+        # Add user message
+        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
+            })
             
-            with st.chat_message("user"):
-                st.markdown(prompt)
+            # Store context chunks for future reference
+            st.session_state.context_chunks = context_chunks
 
-            # Get relevant context from vector store
-            results = self.vector_store.similarity_search(prompt, k=3)
-            context = "\n\n".join([r["metadata"]["text"] for r in results])
+    def generate_response(self, prompt: str, context_chunks: List[Dict]) -> tuple[str, str]:
+        """Generate response using Claude with improved context handling."""
+        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 response with references
-            with st.chat_message("assistant"):
-                with st.spinner("Thinking..."):
-                    response = self.generate_response(prompt, context, results)
-                    st.markdown(response)
-                    st.session_state.messages.append({"role": "assistant", "content": response})
+            # Generate system message using ontology
+            system_message = self._generate_system_message(prompt, context_chunks)
 
-    def generate_response(self, prompt: str, context: str, results: List[Dict]) -> str:
-        """Generate response using Claude with document references."""
-        try:
-            # Call the Claude API for response generation
+            # Call Claude API
             message = self.client.messages.create(
                 model="claude-3-sonnet-20240229",
                 max_tokens=2000,
                 temperature=0.7,
-                messages=[{
-                    "role": "user",
-                    "content": f"""Based on the following context from legal documents, please answer the question.
-                    
-                    Context:
-                    {context}
-                    
-                    Question: {prompt}
-                    
-                    Please provide a detailed response with references to specific parts of the documents when relevant."""
-                }]
+                messages=[
+                    {"role": "system", "content": system_message},
+                    {"role": "user", "content": f"Question: {prompt}\n\nContext:\n{context}"}
+                ]
             )
-            response_content = message.content[0].text
-
-            # Append document references to the response
-            references = [
-                f"{idx + 1}. {res['metadata']['text'][:200]}... (Reference: {res['metadata'].get('reference', 'N/A')})"
-                for idx, res in enumerate(results)
-            ]
-            references_text = "\n\nReferences:\n" + "\n".join(references)
-            return response_content + references_text
+            
+            # 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 "I apologize, but I encountered an error generating the response.", ""
+
+    def _generate_system_message(self, prompt: str, context_chunks: List[Dict]) -> str:
+        """Generate a system message using ontology and document context."""
+        # Get relevant ontology concepts
+        ontology_concepts = self.document_processor._link_to_ontology(prompt)
+        
+        return f"""You are a legal AI assistant analyzing documents with the following context:
+        
+        Document Types Present: {', '.join(set(chunk['metadata']['type'] for chunk in context_chunks))}
+        Jurisdictions: {', '.join(set(chunk['metadata']['jurisdiction'] for chunk in context_chunks))}
+        Relevant Legal Concepts: {', '.join(concept['concept'] for concept in ontology_concepts)}
+        
+        Please provide detailed analysis while:
+        1. Citing specific sections from the provided context
+        2. Incorporating relevant legal concepts and terminology
+        3. Maintaining appropriate legal language and tone
+        4. Providing clear references to source documents
+        """
+
+    def _format_references(self, chunks: List[Dict]) -> str:
+        """Format reference citations in HTML."""
+        references = []
+        for i, chunk in enumerate(chunks, 1):
+            references.append(f"""
+            <div class="document-chunk">
+                <strong>Reference {i}:</strong> {chunk['metadata']['title']}
+                <br/>
+                <em>Section:</em> {chunk['text'][:200]}...
+            </div>
+            """)
+        return "\n".join(references)
 
     def add_analyzed_document(self, doc: Dict):
-        """Add a document to the list of analyzed documents."""
+        """Add a document with improved 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)