Fix DNS resolution errors and restore original UI

streamlit_app.py

@@ -204,143 +204,85 @@ def get_chat_model():
     """Get the chat model for initial RAG."""
     print("Initializing chat model...")
     try:
-        # Very minimal OpenAI initialization for Hugging Face compatibility
+        # Set API key from environment
         openai_api_key = os.environ.get("OPENAI_API_KEY", "")
         if not openai_api_key:
             print("WARNING: OPENAI_API_KEY environment variable not set!")
             raise ValueError("OpenAI API key not found")
+        
+        # Create a wrapper class with a shorter timeout to fail faster on DNS issues
+        class TimeoutChatModel:
+            def __init__(self, api_key):
+                self.api_key = api_key
+                self.timeout = 5  # Short timeout to fail fast on DNS issues
             
-        # First try: OpenAI client
-        try:
-            openai_client = OpenAI(api_key=openai_api_key)
-            
-            # Test the connection
-            print("Testing OpenAI chat API connection...")
-            test_response = openai_client.chat.completions.create(
-                model="gpt-3.5-turbo",
-                messages=[{"role": "user", "content": "test"}],
-                max_tokens=5
-            )
-            print("OpenAI chat API connection successful!")
-            
-            # Create a simplified wrapper that avoids any problematic parameters
-            class SimpleOpenAIWrapper:
-                def invoke(self, messages):
-                    print("Invoking chat model via client...")
-                    # Convert LangChain messages to OpenAI format
-                    openai_messages = []
-                    for msg in messages:
-                        role = "user"
-                        if hasattr(msg, "type"):
-                            role = "assistant" if msg.type == "ai" else "user"
-                        openai_messages.append({
-                            "role": role,
-                            "content": msg.content
-                        })
+            def invoke(self, messages):
+                print("Invoking chat model...")
+                try:
+                    # Convert string input to message format if needed
+                    if isinstance(messages, str):
+                        openai_messages = [{"role": "user", "content": messages}]
+                    else:
+                        # Convert LangChain messages to OpenAI format
+                        openai_messages = []
+                        for msg in messages:
+                            role = "user"
+                            if hasattr(msg, "type"):
+                                role = "assistant" if msg.type == "ai" else "user"
+                            openai_messages.append({
+                                "role": role,
+                                "content": msg.content
+                            })
                     
-                    # Call API directly with absolutely minimal parameters
-                    try:
-                        response = openai_client.chat.completions.create(
-                            model="gpt-3.5-turbo", # Use a minimal, widely supported model
-                            messages=openai_messages
-                        )
-                        
-                        # Create response object
-                        class SimpleResponse:
-                            def __init__(self, content):
-                                self.content = content
-                        
-                        result = SimpleResponse(response.choices[0].message.content)
-                        print(f"Got response of length: {len(result.content)}")
-                        return result
-                    except Exception as e:
-                        print(f"Error calling OpenAI API: {str(e)}")
-                        # Fallback to HTTP request
-                        content = http_chat_request(openai_messages, openai_api_key)
+                    # Direct API call with timeout
+                    import requests
+                    import json
+                    
+                    url = "https://api.openai.com/v1/chat/completions"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {self.api_key}"
+                    }
+                    data = {
+                        "model": "gpt-3.5-turbo",
+                        "messages": openai_messages
+                    }
+                    
+                    response = requests.post(
+                        url, 
+                        headers=headers, 
+                        data=json.dumps(data),
+                        timeout=self.timeout
+                    )
+                    
+                    if response.status_code == 200:
+                        result = response.json()
+                        content = result["choices"][0]["message"]["content"]
+                        print(f"Got response of length: {len(content)}")
                         return type('obj', (object,), {'content': content})
-            
-            return SimpleOpenAIWrapper()
-            
-        except Exception as e:
-            print(f"OpenAI client approach failed: {str(e)}")
-            print("Falling back to direct HTTP requests...")
-            raise  # Continue to HTTP fallback
-            
-    except Exception as e:
-        print(f"Standard chat model approach failed: {str(e)}")
+                    else:
+                        print(f"API request failed with status {response.status_code}")
+                        raise Exception(f"API request failed: {response.text}")
+                except requests.exceptions.Timeout:
+                    print("Timeout connecting to OpenAI API")
+                    raise Exception("Timeout connecting to OpenAI API")
+                except requests.exceptions.ConnectionError as e:
+                    print(f"Connection error to OpenAI API: {str(e)}")
+                    raise Exception(f"Connection error: {str(e)}")
+                except Exception as e:
+                    print(f"Error in chat model: {str(e)}")
+                    raise
         
