src/utils.py - master - vladelo/startups

src/utils.py
346 строк17 KB
vladelo
Make smaller logging
18 авг 2025, 17:21
18 авг 2025, 17:21cd68f3a
import pandas as pd
import logging
from typing import Dict, List, Any, Optional
import os
from pathlib import Path

logger = logging.getLogger(__name__)

class ExcelSearchStrategy:
    """Стратегия поиска в Excel файле с данными о стартапах"""
    
    def __init__(self, excel_path: str = "static_files/scouting_data.xlsx"):
        # Определяем абсолютный путь к файлу
        self.base_dir = Path(__file__).parent
        self.excel_path = self.base_dir / excel_path
        self._df = None
        
        # logger.info(f"Инициализация ExcelSearchStrategy с путем: {self.excel_path}")
    
    def _load_dataframe(self) -> pd.DataFrame:
        """Загружает DataFrame из Excel файла"""
        if self._df is None:
            try:
                # Проверяем существование файла
                if not self.excel_path.exists():
                    logger.error(f"Excel файл не найден: {self.excel_path}")
                    logger.error(f"Текущая директория: {os.getcwd()}")
                    logger.error(f"Содержимое текущей директории: {os.listdir('.')}")
                    
                    # Попробуем найти файл в других местах
                    alternative_paths = [
                        "scouting_data.xlsx",
                        "static_files/scouting_data.xlsx", 
                        "../static_files/scouting_data.xlsx",
                        "./static_files/scouting_data.xlsx"
                    ]
                    
                    for alt_path in alternative_paths:
                        if os.path.exists(alt_path):
                            logger.info(f"Найден файл по альтернативному пути: {alt_path}")
                            self.excel_path = Path(alt_path)
                            break
                    else:
                        logger.error("Excel файл не найден ни по одному из путей")
                        return pd.DataFrame()
                
                logger.info(f"Загружаем Excel файл: {self.excel_path}")
                self._df = pd.read_excel(self.excel_path)
                logger.info(f"Успешно загружен Excel файл: строк={len(self._df)}, колонок={len(self._df.columns)}")
                logger.info(f"Колонки в файле: {list(self._df.columns)}")
                
                # Заполняем NaN значения пустыми строками для корректного поиска
                self._df = self._df.fillna('')
                
            except Exception as e:
                logger.error(f"Ошибка загрузки Excel файла {self.excel_path}: {e}", exc_info=True)
                return pd.DataFrame()
        
        return self._df
    
    async def search(self, search_params: Dict[str, str], limit: Optional[int] = None, year_filters: Optional[Dict[str, int]] = None) -> List[Dict[str, Any]]:
        """
        Поиск стартапов по критериям
        
        Args:
            search_params: Словарь с критериями поиска {колонка: значение}
            limit: Максимальное количество результатов (None = все результаты)
            year_filters: Фильтры по годам {"год_от": 2020, "год_до": 2022}
            
        Returns:
            List[Dict[str, Any]]: Список найденных стартапов
        """
        try:
            # Загружаем данные
            df = self._load_dataframe()
            
            if df.empty:
                logger.warning("DataFrame пуст, возвращаем пустой список")
                return []
            
            logger.info(f"Начинаем поиск по критериям: {search_params}")
            if year_filters:
                logger.info(f"Фильтры по годам: {year_filters}")
            if limit:
                logger.info(f"Лимит результатов: {limit}")
            else:
                logger.info("Лимит не установлен - будут возвращены ВСЕ найденные результаты")
            logger.info(f"Общее количество записей в базе: {len(df)}")
            
            # Применяем фильтры по каждому критерию
            filtered_df = df.copy()
            
            # Сначала применяем обычные фильтры
            for column, search_value in search_params.items():
                if not search_value or search_value.strip() == "" or search_value == "0":
                    logger.info(f"Пропускаем пустой критерий: {column}")
                    continue
                
                # Проверяем существование колонки
                if column not in df.columns:
                    logger.warning(f"Колонка '{column}' не найдена в DataFrame. Доступные: {list(df.columns)}")
                    continue
                
                # Приводим к строке и выполняем поиск по подстроке (нечувствительно к регистру)
                search_value_lower = search_value.lower().strip()
                
                # Создаем маску для поиска
                column_values = filtered_df[column].astype(str).str.lower()
                mask = column_values.str.contains(
                    search_value_lower, 
                    case=False, 
                    na=False, 
                    regex=False
                )
                
                before_count = len(filtered_df)
                filtered_df = filtered_df[mask]
                after_count = len(filtered_df)
                
                logger.info(f"Фильтрация по '{column}' = '{search_value}': {before_count} → {after_count} строк")
                
                # Если после фильтрации не осталось результатов, прерываем
                if filtered_df.empty:
                    logger.info("После фильтрации результатов не осталось")
                    break
            
            # Применяем фильтры по годам
            if year_filters and not filtered_df.empty:
                year_column = 'Год регистрации'
                
                if year_column in filtered_df.columns:
                    before_count = len(filtered_df)
                    
                    # Преобразуем колонку года в числовой формат
                    filtered_df[year_column] = pd.to_numeric(filtered_df[year_column], errors='coerce')
                    
                    # Применяем фильтр "от года"
                    if 'год_от' in year_filters:
                        year_from = year_filters['год_от']
                        filtered_df = filtered_df[filtered_df[year_column] >= year_from]
                        logger.info(f"Фильтр 'от {year_from} года': {before_count} → {len(filtered_df)} строк")
                        before_count = len(filtered_df)
                    
                    # Применяем фильтр "до года"
                    if 'год_до' in year_filters:
                        year_to = year_filters['год_до']
                        filtered_df = filtered_df[filtered_df[year_column] <= year_to]
                        logger.info(f"Фильтр 'до {year_to} года': {before_count} → {len(filtered_df)} строк")
                else:
                    logger.warning(f"Колонка '{year_column}' не найдена для фильтрации по годам")
            
            # Применяем лимит только если он явно задан
            if limit is not None and limit > 0:
                result_df = filtered_df.head(limit)
                logger.info(f"Применен лимит {limit}: показываем {len(result_df)} из {len(filtered_df)} найденных результатов")
            else:
                result_df = filtered_df
                logger.info(f"Лимит не применяется: возвращаем ВСЕ {len(result_df)} найденных результатов")
            
            # Преобразуем в список словарей
            results = result_df.to_dict('records')
            
            logger.info(f"Поиск завершен. Всего найдено {len(filtered_df)} результатов, возвращаем {len(results)}")
            
            return results
            
        except Exception as e:
            logger.error(f"Ошибка при поиске в Excel: {e}", exc_info=True)
            return []
    
    async def get_all_startups(self) -> List[Dict[str, Any]]:
        """
        Возвращает все стартапы из базы данных без фильтрации
        Используется для поиска по ключевым словам
        
        Returns:
            List[Dict[str, Any]]: Список всех стартапов
        """
        try:
            logger.info("Получаем ВСЕ стартапы из базы данных для поиска по ключевым словам")
            
            # Загружаем данные
            df = self._load_dataframe()
            
            if df.empty:
                logger.warning("DataFrame пуст, возвращаем пустой список")
                return []
            
            logger.info(f"Загружено {len(df)} записей из базы данных")
            
            # Преобразуем в список словарей
            results = df.to_dict('records')
            
            logger.info(f"Возвращаем {len(results)} стартапов для поиска по ключевым словам")
            
            return results
            
        except Exception as e:
            logger.error(f"Ошибка при получении всех стартапов: {e}", exc_info=True)
            return []
    
    def get_column_names(self) -> List[str]:
        """Возвращает список названий колонок в Excel файле"""
        try:
            df = self._load_dataframe()
            return list(df.columns) if not df.empty else []
        except Exception as e:
            logger.error(f"Ошибка получения названий колонок: {e}")
            return []
    
    def get_total_records(self) -> int:
        """Возвращает общее количество записей в Excel файле"""
        try:
            df = self._load_dataframe()
            return len(df) if not df.empty else 0
        except Exception as e:
            logger.error(f"Ошибка получения количества записей: {e}")
            return 0


