interactive-retrain-working.sh

#!/bin/bash

# Interactive Retrain Script for Ollama Embeddings - Working Version
set -e

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROGRAM_PATH="$SCRIPT_DIR/target/release/ollama-embeddings"

# Create log file with timestamp in the script directory
LOG_TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
NON_INFO_LOG="$SCRIPT_DIR/interactive-non-info-messages_${LOG_TIMESTAMP}.log"
TEMP_LOG="/tmp/ollama_embeddings_full_${LOG_TIMESTAMP}.log"

# Function to capture output and filter non-INFO messages
run_with_logging() {
    # Run command and tee output to temp file and stdout
    "$@" 2>&1 | tee "$TEMP_LOG"
    
    # Extract non-INFO messages from temp file
    if [ -f "$TEMP_LOG" ]; then
        grep -E '\[(WARN|ERROR|DEBUG|TRACE)\]' "$TEMP_LOG" >> "$NON_INFO_LOG" 2>/dev/null || true
        rm -f "$TEMP_LOG"
    fi
}

echo -e "${CYAN}Interactive Retrain Script for Ollama Embeddings${NC}"
echo -e "${YELLOW}Non-INFO messages will be logged to: $NON_INFO_LOG${NC}"
echo

if [ ! -f "$PROGRAM_PATH" ]; then
    echo -e "${RED}Error: ollama-embeddings binary not found${NC}"
    echo -e "${YELLOW}Please run 'cargo build --release' first.${NC}"
    exit 1
fi

# Check ollama server connectivity
check_ollama_server() {
    local url="${1:-http://localhost:11434}"
    if ! curl -s -f "$url/api/version" >/dev/null 2>&1; then
        echo -e "${YELLOW}Warning: Cannot connect to Ollama server at $url${NC}" >&2
        echo -e "${YELLOW}Make sure Ollama is running and accessible${NC}" >&2
        return 1
    fi
    return 0
}

get_model_dimensions() {
    local model="$1"
    case "$model" in
        "nomic-embed-text:latest"|"nomic-embed-text")
            echo "768"
            ;;
        "mxbai-embed-large:latest"|"mxbai-embed-large")
            echo "1024"
            ;;
        "all-minilm:latest"|"all-minilm")
            echo "384"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

show_embedding_models() {
    echo -e "${CYAN}Common embedding models:${NC}"
    echo "  • nomic-embed-text:latest (768 dim, general purpose)"
    echo "  • mxbai-embed-large:latest (1024 dim, large context)"
    echo "  • all-minilm:latest (384 dim, fast, smaller)"
}

show_retrieval_models() {
    local embed_model="$1"
    echo -e "${CYAN}Common retrieval models:${NC}"
    if [ -n "$embed_model" ] && [ "$embed_model" != "default" ]; then
        echo -e "${YELLOW}⚠️  Recommended: Use the same model as embedding ($embed_model)${NC}"
    fi
    echo "  • nomic-embed-text:latest (768 dimensions)"
    echo "  • mxbai-embed-large:latest (1024 dimensions)"
    echo "  • all-minilm:latest (384 dimensions)"
    echo -e "${RED}Note: Different models must have matching embedding dimensions${NC}"
}

show_prompt_examples() {
    echo -e "${CYAN}Enter a prompt for RAG (Retrieval Augmented Generation):${NC}"
    echo -e "${YELLOW}Examples:${NC}"
    echo "  • What is this code about?"
    echo "  • Summarize the main functionality"
    echo "  • Explain the key concepts"
}

prompt_for_rag() {
    show_prompt_examples
    echo
    read -p "Prompt [default: Analyze and summarize the content]: " rag_prompt
    echo "${rag_prompt:-Analyze and summarize the content}"
}

