Update components/chat_interface.py

components/chat_interface.py

@@ -1,21 +1,41 @@
-# components/chat_interface.py
 import streamlit as st
 from typing import List, Dict
 import anthropic
+import os
 
 class ChatInterface:
     def __init__(self, vector_store):
         self.vector_store = vector_store
-        self.client = anthropic.Anthropic(api_key=st.secrets["Anthropic_API_KEY"])
-        
-        # Initialize chat history
+
+        # Initialize Anthropic client using environment variable
+        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.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
         if "messages" not in st.session_state:
             st.session_state.messages = []
+        if "analyzed_documents" not in st.session_state:
+            st.session_state.analyzed_documents = []
 
     def render(self):
-        """Render chat interface"""
+        """Render chat interface with analyzed documents and history."""
         st.subheader("Chat with Documents")
 
+        # 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 chat history
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
@@ -33,16 +53,17 @@ class ChatInterface:
             results = self.vector_store.similarity_search(prompt, k=3)
             context = "\n\n".join([r["metadata"]["text"] for r in results])
 
-            # Generate response
+            # Generate response with references
             with st.chat_message("assistant"):
                 with st.spinner("Thinking..."):
-                    response = self.generate_response(prompt, context)
+                    response = self.generate_response(prompt, context, results)
                     st.markdown(response)
                     st.session_state.messages.append({"role": "assistant", "content": response})
 
-    def generate_response(self, prompt: str, context: str) -> str:
-        """Generate response using Claude"""
+    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
             message = self.client.messages.create(
                 model="claude-3-sonnet-20240229",
                 max_tokens=2000,
@@ -59,7 +80,20 @@ class ChatInterface:
                     Please provide a detailed response with references to specific parts of the documents when relevant."""
                 }]
             )
-            return message.content[0].text
+            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
         except Exception as e:
             st.error(f"Error generating response: {str(e)}")
             return "I apologize, but I encountered an error generating the response."
+
+    def add_analyzed_document(self, doc: Dict):
+        """Add a document to the list of analyzed documents."""
+        if doc not in st.session_state.analyzed_documents:
+            st.session_state.analyzed_documents.append(doc)