Update app.py

app.py

@@ -1,421 +1,526 @@
 import os
 import logging
-import asyncio
-import time
-from typing import Dict, List, Optional, Any, Tuple
-
-import gradio as gr
 import datasets
 import shodan
-import requests
+import asyncio
+import aiohttp
+import json
+import gradio as gr
+from typing import List, Dict, Any, Optional
 
-# Set up logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+def validate_env_variables():
+    """Validate that required environment variables are set."""
+    required_vars = ["SHODAN_API_KEY", "HF_TOKEN"]
+    missing_vars = [var for var in required_vars if not os.getenv(var)]
+    if missing_vars:
+        raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}")
 
 def load_or_create_dataset():
-    """
-    Load or create the dataset.
-    
-    Returns:
-        HuggingFace dataset
-    """
+    """Load the dataset from HuggingFace or create it if it doesn't exist."""
+    validate_env_variables()
     hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        raise ValueError("HF_TOKEN environment variable is not set")
     
     try:
-        dataset = datasets.load_dataset("latterworks/llama_checker_results", use_auth_token=hf_token)
-        # Convert to in-memory dataset for easier manipulation
-        dataset = dataset['train']
+        logger.info("Attempting to load dataset from HuggingFace Hub")
+        dataset = datasets.load_dataset(
+            "latterworks/llama_checker_results", 
+            use_auth_token=hf_token
+        )
+        if "train" in dataset:
+            return dataset["train"]
+        else:
+            # If there's no train split, just take the first available split
+            return dataset[next(iter(dataset))]
+            
     except FileNotFoundError:
-        # Dataset doesn't exist, create it
-        dataset = datasets.Dataset.from_dict({"ip": [], "port": [], "country": [], "region": [], "org": [], "models": []})
-        dataset.push_to_hub("latterworks/llama_checker_results", token=hf_token)
-        dataset = datasets.load_dataset("latterworks/llama_checker_results", use_auth_token=hf_token)['train']
+        logger.info("Dataset not found, creating a new one")
+        # Create an empty dataset with the required schema
+        empty_dataset = datasets.Dataset.from_dict({
+            "ip": [], 
+            "port": [], 
+            "country": [], 
+            "region": [], 
+            "org": [], 
+            "models": []
+        })
+        
+        # Push the empty dataset to HuggingFace Hub
+        empty_dataset.push_to_hub(
+            "latterworks/llama_checker_results", 
+            token=hf_token
+        )
+        
+        # Load the newly created dataset
+        dataset = datasets.load_dataset(
+            "latterworks/llama_checker_results", 
+            use_auth_token=hf_token
+        )
+        
+        if "train" in dataset:
+            return dataset["train"]
+        else:
+            return dataset[next(iter(dataset))]
+            
     except Exception as e:
-        logging.error(f"Failed to load or create dataset: {e}")
-        raise  # Re-raise the exception to stop the application
-    
-    return dataset
-
+        logger.error(f"Failed to load or create dataset: {e}")
+        raise
 
 def scan_shodan(progress=gr.Progress()) -> List[Dict]:
     """
-    Scan Shodan for Ollama instances.
+    Scan Shodan for Ollama instances using search_cursor for comprehensive result retrieval.
     
     Args:
-        progress: Gradio progress bar
-    
+        progress: Gradio progress bar for visual feedback
+        
     Returns:
-        List of dictionaries containing information about Ollama instances
+        List of Ollama instances from Shodan with comprehensive metadata
     """
-    # Validate Shodan API key exists
-    shodan_api_key = os.getenv("SHODAN_API_KEY")
-    if not shodan_api_key:
-        raise ValueError("SHODAN_API_KEY environment variable is not set")
-    
-    # Get Shodan query
+    # API key fetch - no validation needed as it's centralized at startup
+    api_key = os.getenv("SHODAN_API_KEY")
     shodan_query = os.getenv("SHODAN_QUERY", "product:Ollama port:11434")
     
