Update app.py

app.py

@@ -40,29 +40,19 @@ logging.info("Модель загружена успешно.")
 
 # Jina AI Reranker API
 JINA_API_URL = 'https://api.jina.ai/v1/rerank'
-JINA_API_KEY = os.environ.get("JINA_API_KEY")  # Получение ключа из переменной окружения
+JINA_API_KEY = os.environ.get("JINA_API_KEY")
 if JINA_API_KEY is None:
     raise ValueError("JINA_API_KEY environment variable not set.")
 JINA_RERANKER_MODEL = "jina-reranker-v2-base-multilingual"
 
-
 # Имена таблиц
 embeddings_table = "movie_embeddings"
 query_cache_table = "query_cache"
+movies_table = "Movies"  # Имя таблицы с фильмами
 
 # Максимальный размер таблицы кэша запросов в байтах (50MB)
 MAX_CACHE_SIZE = 50 * 1024 * 1024
 
-# Загружаем данные из файла movies.json
-try:
-    import json
-    with open("movies.json", "r", encoding="utf-8") as f:
-        movies_data = json.load(f)
-    logging.info(f"Загружено {len(movies_data)} фильмов из movies.json")
-except FileNotFoundError:
-    logging.error("Ошибка: Файл movies.json не найден.")
-    movies_data = []
-
 # Очередь для необработанных фильмов
 movies_queue = queue.Queue()
 
@@ -101,32 +91,29 @@ def setup_database():
             # Создаем расширение 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} (
+                CREATE TABLE IF NOT EXISTS {embeddings_table} (
                     movie_id INTEGER PRIMARY KEY,
                     embedding_crc32 BIGINT,
                     string_crc32 BIGINT,
                     model_name TEXT,
                     embedding vector(1024)
                 );
-                CREATE INDEX ON {embeddings_table} (string_crc32);
+                CREATE INDEX IF NOT EXISTS idx_string_crc32 ON {embeddings_table} (string_crc32);
             """)
 
             # Создаем таблицу для кэширования запросов
             cur.execute(f"""