class ExcelFileManager:
    """Менеджер для работы с Excel файлами"""
    
    def __init__(self):
        self.search_strategy = ExcelSearchStrategy()
    
    async def search_startups(self, criteria: Dict[str, str], limit: Optional[int] = None, year_filters: Optional[Dict[str, int]] = None) -> List[Dict[str, Any]]:
        """
        Поиск стартапов через стратегию поиска
        
        Args:
            criteria: Критерии поиска
            limit: Лимит результатов (None = все результаты)
            year_filters: Фильтры по годам
            
        Returns:
            Список найденных стартапов
        """
        return await self.search_strategy.search(criteria, limit, year_filters)
    
    async def get_all_startups(self) -> List[Dict[str, Any]]:
        """
        Получить все стартапы из базы данных
        
        Returns:
            Список всех стартапов
        """
        return await self.search_strategy.get_all_startups()
    
    def get_available_columns(self) -> List[str]:
        """Получить доступные колонки для поиска"""
        return self.search_strategy.get_column_names()
    
    def get_database_info(self) -> Dict[str, Any]:
        """Получить информацию о базе данных"""
        return {
            'total_records': self.search_strategy.get_total_records(),
            'columns': self.search_strategy.get_column_names(),
            'file_path': str(self.search_strategy.excel_path)
        }