-    # Initialize Shodan API
-    api = shodan.Shodan(shodan_api_key)
+    api = shodan.Shodan(api_key)
     
     try:
-        # Search Shodan
-        logger.info(f"Searching Shodan with query: {shodan_query}")
-        results = api.search(shodan_query)
+        logger.info(f"Executing Shodan search_cursor with query: {shodan_query}")
+        
+        # Use search_cursor to handle pagination automatically
+        cursor = api.search_cursor(shodan_query)
         
-        # Process results
+        # Initialize scan metrics
         instances = []
-        total_results = results['total']
-        logger.info(f"Found {total_results} results")
+        processed = 0
+        batch_size = 100  # Process results in batches for progress updates
         
-        # Set up progress bar
-        progress(0, desc="Scanning Shodan for Ollama instances")
+        progress(0, desc="Initializing Shodan data retrieval")
         
-        for i, result in enumerate(results['matches']):
-            progress((i+1)/len(results['matches']), desc=f"Processing result {i+1}/{len(results['matches'])}")
-            
-            instance = {
-                'ip': result['ip_str'],
-                'port': result.get('port', 11434),
-                'country': result.get('location', {}).get('country_name'),
-                'region': result.get('location', {}).get('region_name'),
-                'org': result.get('org'),
-                'models': []
-            }
-            instances.append(instance)
+        # Process all results from the cursor
+        results_batch = []
+        for result in cursor:
+            results_batch.append(result)
+            processed += 1
             
+            # Process in batches for efficiency
+            if len(results_batch) >= batch_size:
+                progress(min(1.0, processed / (processed + 100)), desc=f"Retrieved {processed} Ollama instances")
+                
+                # Extract instance data from batch
+                for result in results_batch:
+                    instances.append({
+                        'ip': result.get('ip_str'),
+                        'port': result.get('port', 11434),
+                        'country': result.get('location', {}).get('country_name'),
+                        'region': result.get('location', {}).get('region_name'),
+                        'org': result.get('org'),
+                        'models': []
+                    })
+                results_batch = []
+        
+        # Process any remaining results
+        if results_batch:
+            for result in results_batch:
+                instances.append({
+                    'ip': result.get('ip_str'),
+                    'port': result.get('port', 11434),
+                    'country': result.get('location', {}).get('country_name'),
+                    'region': result.get('location', {}).get('region_name'),
+                    'org': result.get('org'),
+                    'models': []
+                })
+        
+        logger.info(f"Completed Shodan scan, retrieved {len(instances)} Ollama instances")
         return instances
-    
+        
     except shodan.APIError as e:
-        logger.error(f"Shodan API error: {e}")
-        raise
+        error_msg = str(e)
+        if "Invalid API key" in error_msg:
+            logger.error(f"Shodan authentication failed: Invalid API key")
+            raise ValueError("Invalid Shodan API key. Please check your SHODAN_API_KEY environment variable.")
+        elif "Request rate limit reached" in error_msg:
+            logger.error(f"Shodan rate limit exceeded: {e}")
+            raise ValueError("Shodan API rate limit exceeded. Please wait before trying again.")
+        else:
+            logger.error(f"Shodan API error: {e}")
+            raise
     except Exception as e:
-        logger.error(f"Error during Shodan scan: {e}")
+        logger.error(f"Unhandled exception during Shodan scan: {e}")
         raise
 
+async def check_single_endpoint(session, instance):
+    """Check a single Ollama endpoint for available models."""
+    ip = instance['ip']
+    port = instance['port']
+    url = f"http://{ip}:{port}/api/tags"
+    
+    try:
+        logger.info(f"Checking Ollama endpoint: {url}")
+        
+        # Set a timeout for the request
+        async with session.get(url, timeout=5) as response:
+            if response.status == 200:
+                data = await response.json()
+                models = data.get('models', [])
+                logger.info(f"Found {len(models)} models at {url}")
+                instance['models'] = models
+                return instance
+            else:
+                logger.warning(f"Failed to get models from {url} - Status: {response.status}")
+                return instance
+    except asyncio.TimeoutError:
+        logger.warning(f"Timeout connecting to {url}")
+        return instance
+    except Exception as e:
+        logger.error(f"Error checking {url}: {e}")
+        return instance
 
