Update app.py

app.py

@@ -1,17 +1,23 @@
-import gradio as gr
-from sentence_transformers import SentenceTransformer, util
 import os
 import time
 import threading
 import queue
+from typing import List, Dict, Any, Optional
+import logging
+from urllib.parse import urlparse
+
+import gradio as gr
 import torch
 import psycopg2
 import zlib
 import numpy as np
-from urllib.parse import urlparse
-import logging
+from sentence_transformers import SentenceTransformer, util
 from sklearn.preprocessing import normalize
 
+# Рекомендуется использовать python-dotenv для загрузки переменных окружения
+# from dotenv import load_dotenv
+# load_dotenv()
+
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
@@ -54,12 +60,10 @@ except FileNotFoundError:
     movies_data = []
 
 # Очередь для необработанных фильмов
-movies_queue = queue.Queue()
+movies_queue: queue.Queue = queue.Queue()
 
-# Флаг, указывающий, что обработка фильмов завершена
+# Флаги состояния
 processing_complete = False
-
-# Флаг, указывающий, что выполняется поиск
 search_in_progress = False
 
 # Блокировка для доступа к базе данных
@@ -79,20 +83,19 @@ def get_db_connection():
 
 def setup_database():
     """Настраивает базу данных: создает расширение, таблицы и индексы."""
-    conn = get_db_connection()
-    if conn is None:
-        return
-
-    try:
-        with conn.cursor() as cur:
-            # Создаем расширение pgvector если его нет
-            cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
-
-            # Удаляем существующие таблицы если они есть
-            cur.execute(f"DROP TABLE IF EXISTS {embeddings_table}, {query_cache_table};")
-
-            # Создаем таблицу для хранения эмбеддингов фильмов
-            cur.execute(f"""
+    with get_db_connection() as conn:
+        if conn is None:
+            return
+        try:
+            with conn.cursor() as cur:
+                # Создаем расширение pgvector если его нет
+                cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
+                
+                # Удаляем существующие таблицы если они есть
+                # cur.execute(f"DROP TABLE IF EXISTS {embeddings_table}, {query_cache_table};")
+                
+                # Создаем таблицу для хранения эмбеддингов фильмов
+                cur.execute(f"""
                 CREATE TABLE {embeddings_table} (
                     movie_id INTEGER PRIMARY KEY,
                     embedding_crc32 BIGINT,
@@ -101,10 +104,10 @@ def setup_database():
                     embedding vector(1024)
                 );
                 CREATE INDEX ON {embeddings_table} (string_crc32);
-            """)
-
-            # Создаем таблицу для кэширования запросов
-            cur.execute(f"""
+                """)
+                
+                # Создаем таблицу для кэширования запросов
+                cur.execute(f"""
                 CREATE TABLE {query_cache_table} (
                     query_crc32 BIGINT PRIMARY KEY,
                     query TEXT,
@@ -114,60 +117,51 @@ def setup_database():
                 );
                 CREATE INDEX ON {query_cache_table} (query_crc32);
                 CREATE INDEX ON {query_cache_table} (created_at);
-            """)
-
-        conn.commit()
-        logging.info("База данных успешно настроена.")
-    except Exception as e:
-        logging.error(f"Ошибка при настройке базы данных: {e}")
-        conn.rollback()
-    finally:
-        conn.close()
+                """)
+            
+            conn.commit()
+            logging.info("База данных успешно настроена.")
+        except Exception as e:
+            logging.error(f"Ошибка при настройке базы данных: {e}")
+            conn.rollback()
 
 # Настраиваем базу данных при запуске
 setup_database()
 
-def calculate_crc32(text):
+def calculate_crc32(text: str) -> int:
     """Вычисляет CRC32 для строки."""
     return zlib.crc32(text.encode('utf-8')) & 0xFFFFFFFF
 
-def encode_string(text):
+def encode_string(text: str) -> np.ndarray:
     """Кодирует строку в эмбеддинг."""
     embedding = model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
     return embedding.cpu().numpy()
 
-def get_movies_without_embeddings():
+def get_movies_without_embeddings() -> List[Dict[str, Any]]:
     """Получает список фильмо��, для которых нужно создать эмбеддинги."""
