hotels/process_orel_embeddings.py

#!/usr/bin/env python3
"""
Обработка chunks и embeddings только для Орловской области
"""

import psycopg2
from urllib.parse import unquote
import requests
import json
import time
import logging
from typing import List, Dict, Tuple
import uuid

# Настройка логирования
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('orel_embeddings.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Конфигурация
BGE_API_URL = "http://147.45.146.17:8002/embed"
BGE_API_KEY = "22564b177aa73b6ac0b8642d7773350ff4c01d4983f028beff15ea247f09fa89"
CHUNK_SIZE = 600
CHUNK_OVERLAP = 100
BATCH_SIZE = 8
MAX_RETRIES = 3

class EmbeddingProcessor:
    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.cur = self.conn.cursor()
            logger.info("✅ Подключение к БД установлено")
        except Exception as e:
            logger.error(f"❌ Ошибка подключения к БД: {e}")
            raise

    def create_chunks(self, text: str) -> List[str]:
        """Создание chunks из текста"""
        if not text or len(text.strip()) < 50:
            return []
        
        chunks = []
        start = 0
        
        while start < len(text):
            end = start + CHUNK_SIZE
            
            if end >= len(text):
                chunks.append(text[start:].strip())
                break
                
            # Ищем ближайший пробел или перенос строки
            while end > start and text[end] not in [' ', '\n', '\t']:
                end -= 1
                
            if end == start:  # Если не нашли пробел, берем по символам
                end = start + CHUNK_SIZE
                
            chunk = text[start:end].strip()
            if chunk:
                chunks.append(chunk)
                
            start = end - CHUNK_OVERLAP
            
        return chunks

    def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
        """Получение эмбеддингов для батча текстов"""
        for attempt in range(MAX_RETRIES):
            try:
                response = requests.post(
                    BGE_API_URL,
                    headers={
                        'Authorization': f'Bearer {BGE_API_KEY}',
                        'Content-Type': 'application/json'
                    },
                    json={'text': texts},
                    timeout=30
                )
                
                if response.status_code == 200:
                    result = response.json()
                    return result.get('embeddings', [])
                else:
                    logger.warning(f"⚠️ API вернул статус {response.status_code}: {response.text}")
                    
            except Exception as e:
                logger.warning(f"⚠️ Попытка {attempt + 1} неудачна: {e}")
                if attempt < MAX_RETRIES - 1:
                    time.sleep(2 ** attempt)  # Экспоненциальная задержка
                    
        logger.error(f"❌ Не удалось получить эмбеддинги для батча из {len(texts)} текстов")
        return []

    def process_hotel(self, hotel_id: str) -> bool:
        """Обработка одного отеля"""
        try:
            # Получаем HTML для отеля
            self.cur.execute("""
                SELECT html FROM hotel_website_raw 
                WHERE hotel_id = %s
            """, (hotel_id,))
            
            result = self.cur.fetchone()
            if not result or not result[0]:
                logger.warning(f"⚠️ Нет HTML для отеля {hotel_id}")
                return False
                
            html = result[0]
            
            # Очищаем HTML до текста
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(html, 'html.parser')
            for script in soup(["script", "style"]):
                script.decompose()
            text = soup.get_text()
            
            # Создаем chunks
            chunks = self.create_chunks(text)
            if not chunks:
                logger.warning(f"⚠️ Нет chunks для отеля {hotel_id}")
                return False
                
            logger.info(f"📄 Создано {len(chunks)} chunks")
            
            # Удаляем старые chunks
            self.cur.execute("DELETE FROM hotel_website_chunks WHERE metadata->>'hotel_id' = %s", (hotel_id,))
            
            # Обрабатываем chunks батчами
            total_chunks = 0
            for i in range(0, len(chunks), BATCH_SIZE):
                batch = chunks[i:i + BATCH_SIZE]
                logger.info(f"🔄 Обрабатываем батч {i//BATCH_SIZE + 1}: {len(batch)} chunks")
                
                embeddings = self.get_embeddings_batch(batch)
                if not embeddings:
                    logger.error(f"❌ Не удалось получить эмбеддинги для батча")
                    continue
                    
                # Сохраняем chunks с эмбеддингами
                for j, (chunk, embedding) in enumerate(zip(batch, embeddings)):
                    chunk_id = str(uuid.uuid4())
                    metadata = {
                        'hotel_id': hotel_id,
                        'chunk_index': i + j,
                        'created_at': time.time()
                    }
                    self.cur.execute("""
                        INSERT INTO hotel_website_chunks (id, text, metadata, embedding)
                        VALUES (%s, %s, %s, %s::vector)
                    """, (chunk_id, chunk, json.dumps(metadata), json.dumps(embedding)))
                    total_chunks += 1
                    
                logger.info(f"✅ Батч успешно обработан")
                
            self.conn.commit()
            logger.info(f"✅ Сохранено {total_chunks} chunks для отеля")
            return True
            
        except Exception as e:
            logger.error(f"❌ Ошибка обработки отеля {hotel_id}: {e}")
            self.conn.rollback()
            return False

    def process_orel_region(self):
        """Обработка всех отелей Орловской области"""
        try:
            # Получаем отели Орловской области с HTML но без chunks
            self.cur.execute("""
                SELECT DISTINCT h.id, h.full_name
                FROM hotel_main h
                INNER JOIN hotel_website_raw hwr ON h.id = hwr.hotel_id
                LEFT JOIN hotel_website_chunks hc ON h.id::text = hc.metadata->>'hotel_id'
                WHERE h.region_name = 'Орловская область'
                AND hwr.html IS NOT NULL
                AND hc.id IS NULL
                ORDER BY h.full_name
            """)
            
            hotels = self.cur.fetchall()
            logger.info(f"📊 Найдено {len(hotels)} отелей для обработки")
            
            if not hotels:
                logger.info("✅ Все отели Орловской области уже обработаны!")
                return
                
            for i, (hotel_id, hotel_name) in enumerate(hotels, 1):
                logger.info(f"🔄 Обрабатываем отель {i}/{len(hotels)}: {hotel_name}")
                success = self.process_hotel(hotel_id)
                if success:
                    logger.info(f"✅ Отель {hotel_name} обработан успешно")
                else:
                    logger.error(f"❌ Ошибка обработки отеля {hotel_name}")
                    
        except Exception as e:
            logger.error(f"❌ Ошибка обработки региона: {e}")

    def close(self):
        """Закрытие соединения с БД"""
        if self.cur:
            self.cur.close()
        if self.conn:
            self.conn.close()


def main():
    logger.info("🚀 Запуск обработки Орловской области")
    
    processor = EmbeddingProcessor()
    try:
        processor.process_orel_region()
        logger.info("✅ Обработка завершена!")
    finally:
        processor.close()


if __name__ == "__main__":
    main()