select_config_file() {
    local config_dir="$SCRIPT_DIR/config"
    
    if [ ! -d "$config_dir" ]; then
        echo "" # No config file
        return 0
    fi
    
    # Find all .toml files in config directory
    local toml_files=()
    local toml_basenames=()
    while IFS= read -r -d '' file; do
        toml_files+=("$file")
        toml_basenames+=("$(basename "$file")")
    done < <(find "$config_dir" -maxdepth 1 -name "*.toml" -type f -print0 2>/dev/null | sort -z)
    
    if [ ${#toml_files[@]} -eq 0 ]; then
        echo "" # No config files found
        return 0
    fi
    
    echo >&2
    echo -e "${BLUE}Configuration File Selection${NC}" >&2
    echo -e "${CYAN}Available configuration files:${NC}" >&2
    echo "  0. None (use default settings)" >&2
    
    # Use cached basenames for display
    local i=1
    for basename in "${toml_basenames[@]}"; do
        echo "  $i. $basename" >&2
        ((i++))
    done
    
    echo >&2
    while true; do
        read -p "Select config file [0-$((${#toml_files[@]}))]: " config_choice
        
        if [ -z "$config_choice" ] || [ "$config_choice" = "0" ]; then
            echo "" # No config
            return 0
        fi
        
        if [[ "$config_choice" =~ ^[0-9]+$ ]] && [ "$config_choice" -ge 1 ] && [ "$config_choice" -le ${#toml_files[@]} ]; then
            local selected_file="${toml_files[$((config_choice - 1))]}"
            echo -e "${GREEN}Selected: ${toml_basenames[$((config_choice - 1))]}${NC}" >&2
            echo "$selected_file"
            return 0
        fi
        
        echo -e "${RED}Invalid selection. Please enter a number between 0 and ${#toml_files[@]}${NC}" >&2
    done
}

extract_base_url() {
    local config_file="$1"
    
    if [ -z "$config_file" ] || [ ! -f "$config_file" ]; then
        return 0
    fi
    
    # Extract base_url or url from config file
    # Handles: base_url = "http://..." or url = "http://..."
    local url=$(grep -E '^\s*(base_url|url)\s*=\s*"[^"]*"' "$config_file" 2>/dev/null | head -n1 | sed -E 's/^\s*(base_url|url)\s*=\s*"([^"]*)".*/\2/')
    
    if [ -n "$url" ]; then
        echo "$url"
    fi
}

show_config_info() {
    local config_file="$1"
    local config_type="$2"  # "embed" or "retrieve" or "minimal"
    
    if [ -z "$config_file" ]; then
        return 0
    fi
    
    # Display config file being used
    if [ "$config_type" = "minimal" ]; then
        echo -e "${CYAN}Using minimal config for qwen model${NC}" >&2
    else
        echo -e "${CYAN}Using config: $(basename "$config_file")${NC}" >&2
    fi
    
    # Extract and display URL if present
    local base_url=$(extract_base_url "$config_file")
    if [ -n "$base_url" ]; then
        echo -e "${GREEN}Using URL from config: $base_url${NC}" >&2
    fi
}

display_command_and_confirm() {
    local config_file="$1"
    shift
    local cmd=("$@")
    
    echo >&2
    echo -e "${BLUE}Command to execute:${NC}" >&2
    echo -e "${CYAN}RUST_LOG=info ${cmd[*]}${NC}" >&2
    
    # Show config file contents if provided
    if [ -n "$config_file" ] && [ -f "$config_file" ]; then
        echo >&2
        echo -e "${YELLOW}Config file contents:${NC}" >&2
        if grep -E "^(base_url|url)\s*=" "$config_file" >/dev/null 2>&1; then
            echo -e "${GREEN}Found URL configuration:${NC}" >&2
            grep -E "^(base_url|url)\s*=" "$config_file" | sed 's/^/  /' >&2
        else
            echo -e "${YELLOW}  No base_url/url found (will use defaults)${NC}" >&2
        fi
    fi
    
    echo >&2
    read -p "Press Enter to start execution, or Ctrl+C to cancel..." >&2
    echo >&2
}

show_workflow_summary() {
    local embed_model="$1"
    local retrieve_model="$2"
    local rag_prompt="$3"
    local config_file="$4"
    
    echo
    echo -e "${YELLOW}Summary:${NC}"
    if [ -n "$embed_model" ]; then
        echo -e "  Embedding model: ${CYAN}$embed_model${NC}"
    fi
    echo -e "  Retrieval model: ${CYAN}$retrieve_model${NC}"
    if [ -n "$rag_prompt" ]; then
        echo -e "  RAG Prompt: ${CYAN}$rag_prompt${NC}"
    fi
    if [ -n "$config_file" ]; then
        echo -e "  Config file: ${CYAN}$(basename "$config_file")${NC}"
    else
        echo -e "  Config file: ${CYAN}None (default settings)${NC}"
    fi
    echo
}

validate_model_input() {
    local model="$1"
    local model_type="$2"  # "embedding" or "retrieval"
    
    if [ -z "$model" ]; then
        echo -e "${RED}${model_type^} model is required${NC}" >&2
        return 1
    fi
    
    # Check if model contains suspicious characters
    if [[ "$model" =~ [[:cntrl:]] ]]; then
        echo -e "${RED}Invalid characters in model name${NC}" >&2
        return 1
    fi
    
    return 0
}


warn_dimension_mismatch() {
    local embed_model="$1"
    local retrieve_model="$2"
    
    # Cache dimension lookups
    local embed_dims=$(get_model_dimensions "$embed_model")
    local retrieve_dims=$(get_model_dimensions "$retrieve_model")
    
    # Early return if dimensions match or are unknown
    if [ "$embed_dims" = "unknown" ] || [ "$retrieve_dims" = "unknown" ] || [ "$embed_dims" = "$retrieve_dims" ]; then
        return 0
    fi
    
    echo -e "${RED}⚠️  Warning: Dimension mismatch detected!${NC}" >&2
    echo -e "${RED}   Embedding model ($embed_model): $embed_dims dimensions${NC}" >&2
    echo -e "${RED}   Retrieval model ($retrieve_model): $retrieve_dims dimensions${NC}" >&2
    echo -e "${YELLOW}   This will cause a 'dimension mismatch' error during retrieval.${NC}" >&2
    echo -e "${YELLOW}   Consider using the same model for both operations.${NC}" >&2
    echo >&2
    read -p "Continue anyway? [y/N]: " continue_choice
    case "$continue_choice" in
        y|Y|yes|YES) 
            return 0
            ;;
        *) 
            echo -e "${GREEN}Good choice! Please select compatible models.${NC}" >&2
            return 1
            ;;
    esac
}

prompt_directory() {
    local prompt_text="$1"
    local validation_type="$2"
    
    while true; do
        read -p "$prompt_text: " dir_input
        dir_input="${dir_input/#\~/$HOME}"
        
        if [ -z "$dir_input" ]; then
            echo -e "${RED}Please enter a directory path${NC}" >&2
            continue
        fi
        
        if [ ! -d "$dir_input" ]; then
            echo -e "${RED}Directory does not exist: $dir_input${NC}" >&2
            continue
        fi
        
        case "$validation_type" in
            "embeddings")
                # Check if any files exist first (faster than counting)
                local first_file=$(find "$dir_input" -name "*-encodings.json" -type f -print -quit 2>/dev/null)
                if [ -z "$first_file" ]; then
                    echo -e "${RED}No embedding files found${NC}" >&2
                    continue
                fi
                # Only count if we found at least one file
                local embed_count=$(find "$dir_input" -name "*-encodings.json" -type f -printf '.' 2>/dev/null | wc -c)
                echo -e "${GREEN}Found $embed_count embedding files${NC}" >&2
                ;;
            "input")
                # Use find with -quit for existence check
                local first_file=$(find "$dir_input" -type f -print -quit 2>/dev/null)
                if [ -z "$first_file" ]; then
                    echo -e "${RED}No files found in directory${NC}" >&2
                    continue
                fi
                # Only count if we found at least one file
                local file_count=$(find "$dir_input" -type f -printf '.' 2>/dev/null | wc -c)
                echo -e "${GREEN}Found $file_count files${NC}" >&2
                ;;
        esac
        
        echo "$dir_input"
        return 0
    done
}

