Update app.py

app.py

@@ -13,6 +13,7 @@ import logging
 from sklearn.preprocessing import normalize
 from concurrent.futures import ThreadPoolExecutor
 import requests
+import voyageai
 
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -43,6 +44,9 @@ VOYAGE_API_KEY = os.environ.get("VOYAGE_API_KEY")
 if VOYAGE_API_KEY is None:
     raise ValueError("VOYAGE_API_KEY environment variable not set.")
 
+# Инициализация клиента Voyage AI
+vo = voyageai.Client(api_key=VOYAGE_API_KEY)
+
 # Имена таблиц
 embeddings_table = "movie_embeddings"
 query_cache_table = "query_cache"
@@ -79,7 +83,19 @@ batch_size = 32
 num_threads = 5
 
 # Количество потоков для параллельного реранкинга
-rerank_threads = 5 # Подберите оптимальное значение
+rerank_threads = 3 # Ограничено лимитом RPM
+
+# Лимиты Voyage AI (запросов в минуту, токенов в минуту) - БЕСПЛАТНЫЙ АККАУНТ
+RPM_LIMIT = 3
+TPM_LIMIT = 10000
+
+# Переменные для отслеживания текущего использования
+current_rpm = 0
+current_tpm = 0
+last_reset_time = time.time()
+
+# Среднее количество токенов на описание фильма (можно вычислить один раз при запуске)
+avg_tokens_per_movie = 150  # Замените на более точное значение, если оно известно
 
 def get_db_connection():
     """Устанавливает соединение с базой данных."""
@@ -298,8 +314,64 @@ def get_movie_embeddings(conn):
         logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
     return movie_embeddings
 
+def check_and_wait_for_limits():
+    """Проверяет лимиты RPM и TPM и ожидает, если они исчерпаны."""
+    global current_rpm, current_tpm, last_reset_time
+    
+    elapsed_time = time.time() - last_reset_time
+    
+    if elapsed_time >= 60:
+        current_rpm = 0
+        current_tpm = 0
+        last_reset_time = time.time()
+        logging.info("Лимиты RPM и TPM сброшены.")
+    
+    if current_rpm >= RPM_LIMIT or current_tpm >= TPM_LIMIT:
+        wait_time = 60 - elapsed_time
+        logging.warning(f"Превышены лимиты RPM ({current_rpm}/{RPM_LIMIT}) или TPM ({current_tpm}/{TPM_LIMIT}). Ожидание {wait_time:.2f} секунд...")
+        time.sleep(max(0, wait_time))
+        current_rpm = 0
+        current_tpm = 0
+        last_reset_time = time.time()
+        logging.info("Лимиты RPM и TPM сброшены после ожидания.")
+
+def create_optimized_batches(query, results, max_tokens_per_batch=TPM_LIMIT):
+    """Создает батчи для реранкинга, оптимизированные по количеству токенов."""
+    global avg_tokens_per_movie
+
+    batches = []
+    current_batch = []
+    current_batch_tokens = 0
+
+    query_tokens = vo.count_tokens([query], model="rerank-2")
+
+    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']}"
+            
+            # Считаем токены, но не отправляем запрос если лимит уже исчерпан
+            estimated_movie_tokens = avg_tokens_per_movie 
+            
+            if (current_batch_tokens + query_tokens + estimated_movie_tokens) <= max_tokens_per_batch:
+                current_batch.append((movie_id, _))
+                current_batch_tokens += estimated_movie_tokens
+            else:
+                batches.append(current_batch)
+                current_batch = [(movie_id, _)]
+                current_batch_tokens = estimated_movie_tokens
+
+    if current_batch:
+        batches.append(current_batch)
+
+    return batches
+
 def rerank_batch_voyage(query, batch):
     """Переранжирует пакет результатов с помощью Voyage AI."""
+    global current_rpm, current_tpm
+    
+    check_and_wait_for_limits()
+
     url = "https://api.voyageai.com/v1/rerank"
     headers = {
         "Authorization": f"Bearer {VOYAGE_API_KEY}",
@@ -315,6 +387,8 @@ def rerank_batch_voyage(query, batch):
             documents.append(movie_info)
             movie_ids.append(movie_id)
 
+    
+    
     payload = {
         "query": query,
         "documents": documents,
@@ -324,6 +398,13 @@ def rerank_batch_voyage(query, batch):
     }
 
     try:
+        batch_tokens = vo.count_tokens([query] + documents, model="rerank-2")
+        
+        current_rpm += 1
+        current_tpm += batch_tokens
+
+        logging.info(f"Отправка запроса к Voyage AI. RPM: {current_rpm}/{RPM_LIMIT}, TPM: {current_tpm}/{TPM_LIMIT}, Токенов в запросе: {batch_tokens}")
+        
         response = requests.post(url, headers=headers, json=payload)
         response.raise_for_status()  # Проверка на ошибки HTTP
         response_json = response.json()
@@ -337,6 +418,12 @@ def rerank_batch_voyage(query, batch):
     
     except requests.exceptions.RequestException as e:
         logging.error(f"Ошибка запроса к Voyage AI: {e}")
+        
+        if response.status_code == 429: # Too Many Requests
+            logging.warning("Слишком много запросов к Voyage AI. Ожидание сброса лимитов...")
+            check_and_wait_for_limits()
+            return rerank_batch_voyage(query, batch) # Повторная попытка после ожидания
+        
         return []
     except KeyError as e:
         logging.error(f"Ошибка обработки ответа от Voyage AI: {e}. Полный ответ: {response_json}")
@@ -345,26 +432,20 @@ def rerank_batch_voyage(query, batch):
 def rerank_results(query, results):
     """Переранжирует результаты поиска с помощью Voyage AI."""
     logging.info(f"Начало переранжирования для запроса: '{query}'")
+    
+    # Создаем оптимизированные батчи
+    batches = create_optimized_batches(query, results)
+    
     reranked_results = []
 
     with ThreadPoolExecutor(max_workers=rerank_threads) as executor:
         futures = []
-        batch = []
         batch_num = 0
-        for i, result in enumerate(results):
-            batch.append(result)
-            if len(batch) >= batch_size: # Отправляем на реранк батчами
-                logging.info(f"Отправка на переранжирование батча {batch_num+1} ({len(batch)} фильмов)")
-                future = executor.submit(rerank_batch_voyage, query, batch)
-                futures.append(future)
-                batch = []
-                batch_num += 1
-
-        # Обработка остатка
-        if batch:
+        for batch in batches:
             logging.info(f"Отправка на переранжирование батча {batch_num+1} ({len(batch)} фильмов)")
             future = executor.submit(rerank_batch_voyage, query, batch)
             futures.append(future)
+            batch_num += 1
 
         # Сбор результатов
         for i, future in enumerate(futures):
@@ -422,7 +503,7 @@ def search_movies(query, top_k=20):
                     FROM {embeddings_table} m, query_embedding
                     ORDER BY similarity DESC
                     LIMIT %s
-                """, (query_crc32, int(top_k * 1.1)))  # Уменьшаем лимит до * 1.1
+                """, (query_crc32, int(top_k * 2)))  # Увеличиваем лимит до * 2
 
                 results = cur.fetchall()
             logging.info(f"Найдено {len(results)} предварительных результатов поиска.")