# Функция для проверки доступности файла
def check_excel_file_availability():
    """Проверяет доступность Excel файла и выводит диагностическую информацию"""
    logger.info("=== ДИАГНОСТИКА EXCEL ФАЙЛА ===")
    logger.info(f"Текущая рабочая директория: {os.getcwd()}")
    logger.info(f"Содержимое текущей директории: {os.listdir('.')}")
    
    # Проверяем различные пути
    paths_to_check = [
        "static_files/scouting_data.xlsx",
        "scouting_data.xlsx", 
        "./static_files/scouting_data.xlsx",
        "../static_files/scouting_data.xlsx"
    ]
    
    for path in paths_to_check:
        exists = os.path.exists(path)
        logger.info(f"Путь '{path}': {'СУЩЕСТВУЕТ' if exists else 'НЕ НАЙДЕН'}")
        
        if exists:
            try:
                df = pd.read_excel(path)
                logger.info(f"  ✓ Файл успешно читается, строк: {len(df)}")
                return path
            except Exception as e:
                logger.error(f"  ✗ Ошибка чтения файла: {e}")
    
    logger.error("Excel файл не найден ни по одному из путей!")
    return None


# Пример использования:
if __name__ == "__main__":
    import asyncio
    
    async def test_search():
        # Проверяем доступность файла
        file_path = check_excel_file_availability()
        
        if not file_path:
            print("Excel файл недоступен!")
            return
        
        # Создаем менеджер
        excel_manager = ExcelFileManager()
        
        # Получаем информацию о базе
        db_info = excel_manager.get_database_info()
        print(f"Информация о базе: {db_info}")
        
        # Пример поиска БЕЗ лимита (все результаты)
        search_criteria = {
            'Город': 'Москва',
            'Индустрии': 'финтех'
        }
        
        print("=== ПОИСК БЕЗ ЛИМИТА (все результаты) ===")
        results = await excel_manager.search_startups(search_criteria)  # limit не указан
        print(f"Найдено результатов: {len(results)}")
        
        for i, startup in enumerate(results):
            print(f"{i+1}. {startup.get('Наименование организации', 'Не указано')}")
            print(f"   Город: {startup.get('Город', 'Не указан')}")
            print(f"   Индустрия: {startup.get('Индустрии', 'Не указана')}")
            print()
        
        print("\n=== ПОИСК С ЛИМИТОМ (только 3 результата) ===")
        limited_results = await excel_manager.search_startups(search_criteria, limit=3)
        print(f"Найдено результатов с лимитом: {len(limited_results)}")
        
        for i, startup in enumerate(limited_results):
            print(f"{i+1}. {startup.get('Наименование организации', 'Не указано')}")
        
        print("\n=== ТЕСТ ПОЛУЧЕНИЯ ВСЕХ СТАРТАПОВ ===")
        all_startups = await excel_manager.get_all_startups()
        print(f"Всего стартапов в базе: {len(all_startups)}")
        
        if all_startups:
            print("Первые 3 стартапа:")
            for i, startup in enumerate(all_startups[:3]):
                print(f"{i+1}. {startup.get('Наименование организации', 'Не указано')}")
    
    # Запуск теста
    asyncio.run(test_search())