Update app.py

app.py

@@ -10,7 +10,7 @@ import zlib
 from urllib.parse import urlparse
 
 # Настройки базы данных PostgreSQL
-DATABASE_URL = os.environ.get("DB_URL")
+DATABASE_URL = os.environ.get("DATABASE_URL")
 if DATABASE_URL is None:
     raise ValueError("DATABASE_URL environment variable not set.")
 
@@ -69,13 +69,17 @@ def get_db_connection():
         print(f"Ошибка подключения к базе данных: {e}")
         return None
 
-def create_embeddings_table():
-    """Создает таблицу для хранения эмбеддингов фильмов, если она не существует."""
+def setup_database():
+    """Настраивает базу данных: создает расширение, таблицы и триггер."""
     conn = get_db_connection()
     if conn is None:
         return
 
     with conn.cursor() as cur:
+        # Создаем расширение pgvector
+        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
+
+        # Создаем таблицу для хранения эмбеддингов фильмов
         cur.execute(f"""
             CREATE TABLE IF NOT EXISTS {embeddings_table} (
                 movie_id INTEGER,
@@ -85,16 +89,8 @@ def create_embeddings_table():
                 embedding vector(1024)
             );
         """)
-        conn.commit()
-    conn.close()
-
-def create_query_cache_table():
-    """Создает таблицу для кэширования эмбеддингов запросов, если она не существует."""
-    conn = get_db_connection()
-    if conn is None:
-        return
 
-    with conn.cursor() as cur:
+        # Создаем таблицу для кэширования эмбеддингов запросов
         cur.execute(f"""
             CREATE TABLE IF NOT EXISTS {query_cache_table} (
                 query_crc32 BIGINT PRIMARY KEY,
@@ -106,49 +102,41 @@ def create_query_cache_table():
             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);
         """)
+
+        # Создаем функцию и триггер для автоматического удаления старых записей из таблицы кэша запросов
+        cur.execute(f"""
+            CREATE OR REPLACE FUNCTION manage_query_cache_size()
+            RETURNS TRIGGER AS $$
+            DECLARE
+                table_size BIGINT;
+                row_to_delete RECORD;
+            BEGIN
+                SELECT pg_total_relation_size('{query_cache_table}') INTO table_size;
+                IF table_size > {MAX_CACHE_SIZE} THEN
+                    FOR row_to_delete IN
+                        SELECT query_crc32
+                        FROM {query_cache_table}
+                        ORDER BY created_at ASC
+                    LOOP
+                        DELETE FROM {query_cache_table} WHERE query_crc32 = row_to_delete.query_crc32;
+                        SELECT pg_total_relation_size('{query_cache_table}') INTO table_size;
+                        EXIT WHEN table_size <= {MAX_CACHE_SIZE};
+                    END LOOP;
+                END IF;
+                RETURN NEW;
+            END;
+            $$ LANGUAGE plpgsql;
+
+            CREATE OR REPLACE TRIGGER trg_manage_query_cache_size
+            AFTER INSERT ON {query_cache_table}
+            FOR EACH ROW
+            EXECUTE PROCEDURE manage_query_cache_size();
+        """)
         conn.commit()
     conn.close()
 
-def create_trigger_function():
-    """Создает функцию и триггер для автоматического удаления старых записей из таблицы кэша запросов"""
-    conn = get_db_connection()
-    if conn:
-        with conn.cursor() as cur:
-            cur.execute(f"""
-                CREATE OR REPLACE FUNCTION manage_query_cache_size()
-                RETURNS TRIGGER AS $$
-                DECLARE
-                    table_size BIGINT;
-                    row_to_delete RECORD;
-                BEGIN
-                    SELECT pg_total_relation_size('{query_cache_table}') INTO table_size;
-                    IF table_size > {MAX_CACHE_SIZE} THEN
-                        FOR row_to_delete IN
-                            SELECT query_crc32
-                            FROM {query_cache_table}
-                            ORDER BY created_at ASC
-                        LOOP
-                            DELETE FROM {query_cache_table} WHERE query_crc32 = row_to_delete.query_crc32;
-                            SELECT pg_total_relation_size('{query_cache_table}') INTO table_size;
-                            EXIT WHEN table_size <= {MAX_CACHE_SIZE};
-                        END LOOP;
-                    END IF;
-                    RETURN NEW;
-                END;
-                $$ LANGUAGE plpgsql;
-
-                CREATE OR REPLACE TRIGGER trg_manage_query_cache_size
-                AFTER INSERT ON {query_cache_table}
-                FOR EACH ROW
-                EXECUTE PROCEDURE manage_query_cache_size();
-            """)
-            conn.commit()
-        conn.close()
-
-# Создаем таблицы, индексы и триггер при запуске приложения
-create_embeddings_table()
-create_query_cache_table()
-create_trigger_function()
+# Настраиваем базу данных при запуске приложения
+setup_database()
 
 def calculate_crc32(text):
     """Вычисляет CRC32 для строки."""