-async def check_single_endpoint(ip: str, port: int) -> Optional[List[Dict]]:
+async def check_ollama_endpoints(instances, progress=gr.Progress()):
     """
-    Check a single Ollama endpoint for available models.
+    Check multiple Ollama endpoints for available models.
     
     Args:
-        ip: IP address of the endpoint
-        port: Port number of the endpoint
-    
+        instances: List of Ollama instances from Shodan
+        progress: Gradio progress bar
+        
     Returns:
-        List of models if successful, None otherwise
+        List of Ollama instances with model information
     """
-    url = f"http://{ip}:{port}/api/tags"
+    if not instances:
+        return []
+        
+    progress(0, desc="Checking Ollama endpoints")
     
-    try:
-        # Set a timeout of 5 seconds
-        response = requests.get(url, timeout=5)
-        if response.status_code == 200:
-            data = response.json()
-            return data.get('models', [])
-        else:
-            logger.warning(f"Failed to get models from {ip}:{port}, status code: {response.status_code}")
-            return None
-    except requests.exceptions.RequestException as e:
-        logger.warning(f"Error connecting to {ip}:{port}: {e}")
-        return None
-    except Exception as e:
-        logger.warning(f"Unexpected error checking {ip}:{port}: {e}")
-        return None
-
+    # Set up async HTTP session
+    async with aiohttp.ClientSession() as session:
+        tasks = []
+        for instance in instances:
+            task = check_single_endpoint(session, instance)
+            tasks.append(task)
+        
+        # Process tasks with progress updates
+        updated_instances = []
+        for i, task in enumerate(asyncio.as_completed(tasks)):
+            progress((i + 1) / len(tasks), desc=f"Checking endpoint {i + 1}/{len(tasks)}")
+            instance = await task
+            updated_instances.append(instance)
+    
+    return updated_instances
 
-async def check_ollama_endpoints(instances: List[Dict], dataset, progress=gr.Progress()) -> datasets.Dataset:
+def update_dataset_with_instances(dataset, instances):
     """
-    Check all Ollama endpoints for available models asynchronously.
+    Update the HuggingFace dataset with new Ollama instances.
     
     Args:
-        instances: List of dictionaries containing information about Ollama instances
         dataset: HuggingFace dataset
-        progress: Gradio progress bar
+        instances: List of Ollama instances with model information
         
     Returns:
-        Updated dataset
+        Updated HuggingFace dataset
     """
-    # Validate HF token exists
-    hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        raise ValueError("HF_TOKEN environment variable is not set")
-    
-    # Convert dataset to dictionary for easier manipulation
-    dataset_dict = {f"{item['ip']}:{item['port']}": item for item in dataset}
+    if not instances:
+        logger.warning("No instances to update in dataset")
+        return dataset
+        
+    # Convert dataset to list of dictionaries for easier manipulation
+    dataset_dict = {f"{item['ip']}:{item['port']}": item for item in dataset.to_list()}
     
     # Process each instance
-    progress(0, desc="Checking Ollama endpoints")
-    for i, instance in enumerate(instances):
-        progress((i+1)/len(instances), desc=f"Checking endpoint {i+1}/{len(instances)}")
-        
-        ip = instance['ip']
-        port = instance['port']
-        key = f"{ip}:{port}"
-        
-        # Get models from the endpoint
-        models = await check_single_endpoint(ip, port)
+    updates_count = 0
+    new_instances = []
+    
+    for instance in instances:
+        instance_key = f"{instance['ip']}:{instance['port']}"
         
-        if models:
-            # Update instance with models
-            instance['models'] = models
+        if instance_key in dataset_dict:
+            # Update existing instance
+            dataset_dict[instance_key]['country'] = instance.get('country', dataset_dict[instance_key].get('country'))
+            dataset_dict[instance_key]['region'] = instance.get('region', dataset_dict[instance_key].get('region'))
+            dataset_dict[instance_key]['org'] = instance.get('org', dataset_dict[instance_key].get('org'))
             