-        # Direct HTTP fallback approach
-        try:
-            openai_api_key = os.environ.get("OPENAI_API_KEY", "")
-            if not openai_api_key:
-                raise ValueError("OpenAI API key not found")
-                
-            # Test HTTP connection
-            print("Testing direct HTTP connection to OpenAI chat...")
-            test_message = http_chat_request([{"role": "user", "content": "test"}], openai_api_key)
-            if not test_message:
-                raise ValueError("HTTP chat fallback test failed")
-            print("Direct HTTP chat connection successful!")
-            
-            class HTTPChatModel:
-                def invoke(self, messages):
-                    print("Invoking chat model via HTTP...")
-                    # Convert LangChain messages to OpenAI format
-                    openai_messages = []
-                    for msg in messages:
-                        role = "user"
-                        if hasattr(msg, "type"):
-                            role = "assistant" if msg.type == "ai" else "user"
-                        openai_messages.append({
-                            "role": role,
-                            "content": msg.content
-                        })
-                    
-                    content = http_chat_request(openai_messages, openai_api_key)
-                    return type('obj', (object,), {'content': content})
-            
-            return HTTPChatModel()
-                
-        except Exception as e:
-            print(f"All chat model approaches failed: {str(e)}")
-            # Create dummy for testing
-            class DummyModel:
-                def invoke(self, messages):
-                    print("WARNING: Using dummy model!")
-                    return type('obj', (object,), {'content': 'I apologize, but I cannot access the necessary data to answer this question due to API connectivity issues.'})
-            
-            return DummyModel()
-
-# Add HTTP chat completion function
-def http_chat_request(messages, api_key):
-    """Make a direct HTTP request to OpenAI chat API."""
-    import requests
-    import json
-    
-    print("Using direct HTTP request for chat completion")
-    url = "https://api.openai.com/v1/chat/completions"
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
-    }
-    data = {
-        "model": "gpt-3.5-turbo",
-        "messages": messages
-    }
-    
-    try:
-        response = requests.post(url, headers=headers, data=json.dumps(data))
-        if response.status_code == 200:
-            result = response.json()
-            content = result["choices"][0]["message"]["content"]
-            print(f"Successfully got chat response via HTTP (length: {len(content)})")
-            return content
-        else:
-            print(f"HTTP chat request failed with status {response.status_code}: {response.text}")
-            return "I apologize, but I encountered an error connecting to the AI service."
+        return TimeoutChatModel(openai_api_key)
     except Exception as e:
-        print(f"HTTP chat request exception: {str(e)}")
-        return "I apologize, but I encountered a technical issue while processing your request."
+        print(f"Error initializing chat model: {str(e)}")
+        # Create dummy for testing
+        class DummyModel:
+            def invoke(self, messages):
+                print("WARNING: Using dummy model!")
+                return type('obj', (object,), {'content': 'I apologize, but I cannot access the necessary data to answer this question due to API connectivity issues.'})
+        
+        return DummyModel()
 
 @st.cache_resource
 def get_agent_model():
@@ -354,127 +296,80 @@ def get_embedding_model():
     """Get the embedding model."""
     print("Initializing embedding model...")
     try:
-        # Very minimal OpenAI initialization for Hugging Face compatibility
+        # Set API key from environment
         openai_api_key = os.environ.get("OPENAI_API_KEY", "")
         if not openai_api_key:
             print("WARNING: OPENAI_API_KEY environment variable not set!")
             raise ValueError("OpenAI API key not found")
+        
+        # Create a wrapper class with a shorter timeout to fail faster on DNS issues
+        class TimeoutEmbeddings:
+            def __init__(self, api_key):
+                self.api_key = api_key
+                self.timeout = 5  # Short timeout to fail fast on DNS issues
             
-        # First try: Use OpenAI client
-        try:
-            openai_client = OpenAI(api_key=openai_api_key)
-            
-            # Test the connection
-            print("Testing OpenAI API connection...")
-            response = openai_client.embeddings.create(
-                model="text-embedding-ada-002",
-                input="Test"
-            )
-            print("OpenAI API connection successful!")
-            
-            # Create a wrapper that avoids any problematic parameters
-            class SimpleEmbeddings:
-                def embed_query(self, text):
-                    print(f"Embedding query of length: {len(text)}")
-                    try:
-                        response = openai_client.embeddings.create(
-                            model="text-embedding-ada-002", # Use older, more compatible model
-                            input=text
-                        )
+            def embed_query(self, text):
+                print(f"Embedding query of length: {len(text)}")
+                try:
+                    # Direct API call with timeout
+                    import requests
+                    import json
+                    
+                    url = "https://api.openai.com/v1/embeddings"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {self.api_key}"
+                    }
+                    data = {
+                        "model": "text-embedding-ada-002",
+                        "input": text
+                    }
+                    
+                    response = requests.post(
+                        url, 
+                        headers=headers, 
+                        data=json.dumps(data),
+                        timeout=self.timeout
+                    )
+                    
+                    if response.status_code == 200:
+                        result = response.json()
                         print("Successfully got embedding")
