Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| Scikit-learn Documentation Q&A Bot | |
| A Retrieval-Augmented Generation (RAG) chatbot built with Streamlit | |
| that answers questions about Scikit-learn documentation using ChromaDB | |
| for retrieval and OpenAI for generation. | |
| Author: AI Assistant | |
| Date: September 2025 | |
| """ | |
| import os | |
| import logging | |
| from typing import List, Dict, Any, Optional, Tuple | |
| import streamlit as st | |
| import chromadb | |
| from chromadb.config import Settings | |
| from sentence_transformers import SentenceTransformer | |
| from openai import OpenAI | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| class RAGChatbot: | |
| """ | |
| A Retrieval-Augmented Generation chatbot for Scikit-learn documentation. | |
| This class handles the complete RAG pipeline: retrieval from ChromaDB, | |
| augmentation with context, and generation using OpenAI's API. | |
| """ | |
| def __init__( | |
| self, | |
| db_path: str = './chroma_db', | |
| collection_name: str = 'sklearn_docs', | |
| embedding_model_name: str = 'all-MiniLM-L6-v2' | |
| ): | |
| """ | |
| Initialize the RAG chatbot. | |
| Args: | |
| db_path (str): Path to ChromaDB database | |
| collection_name (str): Name of the ChromaDB collection | |
| embedding_model_name (str): Name of the embedding model | |
| """ | |
| self.db_path = db_path | |
| self.collection_name = collection_name | |
| self.embedding_model_name = embedding_model_name | |
| # Initialize components | |
| self.chroma_client = None | |
| self.collection = None | |
| self.embedding_model = None | |
| self.openai_client = None | |
| # Initialize the retrieval system | |
| self._initialize_retrieval_system() | |
| def _initialize_retrieval_system(self) -> None: | |
| """ | |
| Initialize ChromaDB client and embedding model for retrieval. | |
| """ | |
| try: | |
| # Initialize ChromaDB client | |
| self.chroma_client = chromadb.PersistentClient( | |
| path=self.db_path, | |
| settings=Settings(anonymized_telemetry=False) | |
| ) | |
| # Get collection | |
| self.collection = self.chroma_client.get_collection( | |
| name=self.collection_name | |
| ) | |
| # Load embedding model (same as used for building the database) | |
| self.embedding_model = SentenceTransformer(self.embedding_model_name) | |
| logger.info("RAG retrieval system initialized successfully") | |
| except Exception as e: | |
| logger.error(f"Failed to initialize retrieval system: {e}") | |
| raise | |
| def set_openai_client(self, api_key: str) -> bool: | |
| """ | |
| Initialize OpenAI client with API key. | |
| Args: | |
| api_key (str): OpenAI API key | |
| Returns: | |
| bool: True if successful, False otherwise | |
| """ | |
| try: | |
| self.openai_client = OpenAI(api_key=api_key) | |
| # Test the API key with a simple request | |
| self.openai_client.models.list() | |
| logger.info("OpenAI client initialized successfully") | |
| return True | |
| except Exception as e: | |
| logger.error(f"Failed to initialize OpenAI client: {e}") | |
| st.error(f"Invalid API key or OpenAI connection error: {e}") | |
| return False | |
| def retrieve_relevant_chunks( | |
| self, | |
| query: str, | |
| n_results: int = 3, | |
| min_relevance_score: float = 0.1 | |
| ) -> List[Dict[str, Any]]: | |
| """ | |
| Retrieve relevant text chunks from the vector database. | |
| Args: | |
| query (str): User question/query | |
| n_results (int): Number of chunks to retrieve | |
| min_relevance_score (float): Minimum relevance score threshold | |
| Returns: | |
| List[Dict[str, Any]]: Retrieved chunks with content and metadata | |
| """ | |
| try: | |
| # Query the collection | |
| results = self.collection.query( | |
| query_texts=[query], | |
| n_results=n_results | |
| ) | |
| retrieved_chunks = [] | |
| # Process results | |
| if results['documents'] and results['documents'][0]: | |
| for i in range(len(results['documents'][0])): | |
| chunk_data = { | |
| 'content': results['documents'][0][i], | |
| 'metadata': results['metadatas'][0][i], | |
| 'distance': results['distances'][0][i] if 'distances' in results else None | |
| } | |
| # Filter by relevance score if available | |
| if chunk_data['distance'] is None or chunk_data['distance'] >= min_relevance_score: | |
| retrieved_chunks.append(chunk_data) | |
| logger.info(f"Retrieved {len(retrieved_chunks)} relevant chunks for query: {query[:50]}...") | |
| return retrieved_chunks | |
| except Exception as e: | |
| logger.error(f"Error retrieving chunks: {e}") | |
| st.error(f"Error during retrieval: {e}") | |
| return [] | |
| def create_rag_prompt( | |
| self, | |
| user_question: str, | |
| retrieved_chunks: List[Dict[str, Any]] | |
| ) -> str: | |
| """ | |
| Create an augmented prompt for OpenAI with retrieved context. | |
| Args: | |
| user_question (str): Original user question | |
| retrieved_chunks (List[Dict[str, Any]]): Retrieved relevant chunks | |
| Returns: | |
| str: Augmented prompt for OpenAI | |
| """ | |
| # Build context from retrieved chunks | |
| context_parts = [] | |
| for i, chunk in enumerate(retrieved_chunks, 1): | |
| url = chunk['metadata'].get('url', 'Unknown source') | |
| content = chunk['content'].strip() | |
| context_part = f"--- Context {i} (Source: {url}) ---\n{content}\n" | |
| context_parts.append(context_part) | |
| context = "\n".join(context_parts) | |
| # Create the RAG prompt | |
| rag_prompt = f"""You are an expert AI assistant specializing in Scikit-learn, a popular Python machine learning library. Your task is to answer questions about Scikit-learn based ONLY on the provided context from the official documentation. | |
| CONTEXT: | |
| {context} | |
| USER QUESTION: | |
| {user_question} | |
| INSTRUCTIONS: | |
| 1. Answer the question using ONLY the information provided in the context above | |
| 2. Be accurate, helpful, and specific | |
| 3. If the context doesn't contain enough information to fully answer the question, say so clearly | |
| 4. Include relevant code examples if they appear in the context | |
| 5. Mention specific function names, class names, or parameter names when relevant | |
| 6. Structure your answer clearly with appropriate formatting | |
| ANSWER:""" | |
| return rag_prompt | |
| def generate_answer( | |
| self, | |
| prompt: str, | |
| model: str = "gpt-3.5-turbo", | |
| max_tokens: int = 1000, | |
| temperature: float = 0.1 | |
| ) -> Optional[str]: | |
| """ | |
| Generate answer using OpenAI API. | |
| Args: | |
| prompt (str): Augmented prompt with context | |
| model (str): OpenAI model to use | |
| max_tokens (int): Maximum tokens in response | |
| temperature (float): Temperature for generation | |
| Returns: | |
| Optional[str]: Generated answer or None if failed | |
| """ | |
| try: | |
| response = self.openai_client.chat.completions.create( | |
| model=model, | |
| messages=[ | |
| { | |
| "role": "system", | |
| "content": "You are a helpful AI assistant specializing in Scikit-learn documentation. Provide accurate, helpful answers based only on the provided context." | |
| }, | |
| { | |
| "role": "user", | |
| "content": prompt | |
| } | |
| ], | |
| max_tokens=max_tokens, | |
| temperature=temperature, | |
| top_p=0.9 | |
| ) | |
| answer = response.choices[0].message.content.strip() | |
| logger.info(f"Generated answer of length: {len(answer)}") | |
| return answer | |
| except Exception as e: | |
| logger.error(f"Error generating answer: {e}") | |
| st.error(f"Error generating answer: {e}") | |
| return None | |
| def get_answer( | |
| self, | |
| user_question: str, | |
| n_chunks: int = 3, | |
| model: str = "gpt-3.5-turbo" | |
| ) -> Tuple[Optional[str], List[str]]: | |
| """ | |
| Complete RAG pipeline: retrieve, augment, generate. | |
| Args: | |
| user_question (str): User's question | |
| n_chunks (int): Number of chunks to retrieve | |
| model (str): OpenAI model to use | |
| Returns: | |
| Tuple[Optional[str], List[str]]: Generated answer and source URLs | |
| """ | |
| if not self.openai_client: | |
| st.error("OpenAI client not initialized. Please provide a valid API key.") | |
| return None, [] | |
| # Use a single spinner for the entire process to prevent flickering | |
| with st.spinner("🤖 Generating answer..."): | |
| # Step 1: Retrieve relevant chunks | |
| retrieved_chunks = self.retrieve_relevant_chunks(user_question, n_chunks) | |
| if not retrieved_chunks: | |
| return "I couldn't find relevant information in the Scikit-learn documentation to answer your question. Please try rephrasing your question or ask about a different topic.", [] | |
| # Step 2: Create augmented prompt | |
| rag_prompt = self.create_rag_prompt(user_question, retrieved_chunks) | |
| # Step 3: Generate answer | |
| answer = self.generate_answer(rag_prompt, model) | |
| # Extract source URLs | |
| source_urls = [chunk['metadata'].get('url', 'Unknown') for chunk in retrieved_chunks] | |
| source_urls = list(dict.fromkeys(source_urls)) # Remove duplicates while preserving order | |
| return answer, source_urls | |
| def initialize_session_state(): | |
| """Initialize Streamlit session state variables.""" | |
| if 'chatbot' not in st.session_state: | |
| try: | |
| st.session_state.chatbot = RAGChatbot() | |
| except Exception as e: | |
| st.error(f"Failed to initialize chatbot: {e}") | |
| st.stop() | |
| if 'openai_initialized' not in st.session_state: | |
| st.session_state.openai_initialized = False | |
| if 'chat_history' not in st.session_state: | |
| st.session_state.chat_history = [] | |
| def main(): | |
| """Main Streamlit application.""" | |
| # Page configuration | |
| st.set_page_config( | |
| page_title="Scikit-learn Q&A Bot", | |
| page_icon="🤖", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| # Initialize session state | |
| initialize_session_state() | |
| # Main title and description | |
| st.title("🤖 Scikit-learn Documentation Q&A Bot") | |
| st.markdown(""" | |
| Welcome to the **Scikit-learn Documentation Q&A Bot**! This intelligent assistant can answer your questions about Scikit-learn using the official documentation. | |
| **How it works:** | |
| 1. 🔍 **Retrieval**: Searches through 1,249+ documentation chunks | |
| 2. 📝 **Augmentation**: Provides relevant context to the AI | |
| 3. 🤖 **Generation**: Uses OpenAI to generate accurate answers | |
| """) | |
| # Sidebar for API key and settings | |
| with st.sidebar: | |
| st.header("⚙️ Configuration") | |
| # OpenAI API Key input | |
| api_key = st.text_input( | |
| "🔑 OpenAI API Key", | |
| type="password", | |
| placeholder="sk-...", | |
| help="Enter your OpenAI API key to enable the chatbot" | |
| ) | |
| if api_key and not st.session_state.openai_initialized: | |
| if st.session_state.chatbot.set_openai_client(api_key): | |
| st.session_state.openai_initialized = True | |
| st.success("✅ API key validated!") | |
| st.rerun() | |
| # Model selection | |
| model = st.selectbox( | |
| "🧠 AI Model", | |
| ["gpt-3.5-turbo", "gpt-4", "gpt-4-turbo-preview"], | |
| index=0, | |
| help="Choose the OpenAI model for generating answers" | |
| ) | |
| # Number of context chunks | |
| n_chunks = st.slider( | |
| "📄 Context Chunks", | |
| min_value=1, | |
| max_value=5, | |
| value=3, | |
| help="Number of relevant documentation chunks to use for context" | |
| ) | |
| st.markdown("---") | |
| # Database info | |
| st.header("📊 Database Info") | |
| try: | |
| collection_count = st.session_state.chatbot.collection.count() | |
| st.metric("Total Documents", f"{collection_count:,}") | |
| st.metric("Embedding Model", "all-MiniLM-L6-v2") | |
| st.metric("Vector Dimensions", "384") | |
| except: | |
| st.error("Could not load database info") | |
| st.markdown("---") | |
| # Clear chat history | |
| if st.button("🗑️ Clear Chat History"): | |
| st.session_state.chat_history = [] | |
| st.rerun() | |
| # Main chat interface | |
| col1, col2 = st.columns([2, 1]) | |
| with col1: | |
| st.header("💬 Ask Your Question") | |
| # Question input with better state management | |
| # Initialize question input key if it doesn't exist | |
| if 'question_input' not in st.session_state: | |
| st.session_state.question_input = '' | |
| # Handle selected question from examples | |
| if 'selected_question' in st.session_state: | |
| st.session_state.question_input = st.session_state['selected_question'] | |
| # Clear after setting to prevent re-triggering | |
| del st.session_state['selected_question'] | |
| user_question = st.text_input( | |
| "Enter your question about Scikit-learn:", | |
| placeholder="e.g., How do I perform cross-validation in scikit-learn?", | |
| key="question_input" | |
| ) | |
| # Submit button | |
| submit_button = st.button("🚀 Get Answer", type="primary") | |
| # Process question | |
| if submit_button and user_question: | |
| if not st.session_state.openai_initialized: | |
| st.error("⚠️ Please enter a valid OpenAI API key in the sidebar first.") | |
| else: | |
| # Get answer using RAG | |
| answer, sources = st.session_state.chatbot.get_answer( | |
| user_question, n_chunks, model | |
| ) | |
| if answer: | |
| # Add to chat history | |
| st.session_state.chat_history.append({ | |
| 'question': user_question, | |
| 'answer': answer, | |
| 'sources': sources | |
| }) | |
| # Success message | |
| st.success("✅ Answer generated successfully! Check the chat history below.") | |
| # Display chat history | |
| if st.session_state.chat_history: | |
| st.header("📝 Chat History") | |
| for i, chat in enumerate(reversed(st.session_state.chat_history)): | |
| with st.expander(f"Q: {chat['question'][:60]}...", expanded=(i == 0)): | |
| st.markdown(f"**Question:** {chat['question']}") | |
| st.markdown(f"**Answer:**\n{chat['answer']}") | |
| if chat['sources']: | |
| st.markdown("**Sources:**") | |
| for j, source in enumerate(chat['sources'], 1): | |
| source_name = source.split('/')[-1] if '/' in source else source | |
| st.markdown(f"{j}. [{source_name}]({source})") | |
| with col2: | |
| st.header("💡 Example Questions") | |
| example_questions = [ | |
| "How do I perform cross-validation in scikit-learn?", | |
| "What is the difference between Ridge and Lasso regression?", | |
| "How do I use GridSearchCV for parameter tuning?", | |
| "What clustering algorithms are available in scikit-learn?", | |
| "How do I preprocess data using StandardScaler?", | |
| "What is the difference between classification and regression?", | |
| "How do I handle missing values in my dataset?", | |
| "What is feature selection and how do I use it?", | |
| "How do I visualize decision trees?", | |
| "What is ensemble learning in scikit-learn?" | |
| ] | |
| for question in example_questions: | |
| if st.button(question, key=f"example_{hash(question)}"): | |
| # Set the question and rerun only once | |
| st.session_state['selected_question'] = question | |
| st.rerun() | |
| st.markdown("---") | |
| st.header("ℹ️ Tips") | |
| st.markdown(""" | |
| **For best results:** | |
| - Be specific in your questions | |
| - Ask about scikit-learn functionality | |
| - Include context when possible | |
| - Check the sources for verification | |
| **The bot can help with:** | |
| - API usage and parameters | |
| - Algorithm explanations | |
| - Code examples | |
| - Best practices | |
| - Troubleshooting | |
| """) | |
| if __name__ == "__main__": | |
| # Run the main application | |
| # Note: For deployment environments like HuggingFace Spaces, | |
| # Streamlit warnings about missing ScriptRunContext can be safely ignored | |
| main() |