@@ -158,64 +146,39 @@ def encode_string(text):
     """Кодирует строку в эмбеддинг."""
     return model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
 
-def insert_embedding(conn, movie_id, embedding_string, model_name, embedding):
-    """Вставляет эмбеддинг фильма в базу данных."""
-    embedding_crc32 = calculate_crc32(str(embedding.tolist()))
-    string_crc32 = calculate_crc32(embedding_string)
+def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_name):
+    """
+    Пытается получить эмбеддинг из указанной таблицы по CRC32.
+    Возвращает эмбеддинг, если найден, иначе None.
+    """
     with conn.cursor() as cur:
-        try:
-            cur.execute(
-                f"""
-                INSERT INTO {embeddings_table} (movie_id, embedding_crc32, string_crc32, model_name, embedding)
-                VALUES (%s, %s, %s, %s, %s)
-                ON CONFLICT (embedding_crc32) DO NOTHING;
-                """,
-                (movie_id, embedding_crc32, string_crc32, model_name, embedding.tolist())
-            )
-            conn.commit()
-            return True
-        except Exception as e:
-            print(f"Ошибка при вставке эмбеддинга фильма: {e}")
-            conn.rollback()
-            return False
+        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:
+            return torch.tensor(result[0])
+        else:
+            return None
+
+def insert_embedding(conn, table_name, crc32_column, crc32_value, other_columns, embedding):
+    """Вставляет эмбеддинг в указанную таблицу."""
+    columns = ', '.join([crc32_column] + list(other_columns.keys()) + ['model_name', 'embedding'])
+    placeholders = ', '.join(['%s'] * (len(other_columns) + 3))
+    values = (crc32_value,) + tuple(other_columns.values()) + (model_name, embedding.tolist())
 
-def insert_query_embedding(conn, query, model_name, embedding):
-    """Вставляет эмбеддинг запроса в таблицу кэша."""
-    query_crc32 = calculate_crc32(query)
     with conn.cursor() as cur:
         try:
-            cur.execute(
-                f"""
-                INSERT INTO {query_cache_table} (query_crc32, query, model_name, embedding)
-                VALUES (%s, %s, %s, %s)
-                ON CONFLICT (query_crc32) DO UPDATE SET created_at = DEFAULT; 
-                """,
-                (query_crc32, query, model_name, embedding.tolist())
-            )
+            cur.execute(f"""
+                INSERT INTO {table_name} ({columns})
+                VALUES ({placeholders})
+                ON CONFLICT ({crc32_column}) DO NOTHING;
+                """, values)
             conn.commit()
-            print(f"Эмбеддинг для запроса '{query}' сохранен в кэше.")
             return True
         except Exception as e:
-            print(f"Ошибка при вставке эмбеддинга запроса: {e}")
+            print(f"Ошибка при вставке эмбеддинга в таблицу {table_name}: {e}")
             conn.rollback()
             return False
 
