Create app.py

app.py

@@ -0,0 +1,1424 @@
+import os
+import json
+import time
+import requests
+import numpy as np
+import pandas as pd
+from datetime import datetime
+from typing import Dict, List, Any, Optional, Tuple
+import gradio as gr
+from dotenv import load_dotenv
+
+# Vector DB and embedding imports
+from langchain.vectorstores import FAISS
+from langchain_openai import OpenAIEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.schema import Document
+from langchain_openai import ChatOpenAI
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+
+# Visualization imports
+import plotly.graph_objects as go
+from sklearn.manifold import TSNE
+
+# Load environment variables
+load_dotenv()
+
+# Check if OPENAI_API_KEY is set
+OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
+if not OPENAI_API_KEY:
+    print("⚠️ Warning: OPENAI_API_KEY not found in environment variables.")
+
+# Configuration
+DEFAULT_DATASET_ID = "2457ea29-fc82-48b0-86ec-3b0755de7515"
+DEFAULT_MODEL = "gpt-4o-mini"
+API_BASE_URL = "https://data.cms.gov/data-api/v1"
+INITIAL_SAMPLE_SIZE = 100  # Start with a small sample
+
+# Dataset version mapping
+DATASET_VERSIONS = {
+    # 2025 Data
+    "Q1 2025": "74edb053-bd01-40a0-91a0-4961c1fe6281",
+    
+    # 2024 Data
+    "Q1 2024": "6d6e0e8d-64cf-43fb-9ba8-e2ad9b9bb21e",
+    "Q2 2024": "04405289-5635-4b2a-a64f-c4b6415ab6ff",
+    "Q3 2024": "e87f09c2-5ff7-4ddf-b60c-6130995b15cf",
+    "Q4 2024": "e9d278e4-90e8-47ab-9c5b-af2ca64bf352",
+    
+    # 2023 Data
+    "Q1 2023": "0b6caf2f-8948-4603-922e-d7f0c52c0a45",
+    "Q2 2023": "46339a0c-0f07-40ed-8975-ddb387c367a4",
+    "Q3 2023": "70efac57-6093-4e1d-ad6a-36f8261f53eb",
+    "Q4 2023": "1df8331a-ed44-41ec-971f-158349658949",
+    
+    # 2022 Data
+    "Q1 2022": "5b678653-aa36-455b-9144-1d073ef7991b",
+    
+    # 2021 Data
+    "Q1 2021": "7b409bba-ca00-426e-9493-1dc10e5340cc",
+    
+    # 2020 Data
+    "Q1 2020": "3870b29c-4312-4fb1-a956-71c148ae5b50",
+    
+    # 2019 Data
+    "Q1 2019": "017e6ab7-7e19-4e98-b4fa-30578b47e578",
+    "Q4 2019": "2c209bdb-ed0c-42e0-b027-8a97024b8035"
+}
+
+# US States for reference
+US_STATES = [
+    "", "AL", "AK", "AZ", "AR", "CA", "CO", "CT", "DE", "FL", "GA", 
+    "HI", "ID", "IL", "IN", "IA", "KS", "KY", "LA", "ME", "MD", 
+    "MA", "MI", "MN", "MS", "MO", "MT", "NE", "NV", "NH", "NJ", 
+    "NM", "NY", "NC", "ND", "OH", "OK", "OR", "PA", "RI", "SC", 
+    "SD", "TN", "TX", "UT", "VT", "VA", "WA", "WV", "WI", "WY",
+    "DC", "PR", "VI"
+]
+
+# State names mapping for better UI
+STATE_NAMES = {
+    "": "All States",
+    "AL": "Alabama", "AK": "Alaska", "AZ": "Arizona", "AR": "Arkansas", 
+    "CA": "California", "CO": "Colorado", "CT": "Connecticut", "DE": "Delaware", 
+    "FL": "Florida", "GA": "Georgia", "HI": "Hawaii", "ID": "Idaho", 
+    "IL": "Illinois", "IN": "Indiana", "IA": "Iowa", "KS": "Kansas", 
+    "KY": "Kentucky", "LA": "Louisiana", "ME": "Maine", "MD": "Maryland", 
+    "MA": "Massachusetts", "MI": "Michigan", "MN": "Minnesota", "MS": "Mississippi", 
+    "MO": "Missouri", "MT": "Montana", "NE": "Nebraska", "NV": "Nevada", 
+    "NH": "New Hampshire", "NJ": "New Jersey", "NM": "New Mexico", "NY": "New York", 
+    "NC": "North Carolina", "ND": "North Dakota", "OH": "Ohio", "OK": "Oklahoma", 
+    "OR": "Oregon", "PA": "Pennsylvania", "RI": "Rhode Island", "SC": "South Carolina", 
+    "SD": "South Dakota", "TN": "Tennessee", "TX": "Texas", "UT": "Utah", 
+    "VT": "Vermont", "VA": "Virginia", "WA": "Washington", "WV": "West Virginia", 
+    "WI": "Wisconsin", "WY": "Wyoming", "DC": "District of Columbia", 
+    "PR": "Puerto Rico", "VI": "Virgin Islands"
+}
+
+# Dictionary to store multiple datasets
+rag_systems = {}
+current_dataset_key = None
+
+# Gradio theme configuration
+theme = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="gray",
+    neutral_hue="slate",
+    font=gr.themes.GoogleFont("Inter")
+)
+
+def query_cms_api(version_id, state_filter="", max_records=100):
+    """Query the CMS API with pagination."""
+    url = f"{API_BASE_URL}/dataset/{version_id}/data"
+    all_records = []
+    offset = 0
+    page_size = min(max_records, 100)  # Page size, max 100
+    
+    # Set up filter parameters
+    params = {
+        'size': page_size,
+        'offset': 0
+    }
+    
+    # Add state filter if provided
+    if state_filter and state_filter != "":
+        params[f'filter[STATE_CD]'] = state_filter
+    
+    progress_text = f"Querying CMS API...\n"
+    
+    # Fetch data with pagination
+    while len(all_records) < max_records:
+        params['offset'] = offset
+        
+        try:
+            response = requests.get(url, params=params)
+            
+            if response.status_code != 200:
+                error_msg = f"Error: Status {response.status_code}"
+                return [], error_msg
+            
+            # Parse the response - the API returns a list directly
+            records = response.json()
+            
+            if not records or not isinstance(records, list):
+                if len(all_records) == 0:
+                    return [], "No records found"
+                break
+            
+            progress_text += f"Retrieved {len(records)} records (offset: {offset})\n"
+            all_records.extend(records)
+            
+            # If we got fewer records than requested, we've reached the end
+            if len(records) < page_size:
+                break
+            
+            # Move to next page
+            offset += len(records)
+            
+            # Add delay to be nice to the API
+            time.sleep(0.5)
+            
+        except Exception as e:
+            error_msg = f"Error querying API: {str(e)}"
+            return [], error_msg
+    
+    final_records = all_records[:max_records]
+    success_msg = f"Successfully retrieved {len(final_records)} records"
+    
+    return final_records, success_msg
+
+def process_records(records, version):
+    """Process CMS API records into documents for the RAG system."""
+    # Parse version into quarter and year
+    quarter = "Unknown"
+    year = "Unknown"
+    if ' ' in version:
+        parts = version.split(' ')
+        if len(parts) == 2:
+            quarter, year = parts
+    
+    embeddings = OpenAIEmbeddings()
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    
+    # Convert records to documents
+    documents = []
+    
+    for record in records:
+        # Format the record as text with explicit time information
+        content = [f"Medicare Provider Data from {quarter} {year}"]
+        content.append(f"Time Period: {quarter} of {year}")
+        
+        # Add all fields from the record
+        for key, value in record.items():
+            if value is not None and value != "":
+                content.append(f"{key}: {value}")
+        
+        text = "\n".join(content)
+        
+        # Create metadata with explicit time fields
+        metadata = {
+            'dataset_version': version,
+            'quarter': quarter,
+            'year': year,
+            'record_id': record.get('ENRLMT_ID', 'unknown')
+        }
+        
+        # Add all fields to metadata for better searchability
+        for key, value in record.items():
+            if value is not None and value != "":
+                try:
+                    # Convert complex values to strings to avoid serialization issues
+                    if not isinstance(value, (str, int, float, bool, type(None))):
+                        metadata[key] = str(value)
+                    else:
+                        metadata[key] = value
+                except:
+                    # If there's any issue, convert to string
+                    metadata[key] = str(value)
+        
+        documents.append(Document(page_content=text, metadata=metadata))
+    
+    # Chunk documents
+    chunks = text_splitter.split_documents(documents)
+    
+    # Create vector store
+    vector_store = FAISS.from_documents(chunks, embeddings)
+    
+    return vector_store, len(documents), len(chunks)
+
+def create_progress_callback():
+    """Create a progress callback for long-running operations."""
+    def callback(message):
+        # In a real Gradio app, this would update a progress bar
+        print(f"Progress: {message}")
+    return callback
+
+def validate_api_key():
+    """Validate that the OpenAI API key is set."""
+    api_key = os.getenv('OPENAI_API_KEY')
+    if not api_key:
+        return False, "OpenAI API key not found. Please set it in your environment variables or .env file."
+    return True, "API key validated successfully."
+
+def get_dataset_summary(rag_systems):
+    """Generate a summary of all loaded datasets."""
+    if not rag_systems:
+        return "No datasets currently loaded."
+    
+    summary_lines = ["### Currently Loaded Datasets:\n"]
+    
+    for i, (key, system) in enumerate(rag_systems.items(), 1):
+        meta = system['metadata']
+        summary_lines.append(
+            f"{i}. **{meta['dataset_version']}** - "
+            f"State: {meta['state_filter']} - "
+            f"Records: {meta['record_count']} - "
+            f"Chunks: {meta['chunk_count']}"
+        )
+        
+        if key == current_dataset_key:
+            summary_lines[-1] += " *(Current)*"
+    
+    summary_lines.append(f"\n**Total datasets loaded:** {len(rag_systems)}")
+    
+    return "\n".join(summary_lines)
+
+def format_state_options():
+    """Format state options for Gradio dropdown."""
+    options = []
+    for code in US_STATES:
+        if code == "":
+            options.append(("All States", ""))
+        else:
+            options.append((f"{STATE_NAMES[code]} ({code})", code))
+    return options
+
+def load_dataset_gradio(version, state_filter, max_records, use_sample):
+    """Load data from CMS API and set up the RAG system - Gradio version."""
+    global rag_systems, current_dataset_key
+    
+    # Validate API key first
+    valid, message = validate_api_key()
+    if not valid:
+        return message, get_dataset_summary(rag_systems)
+    
+    # Generate a unique key for this dataset
+    dataset_key = f"{version}_{state_filter}_{max_records}"
+    
+    # Check if dataset already loaded
+    if dataset_key in rag_systems:
+        current_dataset_key = dataset_key
+        return f"✅ Dataset already loaded and set as current: {version} - {STATE_NAMES.get(state_filter, 'All States')}", get_dataset_summary(rag_systems)
+    
+    # Get version ID
+    version_id = DATASET_VERSIONS.get(version)
+    if not version_id:
+        return f"❌ Invalid version: {version}", get_dataset_summary(rag_systems)
+    
+    # Adjust max records if sample
+    actual_max = INITIAL_SAMPLE_SIZE if use_sample else max_records
+    
+    # Status message
+    status_msg = f"🔄 Loading {version} data"
+    if state_filter:
+        status_msg += f" for {STATE_NAMES.get(state_filter, state_filter)}"
+    status_msg += f" (max {actual_max} records)..."
+    
+    try:
+        # Fetch data from API
+        records, api_message = query_cms_api(version_id, state_filter, actual_max)
+        
+        if not records:
+            return f"❌ Failed to load data: {api_message}", get_dataset_summary(rag_systems)
+        
+        status_msg += f"\n✅ {api_message}"
+        
+        # Process records and create vector store
+        status_msg += "\n🔄 Processing records and creating vector store..."
+        vector_store, doc_count, chunk_count = process_records(records, version)
+        
+        # Set up RAG system
+        llm = ChatOpenAI(temperature=0.7, model_name=DEFAULT_MODEL)
+        memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)
+        retriever = vector_store.as_retriever()
+        
+        conversation_chain = ConversationalRetrievalChain.from_llm(
+            llm=llm,
+            retriever=retriever,
+            memory=memory
+        )
+        
+        # Store in the dictionary
+        rag_systems[dataset_key] = {
+            'vector_store': vector_store,
+            'conversation_chain': conversation_chain,
+            'metadata': {
+                'dataset_version': version,
+                'version_id': version_id,
+                'state_filter': STATE_NAMES.get(state_filter, "All States") if state_filter else "All States",
+                'record_count': len(records),
+                'document_count': doc_count,
+                'chunk_count': chunk_count,
+                'loaded_at': datetime.now().isoformat()
+            }
+        }
+        
+        # Set as current dataset
+        current_dataset_key = dataset_key
+        
+        success_msg = f"✅ Successfully loaded {version} - {STATE_NAMES.get(state_filter, 'All States')}\n"
+        success_msg += f"📊 Created {chunk_count} chunks from {len(records)} records"
+        
+        return success_msg, get_dataset_summary(rag_systems)
+        
+    except Exception as e:
+        error_msg = f"❌ Error loading data: {str(e)}"
+        return error_msg, get_dataset_summary(rag_systems)
+
+def switch_dataset_gradio(dataset_index):
+    """Switch to a different dataset - Gradio version."""
+    global rag_systems, current_dataset_key
+    
+    if not rag_systems:
+        return "❌ No datasets loaded.", get_dataset_summary(rag_systems)
+    
+    if not dataset_index:
+        return "❌ Please select a dataset.", get_dataset_summary(rag_systems)
+    
+    try:
+        # Parse the index from the selection (format: "1. Dataset Name")
+        index = int(dataset_index.split(".")[0])
+        
+        if 1 <= index <= len(rag_systems):
+            key = list(rag_systems.keys())[index - 1]
+            current_dataset_key = key
+            meta = rag_systems[key]['metadata']
+            return f"✅ Switched to: {meta['dataset_version']} - {meta['state_filter']}", get_dataset_summary(rag_systems)
+        else:
+            return f"❌ Invalid selection.", get_dataset_summary(rag_systems)
+    except:
+        return "❌ Invalid selection format.", get_dataset_summary(rag_systems)
+
+def remove_dataset_gradio(dataset_index):
+    """Remove a dataset from memory - Gradio version."""
+    global rag_systems, current_dataset_key
+    
+    if not rag_systems:
+        return "❌ No datasets loaded.", get_dataset_summary(rag_systems)
+    
+    if not dataset_index:
+        return "❌ Please select a dataset to remove.", get_dataset_summary(rag_systems)
+    
+    try:
+        # Parse the index from the selection
+        index = int(dataset_index.split(".")[0])
+        
+        if 1 <= index <= len(rag_systems):
+            key = list(rag_systems.keys())[index - 1]
+            meta = rag_systems[key]['metadata']
+            
+            # Remove the dataset
+            del rag_systems[key]
+            
+            # If this was the current dataset, clear the current key
+            if key == current_dataset_key:
+                current_dataset_key = None
+                # Set another dataset as current if available
+                if rag_systems:
+                    current_dataset_key = list(rag_systems.keys())[0]
+            
+            return f"✅ Removed: {meta['dataset_version']} - {meta['state_filter']}", get_dataset_summary(rag_systems)
+        else:
+            return f"❌ Invalid selection.", get_dataset_summary(rag_systems)
+    except Exception as e:
+        return f"❌ Error removing dataset: {str(e)}", get_dataset_summary(rag_systems)
+
+def get_dataset_choices():
+    """Get formatted dataset choices for Gradio dropdown."""
+    if not rag_systems:
+        return []
+    
+    choices = []
+    for i, (key, system) in enumerate(rag_systems.items(), 1):
+        meta = system['metadata']
+        choice_text = f"{i}. {meta['dataset_version']} - {meta['state_filter']} ({meta['record_count']} records)"
+        if key == current_dataset_key:
+            choice_text += " [CURRENT]"
+        choices.append(choice_text)
+    
+    return choices
+
+def clear_all_datasets_gradio():
+    """Clear all loaded datasets - Gradio version."""
+    global rag_systems, current_dataset_key
+    
+    if not rag_systems:
+        return "ℹ️ No datasets to clear.", ""
+    
+    count = len(rag_systems)
+    rag_systems.clear()
+    current_dataset_key = None
+    
+    return f"✅ Cleared {count} dataset(s) from memory.", ""
+
+def get_current_dataset_info():
+    """Get information about the current dataset."""
+    global rag_systems, current_dataset_key
+    
+    if not current_dataset_key or current_dataset_key not in rag_systems:
+        return "No dataset currently selected."
+    
+    meta = rag_systems[current_dataset_key]['metadata']
+    info = f"**Current Dataset:** {meta['dataset_version']} - {meta['state_filter']}\n"
+    info += f"- Records: {meta['record_count']}\n"
+    info += f"- Chunks: {meta['chunk_count']}\n"
+    info += f"- Loaded: {meta['loaded_at'][:19]}"
+    
+    return info
+
+def ask_question_gradio(question, chat_history):
+    """Ask a question to the current dataset - Gradio version."""
+    global rag_systems, current_dataset_key
+    
+    if not current_dataset_key or current_dataset_key not in rag_systems:
+        response = "❌ No dataset selected. Please load a dataset first."
+        chat_history.append((question, response))
+        return "", chat_history
+    
+    # Get the dataset
+    system = rag_systems[current_dataset_key]
+    meta = system['metadata']
+    
+    try:
+        # Use the chain to get a response
+        result = system['conversation_chain'].invoke({"question": question})
+        answer = result["answer"]
+        
+        # Add dataset source information
+        answer += f"\n\n*Source: {meta['dataset_version']} - {meta['state_filter']} ({meta['record_count']} records)*"
+        
+        # Update chat history
+        chat_history.append((question, answer))
+        
+        return "", chat_history
+    
+    except Exception as e:
+        error_response = f"❌ Error processing query: {str(e)}"
+        chat_history.append((question, error_response))
+        return "", chat_history
+
+def ask_global_question_gradio(question, chat_history):
+    """Ask a question that might require knowledge from all loaded datasets."""
+    global rag_systems
+    
+    if not rag_systems:
+        response = "❌ No datasets loaded. Please load datasets first."
+        chat_history.append((question, response))
+        return "", chat_history
+    
+    # Check if this is a global question about the datasets themselves
+    global_keywords = ['how many', 'which years', 'what years', 'what quarters', 'how many years',
+                      'which quarters', 'time period', 'date range', 'all datasets', 'datasets',
+                      'compare', 'comparison', 'difference', 'trend', 'over time']
+    
+    is_global_question = any(keyword in question.lower() for keyword in global_keywords)
+    
+    # Check if the question mentions a specific state
+    mentioned_state = None
+    question_lower = question.lower()
+    
+    # Check for state names
+    for code, name in STATE_NAMES.items():
+        if code and (code.lower() in question_lower or name.lower() in question_lower):
+            mentioned_state = code
+            break
+    
+    try:
+        if mentioned_state and not is_global_question:
+            # Find all datasets for that state
+            suitable_datasets = []
+            
+            for key, system in rag_systems.items():
+                meta = system['metadata']
+                state_filter = meta['state_filter']
+                
+                # Check if this dataset matches the mentioned state
+                if mentioned_state in state_filter or STATE_NAMES[mentioned_state] in state_filter:
+                    suitable_datasets.append(key)
+            
+            if suitable_datasets:
+                response = f"🔄 Found {len(suitable_datasets)} dataset(s) for {STATE_NAMES[mentioned_state]}:\n\n"
+                
+                # Query each suitable dataset
+                all_results = []
+                for dataset_key in suitable_datasets:
+                    system = rag_systems[dataset_key]
+                    meta = system['metadata']
+                    
+                    try:
+                        result = system['conversation_chain'].invoke({"question": question})
+                        answer = result["answer"]
+                        all_results.append({
+                            'dataset': f"{meta['dataset_version']} - {meta['state_filter']}",
+                            'answer': answer
+                        })
+                    except Exception as e:
+                        all_results.append({
+                            'dataset': f"{meta['dataset_version']} - {meta['state_filter']}",
+                            'answer': f"Error: {str(e)}"
+                        })
+                
+                # Format combined response
+                for result in all_results:
+                    response += f"**{result['dataset']}**\n{result['answer']}\n\n---\n\n"
+                
+                chat_history.append((question, response))
+                return "", chat_history
+            else:
+                response = f"ℹ️ No datasets found for {STATE_NAMES[mentioned_state]}. Please load data for this state first."
+                chat_history.append((question, response))
+                return "", chat_history
+        
+        elif is_global_question:
+            # Create a summary of all available datasets
+            dataset_summary = generate_dataset_metadata_summary()
+            
+            # Create a system message that includes this metadata
+            llm = ChatOpenAI(temperature=0.7, model_name=DEFAULT_MODEL)
+            
+            system_message = f"""You are an expert on Medicare Provider data. You have access to multiple datasets spanning different quarters and years.
+
+{dataset_summary}
+
+When answering questions, consider the metadata about all available datasets. For questions about time periods, years, quarters, or trends, use the information about which datasets are loaded."""
+            
+            messages = [
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": question}
+            ]
+            
+            response = llm.invoke(messages)
+            answer = response.content
+            
+            chat_history.append((question, answer))
+            return "", chat_history
+        
+        else:
+            # For non-global questions without specific state mention, use the current dataset
+            return ask_question_gradio(question, chat_history)
+            
+    except Exception as e:
+        error_response = f"❌ Error processing global query: {str(e)}"
+        chat_history.append((question, error_response))
+        return "", chat_history
+
+def generate_dataset_metadata_summary():
+    """Generate a detailed summary of dataset metadata."""
+    if not rag_systems:
+        return "No datasets loaded."
+    
+    summary = "# Available Datasets\n\n"
+    summary += "The following datasets are currently loaded:\n\n"
+    
+    # Group by year
+    years = set()
+    quarters_by_year = {}
+    states = set()
+    
+    for key, system in rag_systems.items():
+        meta = system['metadata']
+        version = meta['dataset_version']
+        state = meta['state_filter']
+        
+        # Extract year from version (e.g., "Q1 2025" -> "2025")
+        if ' ' in version:
+            year = version.split(' ')[1]
+            quarter = version.split(' ')[0]
+            
+            years.add(year)
+            states.add(state)
+            
+            if year not in quarters_by_year:
+                quarters_by_year[year] = set()
+            
+            quarters_by_year[year].add(quarter)
+    
+    # Format the summary
+    summary += "## Years Available\n"
+    summary += ", ".join(sorted(list(years))) + "\n\n"
+    
+    summary += "## Quarters Available by Year\n"
+    for year in sorted(quarters_by_year.keys()):
+        summary += f"- {year}: {', '.join(sorted(list(quarters_by_year[year])))}\n"
+    
+    summary += "\n## States Available\n"
+    summary += ", ".join(sorted(list(states))) + "\n\n"
+    
+    summary += "## Full Dataset List\n"
+    for key, system in rag_systems.items():
+        meta = system['metadata']
+        summary += f"- {meta['dataset_version']} - {meta['state_filter']} ({meta['record_count']} records)\n"
+    
+    return summary
+
+def compare_datasets_gradio(question, dataset_indices):
+    """Compare multiple datasets by asking the same question - Gradio version."""
+    global rag_systems
+    
+    if not rag_systems:
+        return "❌ No datasets loaded. Please load datasets first."
+    
+    if not dataset_indices or len(dataset_indices) < 2:
+        return "❌ Please select at least 2 datasets to compare."
+    
+    # Parse indices and get dataset keys
+    selected_keys = []
+    for selection in dataset_indices:
+        try:
+            index = int(selection.split(".")[0])
+            if 1 <= index <= len(rag_systems):
+                key = list(rag_systems.keys())[index - 1]
+                selected_keys.append(key)
+        except:
+            continue
+    
+    if len(selected_keys) < 2:
+        return "❌ Could not parse selected datasets."
+    
+    comparison_result = f"# Comparison: {question}\n\n"
+    
+    # Query each selected dataset
+    for key in selected_keys:
+        system = rag_systems[key]
+        meta = system['metadata']
+        dataset_name = f"{meta['dataset_version']} - {meta['state_filter']}"
+        
+        comparison_result += f"## {dataset_name}\n\n"
+        
+        try:
+            result = system['conversation_chain'].invoke({"question": question})
+            answer = result["answer"]
+            comparison_result += f"{answer}\n\n"
+        except Exception as e:
+            comparison_result += f"Error: {str(e)}\n\n"
+        
+        comparison_result += "---\n\n"
+    
+    return comparison_result
+
+def analyze_provider_types_gradio(dataset_key=None):
+    """Analyze provider types in a dataset - Gradio version."""
+    global rag_systems, current_dataset_key
+    
+    # Determine which dataset to use
+    target_key = dataset_key if dataset_key and dataset_key in rag_systems else current_dataset_key
+    
+    if not target_key or target_key not in rag_systems:
+        return "❌ No dataset selected."
+    
+    system = rag_systems[target_key]
+    meta = system['metadata']
+    
+    analysis_question = """
+    Analyze the provider types in this dataset:
+    1. What are the most common provider types?
+    2. How many unique provider types are there?
+    3. What percentage of providers fall into each major category?
+    Please provide a detailed breakdown.
+    """
+    
+    try:
+        result = system['conversation_chain'].invoke({"question": analysis_question})
+        
+        analysis = f"# Provider Type Analysis\n"
+        analysis += f"**Dataset:** {meta['dataset_version']} - {meta['state_filter']}\n\n"
+        analysis += result["answer"]
+        
+        return analysis
+    except Exception as e:
+        return f"❌ Error analyzing provider types: {str(e)}"
+
+def clear_chat_history():
+    """Clear the chat history."""
+    return []
+
+def visualize_datasets_gradio(dataset_indices, dimensions, sample_size=1000):
+    """Create a visualization of one or more datasets - Gradio version."""
+    global rag_systems
+    
+    if not rag_systems:
+        return None, "❌ No datasets loaded. Please load datasets first."
+    
+    if not dataset_indices:
+        return None, "❌ Please select at least one dataset to visualize."
+    
+    # Parse indices and get dataset keys
+    selected_keys = []
+    for selection in dataset_indices:
+        try:
+            index = int(selection.split(".")[0])
+            if 1 <= index <= len(rag_systems):
+                key = list(rag_systems.keys())[index - 1]
+                selected_keys.append(key)
+        except:
+            continue
+    
+    if not selected_keys:
+        return None, "❌ Could not parse selected datasets."
+    
+    try:
+        # Create a combined visualization
+        all_vectors = []
+        all_metadata = []
+        all_contents = []
+        all_dataset_labels = []
+        
+        status_msg = f"Processing {len(selected_keys)} dataset(s)...\n"
+        
+        # Collect vectors from all requested datasets
+        for key in selected_keys:
+            vector_store = rag_systems[key]['vector_store']
+            meta = rag_systems[key]['metadata']
+            dataset_label = f"{meta['dataset_version']} - {meta['state_filter']}"
+            
+            # Limit vectors for performance
+            num_vectors = min(sample_size, vector_store.index.ntotal)
+            status_msg += f"- {dataset_label}: {num_vectors} vectors\n"
+            
+            for i in range(num_vectors):
+                all_vectors.append(vector_store.index.reconstruct(i))
+                
+                doc_id = vector_store.index_to_docstore_id[i]
+                document = vector_store.docstore.search(doc_id)
+                
+                all_metadata.append(document.metadata)
+                all_contents.append(document.page_content)
+                all_dataset_labels.append(dataset_label)
+        
+        if not all_vectors:
+            return None, "❌ No vectors to visualize."
+        
+        vectors = np.array(all_vectors)
+        status_msg += f"\nTotal vectors: {len(all_vectors)}\n"
+        
+        # Reduce dimensionality
+        status_msg += f"Reducing dimensionality to {dimensions}D using t-SNE..."
+        tsne = TSNE(n_components=dimensions, random_state=42, perplexity=min(30, len(all_vectors)-1))
+        reduced_vectors = tsne.fit_transform(vectors)
+        
+        # Create color mapping based on dataset
+        unique_labels = list(set(all_dataset_labels))
+        colors = []
+        color_palette = [
+            '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd',
+            '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf'
+        ]
+        color_map = {label: color_palette[i % len(color_palette)] 
+                    for i, label in enumerate(unique_labels)}
+        
+        colors = [color_map[label] for label in all_dataset_labels]
+        
+        # Create hover text
+        hover_texts = []
+        for meta, content, label in zip(all_metadata, all_contents, all_dataset_labels):
+            text = f"<b>Dataset:</b> {label}<br>"
+            
+            # Add key metadata fields
+            key_fields = ['STATE_CD', 'PROVIDER_TYPE_DESC', 'FIRST_NAME', 'LAST_NAME', 'ORG_NAME']
+            for field in key_fields:
+                if field in meta and meta[field]:
+                    text += f"<b>{field}:</b> {meta[field]}<br>"
+            
+            # Add a preview of the content
+            content_preview = content[:200] + "..." if len(content) > 200 else content
+            text += f"<br><b>Preview:</b> {content_preview}"
+            
+            hover_texts.append(text)
+        
+        # Create visualization
+        if dimensions == 2:
+            fig = go.Figure()
+            
+            # Add a trace for each dataset
+            for label in unique_labels:
+                # Get indices for this dataset
+                indices = [i for i, l in enumerate(all_dataset_labels) if l == label]
+                
+                # Add the scatter trace
+                fig.add_trace(go.Scatter(
+                    x=reduced_vectors[indices, 0],
+                    y=reduced_vectors[indices, 1],
+                    mode='markers',
+                    marker=dict(
+                        size=6,
+                        color=color_map[label],
+                        opacity=0.7,
+                        line=dict(width=1, color='white')
+                    ),
+                    text=[hover_texts[i] for i in indices],
+                    hoverinfo='text',
+                    hoverlabel=dict(bgcolor="white", font_size=12),
+                    name=label
+                ))
+            
+            fig.update_layout(
+                title={
+                    'text': 'Medicare Provider Data - 2D Vector Space Visualization',
+                    'font': {'size': 20}
+                },
+                xaxis_title='Dimension 1',
+                yaxis_title='Dimension 2',
+                width=900,
+                height=700,
+                hovermode='closest',
+                template='plotly_white',
+                legend=dict(
+                    yanchor="top",
+                    y=0.99,
+                    xanchor="left",
+                    x=0.01,
+                    bgcolor="rgba(255,255,255,0.8)"
+                )
+            )
+        else:  # 3D
+            fig = go.Figure()
+            
+            # Add a trace for each dataset
+            for label in unique_labels:
+                # Get indices for this dataset
+                indices = [i for i, l in enumerate(all_dataset_labels) if l == label]
+                
+                # Add the scatter trace
+                fig.add_trace(go.Scatter3d(
+                    x=reduced_vectors[indices, 0],
+                    y=reduced_vectors[indices, 1],
+                    z=reduced_vectors[indices, 2],
+                    mode='markers',
+                    marker=dict(
+                        size=5,
+                        color=color_map[label],
+                        opacity=0.7,
+                        line=dict(width=1, color='white')
+                    ),
+                    text=[hover_texts[i] for i in indices],
+                    hoverinfo='text',
+                    hoverlabel=dict(bgcolor="white", font_size=12),
+                    name=label
+                ))
+            
+            fig.update_layout(
+                title={
+                    'text': 'Medicare Provider Data - 3D Vector Space Visualization',
+                    'font': {'size': 20}
+                },
+                scene=dict(
+                    xaxis_title='Dimension 1',
+                    yaxis_title='Dimension 2',
+                    zaxis_title='Dimension 3',
+                    camera=dict(
+                        eye=dict(x=1.5, y=1.5, z=1.5)
+                    )
+                ),
+                width=900,
+                height=700,
+                template='plotly_white',
+                legend=dict(
+                    yanchor="top",
+                    y=0.99,
+                    xanchor="left",
+                    x=0.01,
+                    bgcolor="rgba(255,255,255,0.8)"
+                )
+            )
+        
+        success_msg = f"✅ Successfully created {dimensions}D visualization with {len(all_vectors)} vectors from {len(selected_keys)} dataset(s)"
+        return fig, success_msg
+        
+    except Exception as e:
+        return None, f"❌ Error creating visualization: {str(e)}"
+
+def create_dataset_statistics_plot(dataset_indices):
+    """Create statistical plots for selected datasets."""
+    global rag_systems
+    
+    if not rag_systems:
+        return None, "❌ No datasets loaded."
+    
+    if not dataset_indices:
+        return None, "❌ Please select at least one dataset."
+    
+    # Parse indices and get dataset keys
+    selected_keys = []
+    for selection in dataset_indices:
+        try:
+            index = int(selection.split(".")[0])
+            if 1 <= index <= len(rag_systems):
+                key = list(rag_systems.keys())[index - 1]
+                selected_keys.append(key)
+        except:
+            continue
+    
+    if not selected_keys:
+        return None, "❌ Could not parse selected datasets."
+    
+    try:
+        # Collect statistics
+        dataset_names = []
+        record_counts = []
+        chunk_counts = []
+        
+        for key in selected_keys:
+            meta = rag_systems[key]['metadata']
+            dataset_names.append(f"{meta['dataset_version']}<br>{meta['state_filter']}")
+            record_counts.append(meta['record_count'])
+            chunk_counts.append(meta['chunk_count'])
+        
+        # Create subplots
+        from plotly.subplots import make_subplots
+        
+        fig = make_subplots(
+            rows=1, cols=2,
+            subplot_titles=('Records per Dataset', 'Chunks per Dataset'),
+            specs=[[{'type': 'bar'}, {'type': 'bar'}]]
+        )
+        
+        # Add record count bars
+        fig.add_trace(
+            go.Bar(
+                x=dataset_names,
+                y=record_counts,
+                name='Records',
+                marker_color='lightblue',
+                text=record_counts,
+                textposition='auto',
+            ),
+            row=1, col=1
+        )
+        
+        # Add chunk count bars
+        fig.add_trace(
+            go.Bar(
+                x=dataset_names,
+                y=chunk_counts,
+                name='Chunks',
+                marker_color='lightgreen',
+                text=chunk_counts,
+                textposition='auto',
+            ),
+            row=1, col=2
+        )
+        
+        fig.update_layout(
+            title={
+                'text': 'Dataset Statistics Overview',
+                'font': {'size': 20}
+            },
+            showlegend=False,
+            height=500,
+            template='plotly_white'
+        )
+        
+        fig.update_xaxes(tickangle=-45)
+        
+        return fig, f"✅ Created statistics plot for {len(selected_keys)} dataset(s)"
+        
+    except Exception as e:
+        return None, f"❌ Error creating statistics plot: {str(e)}"
+
+def inspect_dataset_gradio(num_samples):
+    """Display sample documents from the current dataset - Gradio version."""
+    global rag_systems, current_dataset_key
+    
+    if not current_dataset_key or current_dataset_key not in rag_systems:
+        return "❌ No dataset selected. Please load a dataset first."
+    
+    # Get the dataset
+    system = rag_systems[current_dataset_key]
+    vector_store = system['vector_store']
+    meta = system['metadata']
+    
+    inspection_result = f"# Dataset Inspection\n\n"
+    inspection_result += f"**Dataset:** {meta['dataset_version']} - {meta['state_filter']}\n"
+    inspection_result += f"**Total documents:** {vector_store.index.ntotal}\n"
+    inspection_result += f"**Showing:** {min(num_samples, vector_store.index.ntotal)} sample documents\n\n"
+    inspection_result += "---\n\n"
+    
+    for i in range(min(num_samples, vector_store.index.ntotal)):
+        try:
+            doc_id = vector_store.index_to_docstore_id[i]
+            document = vector_store.docstore.search(doc_id)
+            
+            inspection_result += f"### Document {i+1}\n\n"
+            inspection_result += "**Metadata:**\n"
+            
+            # Show key metadata fields
+            key_fields = ['PROVIDER_TYPE_DESC', 'STATE_CD', 'FIRST_NAME', 'LAST_NAME', 
+                         'ORG_NAME', 'NPI', 'ENRLMT_ID']
+            
+            for field in key_fields:
+                if field in document.metadata and document.metadata[field]:
+                    inspection_result += f"- **{field}:** {document.metadata[field]}\n"
+            
+            # Show content preview
+            content_preview = document.page_content[:500] + "..." if len(document.page_content) > 500 else document.page_content
+            inspection_result += f"\n**Content Preview:**\n```\n{content_preview}\n```\n\n"
+            inspection_result += "---\n\n"
+            
+        except Exception as e:
+            inspection_result += f"Error retrieving document {i}: {str(e)}\n\n"
+    
+    return inspection_result
+
+def create_gradio_interface():
+    """Create the main Gradio interface."""
+    
+    with gr.Blocks(theme=theme, title="Medicare Provider Data Analysis System") as app:
+        # Header
+        gr.Markdown(
+            """
+            # 🏥 Medicare Provider Data Analysis System
+            
+            This system allows you to load, query, and analyze Medicare provider data using advanced RAG (Retrieval-Augmented Generation) technology.
+            
+            ---
+            """
+        )
+        
+        # Main tabs
+        with gr.Tabs() as tabs:
+            # Tab 1: Dataset Management
+            with gr.Tab("📊 Dataset Management"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.Markdown("### Load New Dataset")
+                        
+                        version_dropdown = gr.Dropdown(
+                            choices=list(DATASET_VERSIONS.keys()),
+                            label="Select Quarter/Year",
+                            value="Q1 2025"
+                        )
+                        
+                        state_dropdown = gr.Dropdown(
+                            choices=format_state_options(),
+                            label="Select State",
+                            value=""
+                        )
+                        
+                        max_records_slider = gr.Slider(
+                            minimum=100,
+                            maximum=5000,
+                            value=1000,
+                            step=100,
+                            label="Maximum Records"
+                        )
+                        
+                        use_sample_checkbox = gr.Checkbox(
+                            label="Load sample only (100 records)",
+                            value=True
+                        )
+                        
+                        load_button = gr.Button("🔄 Load Dataset", variant="primary")
+                        load_output = gr.Textbox(label="Loading Status", lines=3)
+                    
+                    with gr.Column(scale=1):
+                        gr.Markdown("### Manage Loaded Datasets")
+                        
+                        dataset_summary = gr.Markdown(get_dataset_summary(rag_systems))
+                        
+                        with gr.Row():
+                            dataset_selector = gr.Dropdown(
+                                choices=get_dataset_choices(),
+                                label="Select Dataset",
+                                interactive=True
+                            )
+                            
+                        with gr.Row():
+                            switch_button = gr.Button("↔️ Switch Dataset")
+                            remove_button = gr.Button("🗑️ Remove Dataset")
+                            clear_all_button = gr.Button("🧹 Clear All", variant="stop")
+                        
+                        manage_output = gr.Textbox(label="Status", lines=2)
+                
+                # Wire up dataset management events
+                def update_dataset_selector():
+                    return gr.update(choices=get_dataset_choices())
+                
+                load_button.click(
+                    fn=load_dataset_gradio,
+                    inputs=[version_dropdown, state_dropdown, max_records_slider, use_sample_checkbox],
+                    outputs=[load_output, dataset_summary]
+                ).then(
+                    fn=update_dataset_selector,
+                    outputs=dataset_selector
+                )
+                
+                switch_button.click(
+                    fn=switch_dataset_gradio,
+                    inputs=dataset_selector,
+                    outputs=[manage_output, dataset_summary]
+                )
+                
+                remove_button.click(
+                    fn=remove_dataset_gradio,
+                    inputs=dataset_selector,
+                    outputs=[manage_output, dataset_summary]
+                ).then(
+                    fn=update_dataset_selector,
+                    outputs=dataset_selector
+                )
+                
+                clear_all_button.click(
+                    fn=clear_all_datasets_gradio,
+                    outputs=[manage_output, dataset_summary]
+                ).then(
+                    fn=update_dataset_selector,
+                    outputs=dataset_selector
+                )
+            
+            # Tab 2: Query Interface
+            with gr.Tab("💬 Query & Chat"):
+                gr.Markdown("### Ask Questions About Your Data")
+                
+                current_dataset_info = gr.Markdown(get_current_dataset_info())
+                
+                # Create a timer to update current dataset info
+                timer = gr.Timer(value=2)
+                timer.tick(fn=get_current_dataset_info, outputs=current_dataset_info)
+                
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        chatbot = gr.Chatbot(
+                            label="Conversation",
+                            height=500,
+                            show_copy_button=True
+                        )
+                        
+                        with gr.Row():
+                            question_input = gr.Textbox(
+                                label="Your Question",
+                                placeholder="Ask about provider types, locations, statistics, etc.",
+                                lines=2,
+                                scale=4
+                            )
+                            
+                            with gr.Column(scale=1):
+                                ask_button = gr.Button("📤 Ask Current Dataset", variant="primary")
+                                global_ask_button = gr.Button("🌐 Ask All Datasets")
+                                clear_chat_button = gr.Button("🗑️ Clear Chat")
+                    
+                    with gr.Column(scale=1):
+                        gr.Markdown("### Quick Actions")
+                        
+                        analyze_providers_button = gr.Button("📊 Analyze Provider Types")
+                        
+                        gr.Markdown("### Example Questions")
+                        example_questions = [
+                            "What are the most common provider types?",
+                            "How many providers are in this dataset?",
+                            "Show me all psychiatrists in the data",
+                            "What types of medical facilities are included?",
+                            "Compare provider counts across different quarters"
+                        ]
+                        
+                        for eq in example_questions:
+                            gr.Button(eq, size="sm").click(
+                                lambda q=eq: (q, gr.update()),
+                                outputs=[question_input, chatbot]
+                            )
+                
+                # Wire up query events
+                question_input.submit(
+                    fn=ask_question_gradio,
+                    inputs=[question_input, chatbot],
+                    outputs=[question_input, chatbot]
+                )
+                
+                ask_button.click(
+                    fn=ask_question_gradio,
+                    inputs=[question_input, chatbot],
+                    outputs=[question_input, chatbot]
+                )
+                
+                global_ask_button.click(
+                    fn=ask_global_question_gradio,
+                    inputs=[question_input, chatbot],
+                    outputs=[question_input, chatbot]
+                )
+                
+                clear_chat_button.click(
+                    fn=clear_chat_history,
+                    outputs=chatbot
+                )
+                
+                analyze_providers_button.click(
+                    fn=lambda: ("", [(
+                        "Analyze provider types in the current dataset",
+                        analyze_provider_types_gradio()
+                    )]),
+                    outputs=[question_input, chatbot]
+                )
+            
+            # Tab 3: Comparison & Analysis
+            with gr.Tab("🔍 Compare Datasets"):
+                gr.Markdown("### Compare Multiple Datasets")
+                
+                with gr.Row():
+                    compare_dataset_selector = gr.CheckboxGroup(
+                        choices=get_dataset_choices(),
+                        label="Select Datasets to Compare (choose 2 or more)",
+                        value=[]
+                    )
+                
+                compare_question = gr.Textbox(
+                    label="Comparison Question",
+                    placeholder="Enter a question to ask all selected datasets",
+                    lines=2
+                )
+                
+                compare_button = gr.Button("🔄 Compare Datasets", variant="primary")
+                
+                comparison_output = gr.Markdown(label="Comparison Results")
+                
+                # Update checkbox choices when datasets change
+                def update_compare_selector():
+                    return gr.update(choices=get_dataset_choices())
+                
+                timer.tick(fn=update_compare_selector, outputs=compare_dataset_selector)
+                
+                compare_button.click(
+                    fn=compare_datasets_gradio,
+                    inputs=[compare_question, compare_dataset_selector],
+                    outputs=comparison_output
+                )
+            
+            # Tab 4: Visualization
+            with gr.Tab("📈 Visualizations"):
+                gr.Markdown("### Dataset Visualizations")
+                
+                with gr.Row():
+                    with gr.Column():
+                        viz_dataset_selector = gr.CheckboxGroup(
+                            choices=get_dataset_choices(),
+                            label="Select Datasets to Visualize",
+                            value=[]
+                        )
+                        
+                        viz_dimension = gr.Radio(
+                            choices=[2, 3],
+                            value=2,
+                            label="Visualization Dimensions"
+                        )
+                        
+                        viz_sample_size = gr.Slider(
+                            minimum=100,
+                            maximum=2000,
+                            value=500,
+                            step=100,
+                            label="Sample Size (per dataset)"
+                        )
+                        
+                        create_viz_button = gr.Button("🎨 Create Visualization", variant="primary")
+                        stats_button = gr.Button("📊 Show Statistics")
+                        
+                        viz_status = gr.Textbox(label="Status", lines=2)
+                
+                with gr.Row():
+                    viz_plot = gr.Plot(label="Vector Space Visualization")
+                    stats_plot = gr.Plot(label="Dataset Statistics")
+                
+                # Update visualization selector
+                def update_viz_selector():
+                    return gr.update(choices=get_dataset_choices())
+                
+                timer.tick(fn=update_viz_selector, outputs=viz_dataset_selector)
+                
+                create_viz_button.click(
+                    fn=visualize_datasets_gradio,
+                    inputs=[viz_dataset_selector, viz_dimension, viz_sample_size],
+                    outputs=[viz_plot, viz_status]
+                )
+                
+                stats_button.click(
+                    fn=create_dataset_statistics_plot,
+                    inputs=[viz_dataset_selector],
+                    outputs=[stats_plot, viz_status]
+                )
+            
+            # Tab 5: Dataset Inspector
+            with gr.Tab("🔎 Dataset Inspector"):
+                gr.Markdown("### Inspect Dataset Contents")
+                
+                inspect_current_info = gr.Markdown(get_current_dataset_info())
+                timer.tick(fn=get_current_dataset_info, outputs=inspect_current_info)
+                
+                num_samples_slider = gr.Slider(
+                    minimum=1,
+                    maximum=20,
+                    value=5,
+                    step=1,
+                    label="Number of Sample Documents"
+                )
+                
+                inspect_button = gr.Button("🔍 Inspect Current Dataset", variant="primary")
+                
+                inspection_output = gr.Markdown(label="Dataset Inspection Results")
+                
+                inspect_button.click(
+                    fn=inspect_dataset_gradio,
+                    inputs=num_samples_slider,
+                    outputs=inspection_output
+                )
+            
+            # Tab 6: Settings & Help
+            with gr.Tab("⚙️ Settings & Help"):
+                gr.Markdown(
+                    """
+                    ### System Information
+                    
+                    **Model:** GPT-4 Mini  
+                    **Embedding Model:** OpenAI Embeddings  
+                    **Vector Store:** FAISS  
+                    
+                    ### API Configuration
+                    
+                    This system uses the CMS.gov Data API to fetch Medicare provider information.
+                    
+                    ### Tips for Best Results
+                    
+                    1. **Loading Data**: Start with sample data (100 records) to test queries quickly
+                    2. **State Selection**: Load specific states for focused analysis
+                    3. **Querying**: Be specific in your questions for better results
+                    4. **Comparisons**: Load multiple quarters/states to analyze trends
+                    
+                    ### Common Use Cases
+                    
+                    - **Provider Analysis**: Find specific types of healthcare providers
+                    - **Geographic Distribution**: Analyze providers by state
+                    - **Temporal Trends**: Compare data across different quarters
+                    - **Provider Types**: Understand the distribution of specialties
+                    
+                    ### Troubleshooting
+                    
+                    - **No API Key**: Ensure OPENAI_API_KEY is set in your environment
+                    - **Loading Errors**: Check your internet connection and API limits
+                    - **Query Errors**: Try rephrasing your question or check if data is loaded
+                    """
+                )
+                
+                with gr.Row():
+                    gr.Markdown("### Current Configuration")
+                    config_info = gr.JSON(
+                        value={
+                            "api_key_set": bool(os.getenv('OPENAI_API_KEY')),
+                            "default_model": DEFAULT_MODEL,
+                            "api_base_url": API_BASE_URL,
+                            "datasets_loaded": len(rag_systems)
+                        },
+                        label="System Configuration"
+                    )
+        
+        # Footer
+        gr.Markdown(
+            """
+            ---
+            
+            <center>
+            Medicare Provider Data Analysis System | Powered by LangChain & OpenAI
+            </center>
+            """
+        )
+    
+    return app
+
+# Main execution
+if __name__ == "__main__":
+    # Create and launch the app
+    app = create_gradio_interface()
+    
+    # Launch with appropriate settings
+    app.launch(
+        server_name="0.0.0.0",  # Allow external connections
+        server_port=7860,  # Default Gradio port
+        
+    )