-    conn = get_db_connection()
-    if conn is None:
-        return []
-
-    movies_to_process = []
-    try:
-        with conn.cursor() as cur:
-            # Получаем список ID фильмов, которые уже есть в базе
-            cur.execute(f"SELECT movie_id FROM {embeddings_table}")
-            existing_ids = {row[0] for row in cur.fetchall()}
-
-            # Фильтруем только те фильмы, которых нет в базе
-            for movie in movies_data:
-                if movie['id'] not in existing_ids:
-                    movies_to_process.append(movie)
-
-        logging.info(f"Найдено {len(movies_to_process)} фильмов для обработки.")
-    except Exception as e:
-        logging.error(f"Ошибка при получении списка фильмов для обработки: {e}")
-    finally:
-        conn.close()
-
-    return movies_to_process
+    with get_db_connection() as conn:
+        if conn is None:
+            return []
+        try:
+            with conn.cursor() as cur:
+                # Получаем список ID фильмов, которые уже есть в базе
+                cur.execute(f"SELECT movie_id FROM {embeddings_table}")
+                existing_ids = {row[0] for row in cur.fetchall()}
+                
+                # Фильтруем только те фильмы, которых нет в базе
+                movies_to_process = [movie for movie in movies_data if movie['id'] not in existing_ids]
+                
+            logging.info(f"Найдено {len(movies_to_process)} фильмов для обработки.")
+            return movies_to_process
+        except Exception as e:
+            logging.error(f"Ошибка при получении списка фильмов для обработки: {e}")
+            return []
 
-def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_name):
+def get_embedding_from_db(conn, table_name: str, crc32_column: str, crc32_value: int, model_name: str) -> Optional[np.ndarray]:
     """Получает эмбеддинг из базы данных."""
     try:
         with conn.cursor() as cur:
-            cur.execute(f"SELECT embedding FROM {table_name} WHERE {crc32_column} = %s AND model_name = %s", 
-                       (crc32_value, model_name))
+            cur.execute(f"SELECT embedding FROM {table_name} WHERE {crc32_column} = %s AND model_name = %s", (crc32_value, model_name))
             result = cur.fetchone()
             if result and result[0]:
                 # Нормализуем эмбеддинг после извлечения из БД
@@ -176,17 +170,16 @@ def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_nam
         logging.error(f"Ошибка при получении эмбеддинга из БД: {e}")
     return None
 
-def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32, embedding):
+def insert_embedding(conn, table_name: str, movie_id: int, embedding_crc32: int, string_crc32: int, embedding: np.ndarray) -> bool:
     """Вставляет эмбеддинг в базу данных."""
     try:
         # Нормализуем эмбеддинг перед сохранением
         normalized_embedding = normalize(embedding.reshape(1, -1))[0]
         with conn.cursor() as cur:
             cur.execute(f"""
-                INSERT INTO {table_name} 
-                (movie_id, embedding_crc32, string_crc32, model_name, embedding)
-                VALUES (%s, %s, %s, %s, %s)
-                ON CONFLICT (movie_id) DO NOTHING
+            INSERT INTO {table_name} (movie_id, embedding_crc32, string_crc32, model_name, embedding)
+            VALUES (%s, %s, %s, %s, %s)
+            ON CONFLICT (movie_id) DO NOTHING
             """, (movie_id, embedding_crc32, string_crc32, model_name, normalized_embedding.tolist()))
         conn.commit()
         return True
@@ -198,12 +191,10 @@ def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32,
 def process_movies():
     """Обрабатывает фильмы, создавая для них эмбеддинги."""
     global processing_complete
-
     logging.info("Начало обработки фильмов.")
-
+    
     # Получаем список фильмов, которые нужно обработать
     movies_to_process = get_movies_without_embeddings()
-
     if not movies_to_process:
         logging.info("Все фильмы уже обработаны.")
         processing_complete = True
@@ -213,55 +204,51 @@ def process_movies():
     for movie in movies_to_process:
         movies_queue.put(movie)
 
-    conn = get_db_connection()
-    if conn is None:
-        processing_complete = True
-        return
+    with get_db_connection() as conn:
+        if conn is None:
+            processing_complete = True
+            return
+        try:
+            while not movies_queue.empty():
+                if search_in_progress:
+                    time.sleep(1)
+                    continue
+
+                batch = []
+                while not movies_queue.empty() and len(batch) < batch_size:
+                    try:
+                        movie = movies_queue.get_nowait()
+                        batch.append(movie)
+                    except queue.Empty:
+                        break
 
-    try:
-        while not movies_queue.empty():
-            if search_in_progress:
-                time.sleep(1)
-                continue
-
-            batch = []
-            while not movies_queue.empty() and len(batch) < batch_size:
-                try:
-                    movie = movies_queue.get_nowait()
-                    batch.append(movie)
-                except queue.Empty:
+                if not batch:
                     break
 