-def get_movie_embeddings(conn):
-    """Загружает все эмбеддинги фильмов из базы данных."""
-    movie_embeddings = {}
-    with conn.cursor() as cur:
-        cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
-        rows = cur.fetchall()
-        for row in rows:
-            movie_id, embedding = row
-            # Находим название фильма по его ID
-            for movie in movies_data:
-                if movie['id'] == movie_id:
-                    title = movie["name"]
-                    movie_embeddings[title] = torch.tensor(embedding)
-                    break
-    return movie_embeddings
-
 def process_movies():
     """
     Обрабатывает фильмы из очереди, создавая для них эмбеддинги и сохраняя их в базу данных.
@@ -251,32 +214,45 @@ def process_movies():
         ]
 
         print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
-        embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True)
-
+        
         with db_lock:
-            for movie, embedding, embedding_string in zip(batch, embeddings, embedding_strings):
-                if insert_embedding(conn, movie['id'], embedding_string, model_name, embedding):
-                    print(f"Эмбеддинг для фильма '{movie['name']}' сохранен в базе данных.")
-                else:
-                    print(f"Ошибка сохранения эмбеддинга для фильма '{movie['name']}'.")
+          for movie, embedding_string in zip(batch, embedding_strings):
+              movie_id = movie['id']
+              embedding_crc32 = calculate_crc32(str(embedding.tolist()))
+              string_crc32 = calculate_crc32(embedding_string)
+              
+              # Проверяем, есть ли уже эмбеддинг для этого фильма в базе данных
+              existing_embedding = get_embedding_from_db(conn, embeddings_table, "embedding_crc32", embedding_crc32, model_name)
+              
+              if existing_embedding is None:
+                  # Создаем эмбеддинг, только если его нет в базе данных
+                  embedding = encode_string(embedding_string)
+                  
+                  if insert_embedding(conn, embeddings_table, "embedding_crc32", embedding_crc32, {"movie_id": movie_id, "string_crc32": string_crc32}, embedding):
+                      print(f"Эмбеддинг для фильма '{movie['name']}' сохранен в базе данных.")
+                  else:
+                      print(f"Ошибка сохранения эмбеддинга для фильма '{movie['name']}'.")
+              else:
+                  print(f"Эмбеддинг для фильма '{movie['name']}' уже существует в базе данных.")
 
     conn.close()
     print("Обработка фильмов завершена.")
 
-def get_query_embedding_from_db(conn, query):
-    """
-    Пытается получить эмбеддинг запроса из базы данных по CRC32.
-    Возвращает эмбеддинг, если найден, иначе None.
-    """
-    query_crc32 = calculate_crc32(query)
+def get_movie_embeddings(conn):
+    """Загружает все эмбеддинги фильмов из базы данных."""
+    movie_embeddings = {}
     with conn.cursor() as cur:
-        cur.execute(f"SELECT embedding FROM {query_cache_table} WHERE query_crc32 = %s AND model_name = %s", (query_crc32, model_name))
-        result = cur.fetchone()
-        if result:
-            print(f"Эмбеддинг для запроса '{query}' найден в кэше.")
-            return torch.tensor(result[0])
-        else:
-            return None
+        cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
+        rows = cur.fetchall()
+        for row in rows:
+            movie_id, embedding = row
+            # Находим название фильма по его ID
+            for movie in movies_data:
+                if movie['id'] == movie_id:
+                    title = movie["name"]
+                    movie_embeddings[title] = torch.tensor(embedding)
+                    break
+    return movie_embeddings
 
 def search_movies(query, top_k=10):
     """
@@ -299,14 +275,20 @@ def search_movies(query, top_k=10):
         search_in_progress = False
         return "<p>Ошибка подключения к базе данных.</p>"
 
-    print(f"Начало создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
-    query_embedding_tensor = get_query_embedding_from_db(conn, query)
+    query_crc32 = calculate_crc32(query)
+
+    # Проверяем, есть ли уже эмбеддинг для этого запроса в кэше
+    print(f"Начало поиска эмбеддинга запроса в кэше: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    query_embedding_tensor = get_embedding_from_db(conn, query_cache_table, "query_crc32", query_crc32, model_name)
+    print(f"Окончание поиска эмбеддинга запроса в кэше: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
     if query_embedding_tensor is None:
+        print(f"Начало создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
         query_embedding_tensor = encode_string(query)
+        print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+
         # Вставляем эмбеддинг запроса в базу данных
-        insert_query_embedding(conn, query, model_name, query_embedding_tensor)
-    print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+        insert_embedding(conn, query_cache_table, "query_crc32", query_crc32, {"query": query}, query_embedding_tensor)
 
     with db_lock:
         current_movie_embeddings = get_movie_embeddings(conn)