-            # Update or add to dataset dictionary
-            dataset_dict[key] = instance
+            # Only update models if they were found
+            if instance.get('models'):
+                dataset_dict[instance_key]['models'] = instance['models']
+            
+            updates_count += 1
+        else:
+            # Add new instance
+            new_instances.append(instance)
     
-    # Convert back to dataset
-    updated_dataset = datasets.Dataset.from_dict({
-        "ip": [item['ip'] for item in dataset_dict.values()],
-        "port": [item['port'] for item in dataset_dict.values()],
-        "country": [item.get('country', '') for item in dataset_dict.values()],
-        "region": [item.get('region', '') for item in dataset_dict.values()],
-        "org": [item.get('org', '') for item in dataset_dict.values()],
-        "models": [item.get('models', []) for item in dataset_dict.values()]
-    })
+    # Create updated dataset list
+    updated_dataset_list = list(dataset_dict.values()) + new_instances
     
-    # Push updates to hub
+    # Create a new dataset from the updated list
+    updated_dataset = datasets.Dataset.from_list(updated_dataset_list)
+    
+    # Push updated dataset to HuggingFace Hub
+    hf_token = os.getenv("HF_TOKEN")
     updated_dataset.push_to_hub("latterworks/llama_checker_results", token=hf_token)
     
+    logger.info(f"Updated {updates_count} existing instances and added {len(new_instances)} new instances to dataset")
+    
     return updated_dataset
 
-
-def get_unique_values(dataset) -> Tuple[List[str], List[str], List[str]]:
+def get_unique_values(dataset):
     """
-    Get unique values for family, parameter_size, and name.
+    Get unique values for model attributes to populate dropdown filters.
     
     Args:
         dataset: HuggingFace dataset
-    
+        
     Returns:
-        Tuple of lists containing unique values for family, parameter_size, and name
+        Dictionary with unique values for each attribute
     """
+    # Initialize empty sets
     families = set()
     parameter_sizes = set()
-    names = set()
-    
-    for item in dataset:
-        for model in item.get('models', []):
-            if 'family' in model and model['family']:
-                families.add(model['family'])
-            if 'parameter_size' in model and model['parameter_size']:
-                parameter_sizes.add(model['parameter_size'])
-            if 'name' in model and model['name']:
-                names.add(model['name'])
     
-    # Convert to sorted lists and add empty option
-    families = [''] + sorted(list(families))
-    parameter_sizes = [''] + sorted(list(parameter_sizes))
-    names = sorted(list(names))
+    # Extract unique values from models
+    for instance in dataset:
+        for model in instance.get('models', []):
+            details = model.get('details', {})
+            
+            # Handle both direct details in the model and nested details
+            if isinstance(details, dict):
+                family = details.get('family')
+                parameter_size = details.get('parameter_size')
+            else:
+                family = model.get('family')
+                parameter_size = model.get('parameter_size')
+                
+            if family:
+                families.add(family)
+            
+            if parameter_size:
+                parameter_sizes.add(parameter_size)
     
-    return families, parameter_sizes, names
-
+    return {
+        'families': sorted(list(families)),
+        'parameter_sizes': sorted(list(parameter_sizes))
+    }
 
-def search_models(
-    dataset,
-    family: str = "",
-    parameter_size: str = "",
-    name: str = "",
-    is_admin: bool = False
-) -> Tuple[List[Dict], List[Dict]]:
+def search_models(dataset, family=None, parameter_size=None, name_search=None, is_admin=False):
     """
-    Search models based on criteria.
+    Search for models in the dataset based on filters.
     
     Args:
         dataset: HuggingFace dataset
         family: Filter by model family
         parameter_size: Filter by parameter size
-        name: Filter by model name
-        is_admin: Whether the user is an admin
-    
+        name_search: Filter by model name (substring match)
+        is_admin: Whether to include IP and port information
+        
     Returns:
-        Tuple of (results, selected model info)
+        List of dictionaries with model information
     """
     results = []
     