cleanup_files() {
    local target_dir="$1"
    local pattern="$2"
    local description="$3"
    
    [ ! -d "$target_dir" ] && return 0
    
    # Use find with -quit for quick existence check
    local first_file=$(find "$target_dir" -name "$pattern" -type f -print -quit 2>/dev/null)
    [ -z "$first_file" ] && return 0
    
    # Only count if we found at least one file
    local count=$(find "$target_dir" -name "$pattern" -type f -printf '.' 2>/dev/null | wc -c)
    echo -e "${YELLOW}Found $count existing $description files${NC}" >&2
    read -p "Remove them? [y/N]: " action
    case "$action" in
        y|Y|yes|YES) 
            find "$target_dir" -name "$pattern" -type f -delete 2>/dev/null
            echo -e "${GREEN}Removed $count files${NC}" >&2
            ;;
        *) 
            echo -e "${GREEN}Keeping existing files${NC}" >&2
            ;;
    esac
}

run_embed() {
    local input_dir="$1"
    local embed_model="$2"
    local config_file="$3"
    local dedupe_flag="$4"
    local newest_flag="$5"
    # Cache basename and dirname calculations
    local input_basename=$(basename "$input_dir")
    local input_dirname=$(dirname "$input_dir")
    local output_dir="${input_dirname}/${input_basename}_output"
    
    echo -e "${CYAN}Output directory: $output_dir${NC}" >&2
    cleanup_files "$output_dir" "*-encodings.json" "embedding"
    
    echo -e "${YELLOW}Creating embeddings with model: $embed_model${NC}" >&2
    
    # Build the command array
    local cmd=("$PROGRAM_PATH" --directory "$input_dir" --verbose)
    if [ -n "$embed_model" ] && [ "$embed_model" != "default" ]; then
        cmd+=(--model "$embed_model")
    fi
    
    # Extract and add base_url if config file is provided
    if [ -n "$config_file" ]; then
        # Trim whitespace from config_file path
        config_file=$(echo "$config_file" | tr -d '\n\r' | xargs)
        cmd+=(--config "$config_file")
        
        # Extract URL if file exists
        if [ -f "$config_file" ] && [ -r "$config_file" ]; then
            local base_url=$(extract_base_url "$config_file")
            if [ -n "$base_url" ]; then
                cmd+=(--url "$base_url")
            fi
        else
            echo -e "${RED}ERROR: Config file not found or not readable: $config_file${NC}" >&2
        fi
        
        # Show config info
        show_config_info "$config_file" "embed"
    fi
    
    # Add deduplication flags if provided
    if [ -n "$dedupe_flag" ]; then
        cmd+=("$dedupe_flag")
    fi
    if [ -n "$newest_flag" ]; then
        cmd+=("$newest_flag")
    fi
    
    # Display command and confirm
    display_command_and_confirm "$config_file" "${cmd[@]}"
    
    RUST_LOG=info run_with_logging "${cmd[@]}"
    echo "$output_dir"
}