-                CREATE TABLE {query_cache_table} (
+                CREATE TABLE IF NOT EXISTS {query_cache_table} (
                     query_crc32 BIGINT PRIMARY KEY,
                     query TEXT,
                     model_name TEXT,
                     embedding vector(1024),
                     created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
                 );
-                CREATE INDEX ON {query_cache_table} (query_crc32);
-                CREATE INDEX ON {query_cache_table} (created_at);
+                CREATE INDEX IF NOT EXISTS idx_query_crc32 ON {query_cache_table} (query_crc32);
+                CREATE INDEX IF NOT EXISTS idx_created_at ON {query_cache_table} (created_at);
             """)
 
         conn.commit()
@@ -158,14 +145,26 @@ def get_movies_without_embeddings():
     movies_to_process = []
     try:
         with conn.cursor() as cur:
-            # Получаем список ID фильмов, которые уже есть в базе
+            # Получаем список 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)
+            # Получаем список всех фильмов из таблицы Movies с подготовленной строкой
+            cur.execute(f"""
+                SELECT id, data,
+                    'Название: ' || data->>'name' ||
+                    '\\nГод: ' || data->>'year' ||
+                    '\\nЖанры: ' || (SELECT string_agg(genre->>'name', ', ') FROM jsonb_array_elements(data->'genres') AS genre) ||
+                    '\\nОписание: ' || COALESCE(data->>'description', '')
+                    AS prepared_string
+                FROM {movies_table}
+            """)
+            all_movies = cur.fetchall()
+
+            # Фильтруем только те фильмы, которых нет в таблице эмбеддингов
+            for movie_id, movie_data, prepared_string in all_movies:
+                if movie_id not in existing_ids:
+                    movies_to_process.append((movie_id, movie_data, prepared_string))
 
         logging.info(f"Найдено {len(movies_to_process)} фильмов для об��аботки.")
     except Exception as e:
@@ -215,23 +214,22 @@ def process_batch(batch):
         return
 
     try:
-        for movie in batch:
-            embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genreslist']}\nОписание: {movie['description']}"
-            string_crc32 = calculate_crc32(embedding_string)
+        for movie_id, movie_data, prepared_string in batch:
+            string_crc32 = calculate_crc32(prepared_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 = encode_string(prepared_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']}'")
+                if insert_embedding(conn, embeddings_table, movie_id, embedding_crc32, string_crc32, embedding):
+                    logging.info(f"Сохранен эмбеддинг для '{movie_data['name']}' (ID: {movie_id})")
                 else:
-                    logging.error(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
+                    logging.error(f"Ошибка сохранения эмбеддинга для '{movie_data['name']}' (ID: {movie_id})")
             else:
-                logging.info(f"Эмбеддинг для '{movie['name']}' уже существует")
+                logging.info(f"Эмбеддинг для '{movie_data['name']}' (ID: {movie_id}) уже существует")
     except Exception as e:
         logging.error(f"Ошибка при обработке пакета фильмов: {e}")
     finally:
@@ -281,32 +279,44 @@ def process_movies():
     processing_complete = True
     logging.info("Обработка фильмов завершена")
 
-def get_movie_embeddings(conn):
-    """Загружает все эмбеддинги фильмов из базы данных."""
-    movie_embeddings = {}
+def get_movie_data_from_db(conn, movie_ids):
+    """Получает данные фильмов из таблицы Movies по списку ID."""
+    movie_data_dict = {}
     try:
         with conn.cursor() as cur:
-            cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
-            for movie_id, embedding in cur.fetchall():
-                # Находим название фильма по ID
-                for movie in movies_data:
-                    if movie['id'] == movie_id:
-                        movie_embeddings[movie['name']] = normalize(np.array(embedding).reshape(1, -1))[0]
-                        break
-        logging.info(f"Загружено {len(movie_embeddings)} эмбеддингов фильмов.")
+            cur.execute(f"""
+                SELECT id, data,
+                    'Название: ' || data->>'name' ||
+                    '\\nГод: ' || data->>'year' ||
+                    '\\nЖанры: ' || (SELECT string_agg(genre->>'name', ', ') FROM jsonb_array_elements(data->'genres') AS genre) ||
+                    '\\nОписание: ' || COALESCE(data->>'description', '')
+                    AS prepared_string
+                FROM {movies_table}
+                WHERE id IN %s
+            """, (tuple(movie_ids),))
+            for movie_id, movie_data, prepared_string in cur.fetchall():
+                movie_data_dict[movie_id] = (movie_data, prepared_string)
     except Exception as e:
-        logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
-    return movie_embeddings
+        logging.error(f"Ошибка при получении данных фильмов из БД: {e}")
+    return movie_data_dict
 
 def rerank_with_api(query, results, top_k):
     """Переранжирует результаты с помощью Jina AI Reranker API."""
     logging.info(f"Начало переранжирования для запроса: '{query}'")
+    
+    # Получаем данные фильмов из БД
+    conn = get_db_connection()
+    movie_ids = [movie_id for movie_id, _ in results]
+    movie_data_dict = get_movie_data_from_db(conn, movie_ids)
+    conn.close()
+
     documents = []
     for movie_id, _ in results:
-        movie = next((m for m in movies_data if m['id'] == movie_id), None)
-        if movie:
-            movie_info = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genreslist']}\nОписание: {movie['description']}"
-            documents.append(movie_info)
+        movie_data, prepared_string = movie_data_dict.get(movie_id, (None, None))
+        if movie_data:
+            documents.append(prepared_string)
+        else:
+            logging.warning(f"Данные для фильма с ID {movie_id} не найдены в БД.")
 
     headers = {
         'Content-Type': 'application/json',
@@ -318,18 +328,17 @@ def rerank_with_api(query, results, top_k):
         "top_n": top_k,
         "documents": documents
     }
-    logging.info(f"Отправка данных на реранжировку (len): {len(data)}")
+    logging.info(f"Отправка данных на реранжировку (documents count): {len(data['documents'])}")
 
     try:
         response = requests.post(JINA_API_URL, headers=headers, json=data)
-        response.raise_for_status()  # Проверка на ошибки HTTP
+        response.raise_for_status()
         result = response.json()
-        logging.info(f"Ответ от API реранжировщика (len): {len(result)}")
+        logging.info(f"Ответ от API реранжировщика получен.")
 
         reranked_results = []
         if 'results' in result:
             for item in result['results']:
-              
                 index = item['index']
                 movie_id = results[index][0]
                 reranked_results.append((movie_id, item['relevance_score']))
@@ -337,7 +346,6 @@ def rerank_with_api(query, results, top_k):
             logging.warning("Ответ от API не содержит ключа 'results'.")
 
         logging.info("Переранжирование завершено.")
-        # time.sleep(0.1)
         return reranked_results
 
     except requests.exceptions.RequestException as e:
@@ -394,20 +402,29 @@ def search_movies(query, top_k=25):
         except Exception as e:
             logging.error(f"Ошибка при выполнении поискового запроса: {e}")
             results = []
+        finally:
+            conn.close()
 
         # Переранжируем результаты с помощью API
         reranked_results = rerank_with_api(query, results, top_k)
 
+        conn = get_db_connection()
+        movie_ids = [movie_id for movie_id, _ in reranked_results]
+        movie_data_dict = get_movie_data_from_db(conn, movie_ids)
+        conn.close()
+
         output = ""
         for movie_id, score in reranked_results:
-            # Находим фильм по ID
-            movie = next((m for m in movies_data if m['id'] == movie_id), None)
-            if movie:
-                output += f"<h3>{movie['name']} ({movie['year']})</h3>\n"
-                output += f"<p><strong>Жанры:</strong> {movie['genreslist']}</p>\n"
-                output += f"<p><strong>Описание:</strong> {movie['description']}</p>\n"
+            # Находим данные фильма
+            movie_data, _ = movie_data_dict.get(movie_id, (None, None))
+            if movie_data:
+                output += f"<h3>{movie_data['name']} ({movie_data['year']})</h3>\n"
+                output += f"<p><strong>Жанры:</strong> {', '.join([genre['name'] for genre in movie_data['genres']])}</p>\n"
+                output += f"<p><strong>Описание:</strong> {movie_data.get('description', '')}</p>\n"
                 output += f"<p><strong>Релевантность (reranker score):</strong> {score:.4f}</p>\n"
                 output += "<hr>\n"
+            else:
+                logging.warning(f"Данные для фильма с ID {movie_id} не найдены в БД.")
 
         search_time = time.time() - start_time
         logging.info(f"Поиск выполнен за {search_time:.2f} секунд.")
@@ -419,8 +436,6 @@ def search_movies(query, top_k=25):
         return "<p>Произошла ошибка при выполнении поиска.</p>"
 
     finally:
-        if conn:
-            conn.close()
         search_in_progress = False
 
 # Запускаем обработку фильмов в отдельном потоке