-                        return response.data[0].embedding
-                    except Exception as e:
-                        print(f"Error in embed_query: {str(e)}")
-                        # Fall back to direct HTTP request
-                        return http_embed_request(text, openai_api_key)
-                
-                def embed_documents(self, texts):
-                    print(f"Embedding {len(texts)} documents")
-                    results = []
-                    for i, text in enumerate(texts):
-                        results.append(self.embed_query(text))
-                    return results
-            
-            return SimpleEmbeddings()
-            
-        except Exception as e:
-            print(f"OpenAI client failed: {str(e)}")
-            print("Falling back to direct HTTP requests...")
-            raise  # Continue to HTTP fallback
+                        return result["data"][0]["embedding"]
+                    else:
+                        print(f"API request failed with status {response.status_code}")
+                        raise Exception(f"API request failed: {response.text}")
+                except requests.exceptions.Timeout:
+                    print("Timeout connecting to OpenAI API - using dummy embedding")
+                    return [0.0] * 1536
+                except requests.exceptions.ConnectionError:
+                    print("Connection error to OpenAI API - using dummy embedding")
+                    return [0.0] * 1536
+                except Exception as e:
+                    print(f"Error getting embeddings: {str(e)}")
+                    return [0.0] * 1536
             
+            def embed_documents(self, texts):
+                print(f"Embedding {len(texts)} documents")
+                results = []
+                for i, text in enumerate(texts):
+                    results.append(self.embed_query(text))
+                return results
+        
+        return TimeoutEmbeddings(openai_api_key)
     except Exception as e:
-        print(f"Standard embedding approach failed: {str(e)}")
+        print(f"Error initializing embedding model: {str(e)}")
         
-        # Direct HTTP fallback approach
-        try:
-            openai_api_key = os.environ.get("OPENAI_API_KEY", "")
-            if not openai_api_key:
-                raise ValueError("OpenAI API key not found")
-                
-            # Test the connection with direct HTTP
-            print("Testing direct HTTP connection to OpenAI...")
-            test_embedding = http_embed_request("Test", openai_api_key)
-            if not test_embedding:
-                raise ValueError("HTTP fallback test failed")
-            print("Direct HTTP connection successful!")
-            
-            class HTTPEmbeddings:
-                def embed_query(self, text):
-                    print(f"HTTP embedding query of length: {len(text)}")
-                    return http_embed_request(text, openai_api_key)
-                
-                def embed_documents(self, texts):
-                    print(f"HTTP embedding {len(texts)} documents")
-                    results = []
-                    for text in texts:
-                        results.append(self.embed_query(text))
-                    return results
+        # Create dummy for testing
+        class DummyEmbeddings:
+            def embed_query(self, text):
+                print("WARNING: Using dummy embeddings!")
+                return [0.0] * 1536
             
-            return HTTPEmbeddings()
-                
-        except Exception as e:
-            print(f"All embedding approaches failed: {str(e)}")
-            # Last resort: Dummy implementation
-            print("Using dummy embeddings as last resort")
-            class DummyEmbeddings:
-                def embed_query(self, text):
-                    print("WARNING: Using dummy embeddings!")
-                    return [0.0] * 1536
-                
-                def embed_documents(self, texts):
-                    return [[0.0] * 1536 for _ in range(len(texts))]
-            
-            return DummyEmbeddings()
-
-# Add HTTP fallback function
-def http_embed_request(text, api_key):
-    """Make a direct HTTP request to OpenAI embeddings API."""
-    import requests
-    import json
-    
-    print("Using direct HTTP request for embedding")
-    url = "https://api.openai.com/v1/embeddings"
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
-    }
-    data = {
-        "model": "text-embedding-ada-002",
-        "input": text
-    }
-    
-    try:
-        response = requests.post(url, headers=headers, data=json.dumps(data))
-        if response.status_code == 200:
-            result = response.json()
-            print("Successfully got embedding via HTTP")
-            return result["data"][0]["embedding"]
-        else:
-            print(f"HTTP request failed with status {response.status_code}: {response.text}")
-            return None
-    except Exception as e:
-        print(f"HTTP request exception: {str(e)}")
-        return None
+            def embed_documents(self, texts):
+                return [[0.0] * 1536 for _ in range(len(texts))]
+        
+        return DummyEmbeddings()
 
 @st.cache_resource
 def setup_qdrant_client():
@@ -601,7 +496,7 @@ def setup_retriever():
     return QdrantRetriever()
 
 def rag_chain_node(query, run_manager):
-    """A LangGraph node for retrieval augmented generation. Returns a string."""
+    """A LangGraph node for retrieval augmented generation. Returns a string and sources."""
     print("Starting rag_chain_node...")
     # Log the query
     print(f"Query: {query}")