run_retrieve() {
    local embed_dir="$1"
    local retrieve_model="$2"
    local prompt="${3:-Analyze and summarize the content}"
    local config_file="$4"
    
    cleanup_files "$embed_dir" "*_retrained_*.json" "retrieved"
    
    echo -e "${YELLOW}Retrieving embeddings with model: $retrieve_model${NC}" >&2
    echo -e "${YELLOW}Note: Some models may show 'not found in registry' warnings.${NC}" >&2
    echo -e "${YELLOW}This is normal and doesn't affect processing quality.${NC}" >&2
    echo "" >&2
    
    # Build the command array
    local cmd=("$PROGRAM_PATH" --retrieve-from "$embed_dir" --model "$retrieve_model" --prompt "$prompt" --verbose)
    
    # Use provided config file or check for qwen models
    if [ -n "$config_file" ]; then
        # Trim whitespace from config_file path
        config_file=$(echo "$config_file" | tr -d '\n\r' | xargs)
        cmd+=(--config "$config_file")
        
        if [ -f "$config_file" ] && [ -r "$config_file" ]; then
            local base_url=$(extract_base_url "$config_file")
            if [ -n "$base_url" ]; then
                cmd+=(--url "$base_url")
            fi
        fi
        
        # Show config info
        show_config_info "$config_file" "retrieve"
    elif [[ "$retrieve_model" == *"qwen"* ]]; then
        if [ -f "$SCRIPT_DIR/config/minimal-config.toml" ]; then
            config_file="$SCRIPT_DIR/config/minimal-config.toml"
            cmd+=(--config "$config_file")
            local base_url=$(extract_base_url "$config_file")
            if [ -n "$base_url" ]; then
                cmd+=(--url "$base_url")
            fi
            show_config_info "$config_file" "minimal"
        fi
    fi
    
    # Display command and confirm
    display_command_and_confirm "$config_file" "${cmd[@]}"
    
    RUST_LOG=info run_with_logging "${cmd[@]}"
}

