hotels/test_semantic_search_chukotka.py

#!/usr/bin/env python3
"""
Тест семантического поиска по Чукотскому автономному округу
на основе готовой базы с эмбеддингами
"""

import psycopg2
from urllib.parse import unquote
import requests
import json
import time

# API настройки
BGE_API_URL = "http://147.45.146.17:8002/embed"
BGE_API_KEY = "22564b177aa73b6ac0b8642d7773350ff4c01d4983f028beff15ea247f09fa89"

class ChukotkaAnalyzer:
    def __init__(self):
        self.conn = None
        self.cur = None
        self.connect_db()
        
    def connect_db(self):
        """Подключение к базе данных"""
        try:
            self.conn = psycopg2.connect(
                host='147.45.189.234',
                port=5432,
                database='default_db',
                user='gen_user',
                password=unquote('2~~9_%5EkVsU%3F2%5CS')
            )
            self.conn.autocommit = True
            self.cur = self.conn.cursor()
            print("✅ Подключение к БД установлено")
        except Exception as e:
            print(f"❌ Ошибка подключения к БД: {e}")
            raise
    
    def get_chukotka_stats(self):
        """Получение статистики по Чукотке"""
        self.cur.execute("""
            SELECT 
                COUNT(DISTINCT metadata->>'hotel_id') as hotels_count,
                COUNT(*) as total_chunks,
                AVG(LENGTH(text)) as avg_chunk_length
            FROM hotel_website_chunks
            WHERE metadata->>'region_name' = 'Чукотский автономный округ';
        """)
        
        result = self.cur.fetchone()
        return {
            'hotels_count': result[0],
            'total_chunks': result[1],
            'avg_chunk_length': result[2]
        }
    
    def get_chukotka_hotels(self):
        """Получение списка отелей Чукотки"""
        self.cur.execute("""
            SELECT DISTINCT
                metadata->>'hotel_id' as hotel_id,
                metadata->>'hotel_name' as hotel_name,
                COUNT(*) as chunks_count
            FROM hotel_website_chunks
            WHERE metadata->>'region_name' = 'Чукотский автономный округ'
            GROUP BY metadata->>'hotel_id', metadata->>'hotel_name'
            ORDER BY chunks_count DESC;
        """)
        
        return self.cur.fetchall()
    
    def generate_query_embedding(self, query: str):
        """Генерация эмбеддинга для поискового запроса"""
        try:
            headers = {
                "X-API-Key": BGE_API_KEY,
                "Content-Type": "application/json"
            }
            
            payload = {"text": query}
            response = requests.post(BGE_API_URL, json=payload, headers=headers, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                return result.get('embeddings', [[]])[0]
            else:
                print(f"❌ Ошибка API: {response.status_code}")
                return None
                
        except Exception as e:
            print(f"❌ Ошибка генерации эмбеддинга: {e}")
            return None
    
    def search_chukotka(self, query: str, limit: int = 5):
        """Семантический поиск по Чукотке"""
        query_embedding = self.generate_query_embedding(query)
        if not query_embedding:
            return []
        
        embedding_str = json.dumps(query_embedding)
        
        self.cur.execute("""
            SELECT 
                metadata->>'hotel_name' as hotel_name,
                metadata->>'url' as url,
                LEFT(text, 150) as sample_text,
                LENGTH(text) as text_length,
                embedding <-> %s::vector as distance
            FROM hotel_website_chunks
            WHERE metadata->>'region_name' = 'Чукотский автономный округ'
            AND embedding IS NOT NULL
            ORDER BY embedding <-> %s::vector
            LIMIT %s;
        """, (embedding_str, embedding_str, limit))
        
        return self.cur.fetchall()
    
    def analyze_hotel_criteria(self, hotel_id: str):
        """Анализ отеля по критериям аудита"""
        criteria_queries = {
            'Юридическая идентификация': 'инн огрн егрюл организация',
            'Контактная информация': 'телефон адрес email контакты',
            'Политика конфиденциальности': 'политика конфиденциальности персональные данные',
            'Условия бронирования': 'бронирование условия отмена возврат',
            'Услуги отеля': 'услуги сервис завтрак wi-fi парковка',
            'Доступность': 'доступность инвалиды коляска лифт'
        }
        
        results = {}
        
        for criteria, query in criteria_queries.items():
            self.cur.execute("""
                SELECT 
                    embedding <-> %s::vector as distance,
                    LEFT(text, 200) as sample_text
                FROM hotel_website_chunks
                WHERE metadata->>'hotel_id' = %s
                AND embedding IS NOT NULL
                ORDER BY embedding <-> %s::vector
                LIMIT 1;
            """, (json.dumps(self.generate_query_embedding(query)), hotel_id, json.dumps(self.generate_query_embedding(query))))
            
            result = self.cur.fetchone()
            if result:
                distance, sample_text = result
                relevance = "🟢 Высокая" if distance < 0.9 else "🟡 Средняя" if distance < 1.0 else "🔴 Низкая"
                results[criteria] = {
                    'distance': distance,
                    'relevance': relevance,
                    'sample_text': sample_text
                }
        
        return results
    
    def close(self):
        """Закрытие соединения с БД"""
        if self.cur:
            self.cur.close()
        if self.conn:
            self.conn.close()

def main():
    print("="*70)
    print("🏔️ АНАЛИЗ ЧУКОТСКОГО АВТОНОМНОГО ОКРУГА")
    print("="*70)
    
    analyzer = ChukotkaAnalyzer()
    
    try:
        # Статистика по региону
        stats = analyzer.get_chukotka_stats()
        print(f"\n📊 СТАТИСТИКА ПО ЧУКОТКЕ:")
        print(f"   Отелей: {stats['hotels_count']}")
        print(f"   Chunks: {stats['total_chunks']}")
        print(f"   Средняя длина chunk: {stats['avg_chunk_length']:.0f} символов")
        
        # Список отелей
        hotels = analyzer.get_chukotka_hotels()
        print(f"\n🏨 ОТЕЛИ ЧУКОТКИ:")
        print("-" * 70)
        for hotel_id, hotel_name, chunks_count in hotels:
            print(f"   🏨 {hotel_name}")
            print(f"      ID: {hotel_id}")
            print(f"      Chunks: {chunks_count}")
            print()
        
        # Тестовые поисковые запросы
        test_queries = [
            "телефон отеля",
            "услуги и сервисы",
            "бронирование номеров",
            "адрес и контакты",
            "политика конфиденциальности",
            "завтрак и питание"
        ]
        
        print("🔍 ТЕСТИРОВАНИЕ СЕМАНТИЧЕСКОГО ПОИСКА:")
        print("-" * 70)
        
        for query in test_queries:
            print(f"\n🔍 Запрос: '{query}'")
            results = analyzer.search_chukotka(query, 3)
            
            for i, (hotel_name, url, sample_text, text_length, distance) in enumerate(results, 1):
                if distance < 0.9:
                    relevance = "🟢 Отлично"
                elif distance < 1.0:
                    relevance = "🟡 Хорошо"
                else:
                    relevance = "🔴 Слабо"
                
                print(f"   {i}. Distance: {distance:.4f} {relevance}")
                print(f"      Отель: {hotel_name[:50]}...")
                print(f"      Текст: {sample_text}...")
                print()
        
        # Анализ одного отеля по критериям
        if hotels:
            test_hotel_id, test_hotel_name, _ = hotels[0]
            print(f"\n📋 АНАЛИЗ ОТЕЛЯ ПО КРИТЕРИЯМ:")
            print(f"🏨 {test_hotel_name}")
            print("-" * 70)
            
            criteria_results = analyzer.analyze_hotel_criteria(test_hotel_id)
            
            for criteria, data in criteria_results.items():
                print(f"{data['relevance']} {criteria}")
                print(f"   Distance: {data['distance']:.4f}")
                print(f"   Найденный текст: {data['sample_text'][:100]}...")
                print()
        
        print("="*70)
        print("✅ АНАЛИЗ ЗАВЕРШЁН!")
        print("="*70)
        
    except Exception as e:
        print(f"❌ Ошибка анализа: {e}")
    finally:
        analyzer.close()

if __name__ == "__main__":
    main()