@@ -618,10 +513,25 @@ def rag_chain_node(query, run_manager):
     print(f"Retrieved {len(relevant_docs)} documents")
     
     # Print document sources for debugging
+    sources = []
     for i, doc in enumerate(relevant_docs):
         source = doc.metadata.get("source", "Unknown")
         page = doc.metadata.get("page", "Unknown")
         print(f"Document {i+1} source: {source}, Page: {page}")
+        
+        # Extract source information for display
+        source_path = source
+        filename = source_path.split("/")[-1] if "/" in source_path else source_path
+        
+        # Remove .pdf extension if present
+        if filename.lower().endswith('.pdf'):
+            filename = filename[:-4]
+        
+        sources.append({
+            "title": f"Ron Kohavi: {filename}",
+            "page": page,
+            "type": "pdf"
+        })
     
     # Format documents to include in the prompt
     formatted_docs = "\n\n".join([f"Document from {doc.metadata.get('source', 'Unknown')}, Page {doc.metadata.get('page', 'Unknown')}:\n{doc.page_content}" for doc in relevant_docs])
@@ -645,7 +555,7 @@ Answer:"""
     # Generate response
     response = chat_model.invoke(rag_prompt)
     print("Successfully generated response")
-    return response.content
+    return response.content, sources
 
 def evaluate_response(query, response):
     """
@@ -944,41 +854,92 @@ def execute_agent(agent, query):
 
 # Streamlit UI
 st.set_page_config(
-    page_title="AB Testing RAG Agent",
+    page_title="📊 AB Testing RAG Agent",
     page_icon="📊",
     layout="wide"
 )
 
 def main():
     """Main function for the Streamlit app."""
-    st.title("A/B Testing RAG Assistant")
-    st.write("Ask me about A/B testing concepts in the pdfs!")
+    st.title("📊 AB Testing RAG Agent")
+    st.markdown("""
+This specialized agent can answer questions about A/B Testing using a collection of Ron Kohavi's work. If it can't fully answer your A/B Testing questions using this collection, it will then automatically search Arxiv. Let's begin!
+""")
+
+    # Initialize chat history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # Display chat history
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+            
+            # Display sources if available
+            if "sources" in message and message["sources"]:
+                st.markdown("#### Sources")
+                for i, source in enumerate(message["sources"]):
+                    title = source.get("title", "Unknown")
+
+                    # Display differently based on source type
+                    if source.get("type") == "arxiv":
+                        authors = source.get("authors", "Unknown authors")
+                        st.markdown(f"**{i+1}. {title}**\nAuthors: {authors}")
+                    else:
+                        # PDF source with page number
+                        page = source.get("page", "Unknown")
+                        st.markdown(f"**{i+1}. {title}** (Page: {page})")
 
-    # Add a text input for the query
-    query = st.text_input("Ask a question:")
+    # Input for new question
+    query = st.chat_input("Ask a question about A/B Testing")
 
-    # Process the query when submitted
-    if st.button("Submit") or query:
-        if query:
-            with st.spinner("Thinking..."):
+    if query:
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": query})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.markdown(query)
+        
+        # Display assistant response
+        with st.chat_message("assistant"):
+            message_placeholder = st.empty()
+            
+            with st.status("Processing your query...", expanded=True) as status:
                 try:
-                    # First try the RAG approach
+                    # Use the RAG approach with a timeout
                     st.write("Starting with Initial RAG...")
-                    response = rag_chain_node(query, None)
+                    print("Starting RAG process for query:", query)
+                    
+                    # Step 1: Initial RAG
+                    response, sources = rag_chain_node(query, None)
                     
-                    # Display the initial response
-                    st.write(response)
+                    # Display the processed response
+                    message_placeholder.markdown(response)
                     
-                    # For debugging only
-                    print("Initial response complete")
+                    # Add assistant message to chat history
+                    st.session_state.messages.append({
+                        "role": "assistant", 
+                        "content": response,
+                        "sources": sources
+                    })
                     
+                    status.update(label="Completed!", state="complete", expanded=False)
                 except Exception as e:
-                    st.error(f"An error occurred during document retrieval or response generation: {str(e)}")
-                    print(f"Error in main function: {str(e)}")
-                    import traceback
-                    traceback.print_exc()
-        else:
-            st.write("Please enter a question.")
+                    error_msg = str(e)
+                    if "Name or service not known" in error_msg:
+                        response = "I'm having trouble connecting to the language model API due to network restrictions. The Hugging Face environment may be blocking external API calls."
+                    else:
+                        response = f"An error occurred: {error_msg}"
+                    
+                    message_placeholder.markdown(response)
+                    st.session_state.messages.append({
+                        "role": "assistant", 
+                        "content": response,
+                        "sources": []
+                    })
+                    status.update(label="Error", state="error", expanded=False)
 
 if __name__ == "__main__":
-    main() 
+    if query:
+        main()