workflow_full() {
    local input_dir=$(prompt_directory "Input directory" "input")
    
    echo
    echo -e "${BLUE}File Deduplication Options${NC}"
    echo -e "${CYAN}These options help handle duplicate files in your directory${NC}"
    echo
    echo -e "${YELLOW}--deduplicate-files:${NC} Remove duplicate file references (symlinks, hardlinks)"
    echo "  • Detects when the same physical file is accessed via different paths"
    echo "  • Useful for directories with symlinks or complex structures"
    echo
    echo -e "${YELLOW}--detect-newest:${NC} Keep only the newest file when multiple files have same name"
    echo "  • Compares modification timestamps across different directories"
    echo "  • Useful for backup/archive directories with multiple versions"
    echo
    read -p "Enable --deduplicate-files? [y/N]: " dedupe_choice
    local dedupe_flag=""
    case "$dedupe_choice" in
        y|Y|yes|YES) 
            dedupe_flag="--deduplicate-files"
            echo -e "${GREEN}✓ Path-based deduplication enabled${NC}"
            ;;
        *) 
            echo -e "${YELLOW}Path-based deduplication disabled${NC}"
            ;;
    esac
    
    read -p "Enable --detect-newest? [y/N]: " newest_choice
    local newest_flag=""
    case "$newest_choice" in
        y|Y|yes|YES) 
            newest_flag="--detect-newest"
            echo -e "${GREEN}✓ Filename-based newest detection enabled${NC}"
            ;;
        *) 
            echo -e "${YELLOW}Filename-based newest detection disabled${NC}"
            ;;
    esac
    
    echo
    echo -e "${BLUE}Model Selection${NC}"
    echo -e "${CYAN}You can use different models for embedding and retraining${NC}"
    show_embedding_models
    echo
    read -p "Embedding model [default]: " embed_model
    [ -z "$embed_model" ] && embed_model="default"
    
    echo
    show_retrieval_models "$embed_model"
    echo
    read -p "Retrieval model: " retrain_model
    
    validate_model_input "$retrain_model" "retrieval" || return 1
    
    echo
    local rag_prompt=$(prompt_for_rag)
    
    # Select config file
    local config_file=$(select_config_file)
    
    # Show summary
    show_workflow_summary "$embed_model" "$retrain_model" "$rag_prompt" "$config_file"
    
    # Check model compatibility
    warn_dimension_mismatch "$embed_model" "$retrain_model" || return 1
    
    local output_dir=$(run_embed "$input_dir" "$embed_model" "$config_file" "$dedupe_flag" "$newest_flag")
    run_retrieve "$output_dir" "$retrain_model" "$rag_prompt" "$config_file"
    
    echo -e "${GREEN}Full workflow completed!${NC}"
    echo -e "${CYAN}Compare embeddings from different models in: $output_dir${NC}"
    
    # Show log file info if it has content
    if [ -f "$NON_INFO_LOG" ] && [ -s "$NON_INFO_LOG" ]; then
        local log_lines=$(wc -l < "$NON_INFO_LOG")
        echo -e "${YELLOW}Non-INFO messages logged: $log_lines lines in $NON_INFO_LOG${NC}"
    fi
}