-    for item in dataset:
-        for model in item.get('models', []):
+    for instance in dataset:
+        ip = instance.get('ip')
+        port = instance.get('port')
+        country = instance.get('country')
+        region = instance.get('region')
+        org = instance.get('org')
+        
+        for model in instance.get('models', []):
+            # Extract model details
+            model_name = model.get('name', '')
+            
+            # Handle both direct details in the model and nested details
+            details = model.get('details', {})
+            if isinstance(details, dict):
+                model_family = details.get('family', '')
+                model_param_size = details.get('parameter_size', '')
+                model_quant_level = details.get('quantization_level', '')
+            else:
+                model_family = model.get('family', '')
+                model_param_size = model.get('parameter_size', '')
+                model_quant_level = model.get('quantization_level', '')
+                
+            model_size_bytes = model.get('size', 0)
+            model_size_gb = model_size_bytes / (1024 * 1024 * 1024) if model_size_bytes else 0
+            
             # Apply filters
-            if family and model.get('family', '') != family:
+            if family and model_family != family:
                 continue
-            if parameter_size and model.get('parameter_size', '') != parameter_size:
+                
+            if parameter_size and model_param_size != parameter_size:
                 continue
-            if name and name.lower() not in model.get('name', '').lower():
+                
+            if name_search and name_search.lower() not in model_name.lower():
                 continue
             
-            # Create result with model info
+            # Create result object
             result = {
-                'name': model.get('name', ''),
-                'family': model.get('family', ''),
-                'parameter_size': model.get('parameter_size', ''),
-                'quantization_level': model.get('quantization_level', ''),
-                'size': round(model.get('size', 0) / (1024**3), 2)  # Convert to GB
+                'name': model_name,
+                'family': model_family,
+                'parameter_size': model_param_size,
+                'quantization_level': model_quant_level,
+                'size_gb': round(model_size_gb, 2),
+                'country': country,
+                'region': region,
+                'org': org,
             }
             
-            # Add IP and port only for admin users
+            # Include full model info for details view
+            result['full_model_info'] = json.dumps(model, indent=2)
+            
+            # Include IP and port for admin users only
             if is_admin:
-                result['ip'] = item['ip']
-                result['port'] = item['port']
+                result['ip'] = ip
+                result['port'] = port
             
             results.append(result)
     
-    # For empty result, return empty JSON info
-    selected_model_info = [{}]
-    
-    return results, selected_model_info
-
-
-def get_model_info(model_row: Dict) -> Dict:
-    """
-    Get detailed information about a selected model.
-    
-    Args:
-        model_row: Selected model row from the results
-    
-    Returns:
-        Dictionary containing detailed model information
-    """
-    return model_row
-
+    return results
 
 def create_interface():
-    """
-    Create Gradio interface for the application.
-    
-    Returns:
-        Gradio interface
-    """
-    # Load or create dataset
-    dataset = load_or_create_dataset()
-    
-    # Check for admin mode
-    is_admin = os.getenv("ADMIN_MODE", "false").lower() == "true"
-    
-    # Get unique values for dropdown menus
-    families, parameter_sizes, names = get_unique_values(dataset)
-    
-    # Get initial search results
-    initial_results, initial_model_info = search_models(dataset, is_admin=is_admin)
-    
-    # Function to run Shodan scan
-    def run_shodan_scan(progress=gr.Progress()):
-        nonlocal dataset
-        instances = scan_shodan(progress)
-        dataset = asyncio.run(check_ollama_endpoints(instances, dataset, progress))
+    """Create the Gradio interface for the application."""
+    try:
+        # Load dataset once at startup
+        dataset = load_or_create_dataset()
         
-        # Update unique values
-        updated_families, updated_parameter_sizes, updated_names = get_unique_values(dataset)
+        # Get unique values for dropdowns once at startup
+        unique_values = get_unique_values(dataset)
         