-            if not batch:
-                break
-
-            logging.info(f"Обработка пакета из {len(batch)} фильмов...")
-
-            for movie in batch:
-                embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
-                string_crc32 = calculate_crc32(embedding_string)
-
-                # Проверяем существующий эмбеддинг
-                existing_embedding = get_embedding_from_db(conn, embeddings_table, "string_crc32", string_crc32, model_name)
-
-                if existing_embedding is None:
-                    embedding = encode_string(embedding_string)
-                    embedding_crc32 = calculate_crc32(str(embedding.tolist()))
-
-                    if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
-                        logging.info(f"Сохранен эмбеддинг для '{movie['name']}'")
+                logging.info(f"Обработка пакета из {len(batch)} фильмов...")
+                for movie in batch:
+                    embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
+                    string_crc32 = calculate_crc32(embedding_string)
+
+                    # Проверяем существующий эмбеддинг
+                    existing_embedding = get_embedding_from_db(conn, embeddings_table, "string_crc32", string_crc32, model_name)
+                    if existing_embedding is None:
+                        embedding = encode_string(embedding_string)
+                        embedding_crc32 = calculate_crc32(str(embedding.tolist()))
+                        if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
+                            logging.info(f"Сохранен эмбеддинг для '{movie['name']}'")
+                        else:
+                            logging.error(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
                     else:
-                        logging.error(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
-                else:
-                    logging.info(f"Эмбеддинг для '{movie['name']}' уже существует")
-    except Exception as e:
-        logging.error(f"Ошибка при обработке фильмов: {e}")
-    finally:
-        conn.close()
-        processing_complete = True
-        logging.info("Обработка фильмов завершена")
+                        logging.info(f"Эмбеддинг для '{movie['name']}' уже существует")
 
-def get_movie_embeddings(conn):
+        except Exception as e:
+            logging.error(f"Ошибка при обработке фильмов: {e}")
+        finally:
+            processing_complete = True
+            logging.info("Обработка фильмов завершена")
+
+def get_movie_embeddings(conn) -> Dict[str, np.ndarray]:
     """Загружает все эмбеддинги фильмов из базы данных."""
     movie_embeddings = {}
     try:
@@ -278,96 +265,97 @@ def get_movie_embeddings(conn):
         logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
     return movie_embeddings
 
-def search_movies(query, top_k=10):
+def clean_query_cache(conn):
+    """Очищает устаревшие записи из кэша запросов."""
+    try:
+        with conn.cursor() as cur:
+            # Получаем общий размер кэша
+            cur.execute(f"SELECT pg_total_relation_size('{query_cache_table}')")
+            total_size = cur.fetchone()[0]
+            
+            if total_size > MAX_CACHE_SIZE:
+                # Удаляем старые записи, пока размер не станет меньше максимального
+                cur.execute(f"""
+                DELETE FROM {query_cache_table}
+                WHERE ctid IN (
+                    SELECT ctid
+                    FROM {query_cache_table}
+                    ORDER BY created_at ASC
+                    LIMIT (SELECT COUNT(*) / 2 FROM {query_cache_table})
+                )
+                """)
+                conn.commit()
+                logging.info("Кэш запросов очищен.")
+    except Exception as e:
+        logging.error(f"Ошибка при очистке кэша запросов: {e}")
+        conn.rollback()
+
+def search_movies(query: str, top_k: int = 10) -> List[Dict[str, Any]]:
     """Выполняет поиск фильмов по запросу."""
     global search_in_progress
     search_in_progress = True
-    start_time = time.time()
-
+    
     try:
-        conn = get_db_connection()
-        if conn is None:
-            return "<p>Ошибка подключения к базе данных</p>"
+        with get_db_connection() as conn:
+            if conn is None:
+                return []
 
-        query_crc32 = calculate_crc32(query)
-        query_embedding = get_embedding_from_db(conn, query_cache_table, "query_crc32", query_crc32, model_name)
+            clean_query_cache(conn)
 
-        if query_embedding is None:
-            query_embedding = encode_string(query)
+            query_crc32 = calculate_crc32(query)
+            query_embedding = get_embedding_from_db(conn, query_cache_table, "query_crc32", query_crc32, model_name)
 
-            try:
-                with conn.cursor() as cur:
-                    cur.execute(f"""
-                        INSERT INTO {query_cache_table} (query_crc32, query, model_name, embedding)
-                        VALUES (%s, %s, %s, %s)
-                        ON CONFLICT (query_crc32) DO NOTHING
-                    """, (query_crc32, query, model_name, query_embedding.tolist()))
-                conn.commit()
-                logging.info(f"Сохранен новый эмбеддинг запроса: {query}")
-            except Exception as e:
-                logging.error(f"Ошибка при сохранении эмбеддинга запроса: {e}")
-                conn.rollback()
+            if query_embedding is None:
+                query_embedding = encode_string(query)
+                insert_embedding(conn, query_cache_table, -1, -1, query_crc32, query_embedding)
 
-        # Используем косинусное расстояние для поиска
-        try:
-            with conn.cursor() as cur:
-                cur.execute(f"""
-                    WITH query_embedding AS (
-                        SELECT embedding
-                        FROM {query_cache_table}
-                        WHERE query_crc32 = %s
-                    )
-                    SELECT m.movie_id, 1 - (m.embedding <=> (SELECT embedding FROM query_embedding)) as similarity
-                    FROM {embeddings_table} m, query_embedding
-                    ORDER BY similarity DESC
-                    LIMIT %s
-                """, (query_crc32, top_k))
-
-                results = cur.fetchall()
-            logging.info(f"Найдено {len(results)} результатов поиска.")
-        except Exception as e:
-            logging.error(f"Ошибка при выполнении поискового запроса: {e}")
-            results = []
-
-        results_html = "<ol>"
-        for movie_id, similarity in results:
-            # Находим название фильма по ID
-            movie_title = None
-            for movie in movies_data:
-                if movie['id'] == movie_id:
-                    movie_title = movie['name']
-                    break
+            movie_embeddings = get_movie_embeddings(conn)
 
-            if movie_title:
-                results_html += f"<li><strong>{movie_title}</strong> (Сходство: {similarity:.4f})</li>"
-        results_html += "</ol>"
+            # Вычисляем косинусное сходство
+            similarities = util.cos_sim(query_embedding, list(movie_embeddings.values()))[0]
 
-        search_time = time.time() - start_time
-        logging.info(f"Поиск выполнен за {search_time:.2f} секунд.")
-
-        return f"<p>Время поиска: {search_time:.2f} сек</p>{results_html}"
+            # Сортируем результаты
+            top_results = sorted(zip(similarities, movie_embeddings.keys()), key=lambda x: x[0], reverse=True)[:top_k]
 
+            results = []
+            for score, movie_name in top_results:
+                movie = next((m for m in movies_data if m['name'] == movie_name), None)
+                if movie:
+                    results.append({
+                        "name": movie['name'],
+                        "year": movie['year'],
+                        "genres": movie['genresList'],
+                        "description": movie['description'],
+                        "score": float(score)
+                    })
+
+        return results
     except Exception as e:
-        logging.error(f"Ошибка при выполнении поиска: {e}")
-        return "<p>Произошла ошибка при выполнении поиска.</p>"
-
+        logging.error(f"Ошибка при поиске фильмов: {e}")
+        return []
     finally:
-        if conn:
-            conn.close()
         search_in_progress = False
 
 # Запускаем обработку фильмов в отдельном потоке
-processing_thread = threading.Thread(target=process_movies)
-processing_thread.start()
+threading.Thread(target=process_movies, daemon=True).start()
 
 # Создаем интерфейс Gradio
+def gradio_search(query: str) -> str:
+    results = search_movies(query)
+    output = ""
+    for movie in results:
+        output += f"Название: {movie['name']} ({movie['year']})\n"
+        output += f"Жанры: {', '.join(movie['genres'])}\n"
+        output += f"Описание: {movie['description']}\n"
+        output += f"Релевантность: {movie['score']:.2f}\n\n"
+    return output
+
 iface = gr.Interface(
-    fn=search_movies,
-    inputs=gr.Textbox(lines=2, placeholder="Введите запрос для поиска фильмов..."),
-    outputs=gr.HTML(label="Результаты поиска"),
-    title="Семантический поиск фильмов",
-    description="Введите описание фильма, который вы ищете, и система найдет наиболее похожие фильмы."
+    fn=gradio_search,
+    inputs="text",
+    outputs="text",
+    title="Поиск фильмов",
+    description="Введите запрос для поиска фильмов"
 )
 
-# Запускаем интерфейс
 iface.launch()