workflow_retrieve_only() {
    local embed_dir=$(prompt_directory "Embeddings directory" "embeddings")
    
    echo
    echo -e "${BLUE}Retrieval Model Selection${NC}"
    echo -e "${CYAN}Select a model for retrieving from existing embeddings${NC}"
    show_retrieval_models ""
    echo
    echo -e "${YELLOW}Note: Some models may show 'not found in registry' warnings.${NC}"
    echo -e "${YELLOW}This is normal and doesn't affect processing quality.${NC}"
    echo
    read -p "Retrieval model: " retrain_model
    
    validate_model_input "$retrain_model" "retrieval" || return 1
    
    echo
    local rag_prompt=$(prompt_for_rag)
    
    # Select config file
    local config_file=$(select_config_file)
    
    # Show summary
    show_workflow_summary "" "$retrain_model" "$rag_prompt" "$config_file"
    
    run_retrieve "$embed_dir" "$retrain_model" "$rag_prompt" "$config_file"
    
    echo -e "${GREEN}Retrieval completed!${NC}"
    echo -e "${CYAN}New retrieved embeddings created with model: $retrain_model${NC}"
    
    # Show log file info if it has content
    if [ -f "$NON_INFO_LOG" ] && [ -s "$NON_INFO_LOG" ]; then
        local log_lines=$(wc -l < "$NON_INFO_LOG")
        echo -e "${YELLOW}Non-INFO messages logged: $log_lines lines in $NON_INFO_LOG${NC}"
    fi
}

main() {
    if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
        echo "Usage: $0 [--help] [--tests]"
        echo "Interactive script for creating and retraining embeddings"
        exit 0
    fi
    
    if [ "$1" = "--tests" ]; then
        echo -e "${CYAN}Running test workflow${NC}"
        local test_input="$SCRIPT_DIR/tests/embed/test_files"
        local test_output="$SCRIPT_DIR/tests/embed/test_files_output"
        
        if [ ! -d "$test_input" ]; then
            echo -e "${RED}Test directory not found: $test_input${NC}"
            exit 1
        fi
        
        echo -e "${YELLOW}Creating test embeddings with nomic-embed-text:latest...${NC}"
        RUST_LOG=info run_with_logging "$PROGRAM_PATH" --directory "$test_input" --verbose --model "nomic-embed-text:latest"
        
        echo -e "${YELLOW}Retrieving test embeddings with nomic-embed-text:latest (matching dimensions)...${NC}"
        RUST_LOG=info run_with_logging "$PROGRAM_PATH" --retrieve-from "$test_output" --model "nomic-embed-text:latest" --prompt "What is this code about?" --verbose
        
        echo -e "${GREEN}Test workflow completed${NC}"
        
        # Show log file info if it has content
        if [ -f "$NON_INFO_LOG" ] && [ -s "$NON_INFO_LOG" ]; then
            local log_lines=$(wc -l < "$NON_INFO_LOG")
            echo -e "${YELLOW}Non-INFO messages logged: $log_lines lines in $NON_INFO_LOG${NC}"
        fi
        
        exit 0
    fi
    
    # Check server connectivity early
    check_ollama_server
    
    echo -e "${BLUE}Workflow Selection${NC}"
    echo "1. Full workflow - embed then retrieve with different models"
    echo "2. Retrieve only - use existing embeddings with new model"
    echo "3. Test mode - demo with different embed and retrieve models"
    
    while true; do
        read -p "Select workflow [1-3]: " workflow
        case "$workflow" in
            1)
                workflow_full && break
                ;;
            2)
                workflow_retrieve_only && break
                ;;
            3)
                exec "$0" --tests
                ;;
            *)
                echo -e "${RED}Please select 1, 2, or 3${NC}"
                ;;
        esac
    done
}

main "$@"