-        # Update search results
-        updated_results, updated_model_info = search_models(dataset, is_admin=is_admin)
+        # Get all models to display on initial load
+        initial_results = search_models(dataset)
         
-        return (
-            updated_families, updated_parameter_sizes,
-            updated_results, updated_model_info
-        )
-    
-    # Function to run model search
-    def run_search(family, parameter_size, name):
-        results, model_info = search_models(dataset, family, parameter_size, name, is_admin=is_admin)
-        return results, model_info
-    
-    # Function to get model details when a row is selected
-    def select_model(evt: gr.SelectData, results):
-        if evt.index[0] < len(results):
-            selected = results[evt.index[0]]
-            return selected
-        return {}
-    
-    # Create Gradio interface
-    with gr.Blocks(title="Ollama Instance Scanner") as interface:
-        gr.Markdown("# Ollama Instance Scanner")
-        
-        with gr.Tabs():
-            # Browse Models tab
-            with gr.TabItem("Browse Models"):
-                with gr.Row():
-                    with gr.Column():
-                        family_dropdown = gr.Dropdown(
-                            choices=families,
-                            label="Model Family",
-                            value=""
-                        )
-                        parameter_size_dropdown = gr.Dropdown(
-                            choices=parameter_sizes,
-                            label="Parameter Size",
-                            value=""
-                        )
-                        name_search = gr.Textbox(
-                            label="Model Name",
-                            placeholder="Search by name..."
+        # Create Gradio interface
+        with gr.Blocks(title="Ollama Instance Scanner") as interface:
+            gr.Markdown("# Ollama Instance Scanner")
+            gr.Markdown("Browse publicly accessible Ollama instances and their models")
+            
+            with gr.Tabs() as tabs:
+                # Browse Models Tab
+                with gr.TabItem("Browse Models"):
+                    with gr.Row():
+                        with gr.Column(scale=1):
+                            family_dropdown = gr.Dropdown(
+                                choices=["All"] + unique_values['families'],
+                                value="All",
+                                label="Filter by Family"
+                            )
+                            parameter_size_dropdown = gr.Dropdown(
+                                choices=["All"] + unique_values['parameter_sizes'],
+                                value="All",
+                                label="Filter by Parameter Size"
+                            )
+                            name_search = gr.Textbox(
+                                label="Search by Name", 
+                                placeholder="Enter model name..."
+                            )
+                            search_button = gr.Button("Search")
+                    
+                    with gr.Row():
+                        results_table = gr.DataFrame(
+                            value=initial_results,
+                            headers=["name", "family", "parameter_size", "quantization_level", "size_gb", "country", "region", "org"],
+                            label="Search Results"
                         )
-                        search_button = gr.Button("Search")
+                    
+                    with gr.Row():
+                        model_details = gr.JSON(label="Model Details")
                 
-                results_df = gr.DataFrame(
-                    value=initial_results,
-                    label="Search Results",
-                    headers=["name", "family", "parameter_size", "quantization_level", "size"],
-                    row_count=10,
-                    interactive=False
-                )
+                # Shodan Scan Tab (Admin only)
+                with gr.TabItem("Shodan Scan (Admin Only)"):
+                    gr.Markdown("## Shodan Scan")
+                    gr.Markdown("This tab allows scanning for Ollama instances using Shodan. You need a valid Shodan API key set as an environment variable.")
+                    
+                    shodan_scan_button = gr.Button("Start Shodan Scan")
+                    scan_status = gr.Textbox(label="Scan Status", interactive=False)
+            
+            # Define event handlers
+            def on_search_click(family, parameter_size, name_search):
+                # Use "All" as a signal not to filter
+                family_filter = None if family == "All" else family
+                param_size_filter = None if parameter_size == "All" else parameter_size
+                name_filter = None if not name_search else name_search
                 
-                model_info = gr.JSON(
-                    value=initial_model_info[0] if initial_model_info else {},
-                    label="Model Details"
-                )
+                # Check if admin mode is enabled (would need to implement proper authentication)
+                is_admin = False  # This should be based on proper authentication
                 
-                # Event handlers
-                search_button.click(
-                    fn=run_search,
-                    inputs=[family_dropdown, parameter_size_dropdown, name_search],
-                    outputs=[results_df, model_info]
-                )
+                # Search for models
+                results = search_models(dataset, family_filter, param_size_filter, name_filter, is_admin)
                 
-                results_df.select(
-                    fn=select_model,
-                    inputs=[results_df],
-                    outputs=[model_info]
-                )
+                # Return results
+                return results
             
-            # Shodan Scan tab
-            with gr.TabItem("Shodan Scan"):
-                # Check if Shodan API key is available
-                shodan_api_key = os.getenv("SHODAN_API_KEY")
-                
-                if shodan_api_key:
-                    scan_button = gr.Button("Start Scan")
-                    scan_output = gr.Markdown("Press the button to start scanning Shodan for Ollama instances.")
+            def on_table_select(evt: gr.SelectData, results):
+                if evt.index[0] < len(results):
+                    selected_row = results[evt.index[0]]
+                    return selected_row.get('full_model_info', {})
+                return {}
+            
+            async def run_shodan_scan():
+                try:
+                    # Verify Shodan API Key exists
+                    if not os.getenv("SHODAN_API_KEY"):
+                        return "Error: SHODAN_API_KEY environment variable is not set."
+                    
+                    # Perform Shodan scan
+                    instances = scan_shodan()
                     
-                    # Event handlers
-                    scan_button.click(
-                        fn=run_shodan_scan,
-                        outputs=[
-                            family_dropdown, parameter_size_dropdown,
-                            results_df, model_info
-                        ]
-                    )
-                else:
-                    gr.Markdown("## Shodan API key not configured")
-                    gr.Markdown(
-                        "To use the Shodan scan feature, you need to set the `SHODAN_API_KEY` "
-                        "environment variable in your Hugging Face Space settings."
-                    )
+                    if not instances:
+                        return "No Ollama instances found in Shodan scan."
+                    
+                    # Check Ollama endpoints
+                    updated_instances = await check_ollama_endpoints(instances)
+                    
+                    # Update dataset
+                    nonlocal dataset
+                    dataset = update_dataset_with_instances(dataset, updated_instances)
+                    
+                    # Update unique values
+                    nonlocal unique_values
+                    unique_values = get_unique_values(dataset)
+                    
+                    # Update dropdown choices
+                    family_dropdown.choices = ["All"] + unique_values['families']
+                    parameter_size_dropdown.choices = ["All"] + unique_values['parameter_sizes']
+                    
+                    return f"Scan completed successfully. Found {len(instances)} instances, {sum(1 for i in updated_instances if i.get('models'))} with accessible models."
+                except Exception as e:
+                    logger.error(f"Error in Shodan scan: {e}")
+                    return f"Error: {str(e)}"
+            
+            # Connect event handlers
+            search_button.click(
+                on_search_click,
+                inputs=[family_dropdown, parameter_size_dropdown, name_search],
+                outputs=[results_table]
+            )
+            
+            results_table.select(
+                on_table_select,
+                inputs=[results_table],
+                outputs=[model_details]
+            )
+            
+            shodan_scan_button.click(
+                run_shodan_scan,
+                inputs=[],
+                outputs=[scan_status]
+            )
+        
+        return interface
     
-    return interface
-
+    except Exception as e:
+        logger.error(f"Failed to create Gradio interface: {e}")
+        raise
 
 def main():
     """Main function to run the application."""
-    # Create and launch interface
-    interface = create_interface()
-    interface.launch()
-
+    try:
+        interface = create_interface()
+        if interface:
+            interface.launch()
+        else:
+            logger.error("Failed to create interface")
+    except Exception as e:
+        logger.error(f"Application failed: {e}")
 
 if __name__ == "__main__":
     main()