Update app.py

app.py

@@ -13,7 +13,6 @@ 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')
@@ -39,13 +38,10 @@ logging.info(f"Загрузка модели {model_name}...")
 model = SentenceTransformer(model_name)
 logging.info("Модель загружена успешно.")
 
-# Voyage AI API Key
-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)
+# Jina AI 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.")
 
 # Имена таблиц
 embeddings_table = "movie_embeddings"
@@ -82,20 +78,8 @@ batch_size = 32
 # Количество потоков для параллельной обработки
 num_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  # Замените на более точное значение, если оно известно
+# Количество потоков для параллельного реранкинга (Jina AI не имеет жестких ограничений)
+rerank_threads = 5
 
 def get_db_connection():
     """Устанавливает соединение с базой данных."""
@@ -314,147 +298,55 @@ 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"
+def rerank_with_jina(query, results, top_n):
+    """Переранжирует результаты с помощью Jina AI."""
+    url = 'https://api.jina.ai/v1/rerank'
     headers = {
-        "Authorization": f"Bearer {VOYAGE_API_KEY}",
-        "content-type": "application/json"
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {JINA_API_KEY}'
     }
-    
+
     documents = []
-    movie_ids = []
-    for movie_id, _ in batch:
+    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_ids.append(movie_id)
+            documents.append({"text": movie_info})
 
-    
-    
-    payload = {
+    data = {
+        "model": "jina-reranker-v2-base-multilingual",
         "query": query,
-        "documents": documents,
-        "model": "rerank-2", # Можно использовать rerank-2-lite для более быстрой, но менее точной модели
-        "return_documents": False,
-        "truncation": True
+        "top_n": top_n,
+        "documents": documents
     }
-
+    
+    logging.info(f"Отправка запроса к Jina AI для переранжирования {len(documents)} документов...")
+    
     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 = requests.post(url, headers=headers, json=data)
+        response.raise_for_status()
         response_json = response.json()
         
         reranked_results = []
-        for item in response_json['data']:
-            reranked_results.append((movie_ids[item['index']], item['relevance_score']))
-        
-        logging.info(f"Voyage AI: Успешно переранжирован батч. Задействовано токенов: {response_json['usage']['total_tokens']}")
+        for result in response_json['results']:
+            reranked_results.append((results[result['index']][0], result['relevance_score']))
+
+        logging.info(f"Jina AI: Успешно переранжировано. Задействовано токенов: {response_json['usage']['total_tokens']}")
         return reranked_results
     
     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) # Повторная попытка после ожидания
-        
+        logging.error(f"Ошибка запроса к Jina AI: {e}")
         return []
     except KeyError as e:
-        logging.error(f"Ошибка обработки ответа от Voyage AI: {e}. Полный ответ: {response_json}")
+        logging.error(f"Ошибка обработки ответа от Jina AI: {e}. Полный ответ: {response_json}")
         return []
 
-def rerank_results(query, results):
-    """Переранжирует результаты поиска с помощью Voyage AI."""
+def rerank_results(query, results, top_k):
+    """Переранжирует результаты поиска с помощью Jina AI."""
     logging.info(f"Начало переранжирования для запроса: '{query}'")
-    
-    # Создаем оптимизированные батчи
-    batches = create_optimized_batches(query, results)
-    
-    reranked_results = []
-
-    with ThreadPoolExecutor(max_workers=rerank_threads) as executor:
-        futures = []
-        batch_num = 0
-        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):
-            try:
-                batch_result = future.result()
-                reranked_results.extend(batch_result)
-                logging.info(f"Завершен реранк батча {i+1}")
-            except Exception as e:
-                logging.error(f"Ошибка при переранжировании батча {i+1}: {e}")
+
+    # Jina AI не имее�� жестких ограничений, поэтому можем обрабатывать все результаты за раз
+    reranked_results = rerank_with_jina(query, results, top_k)
 
     reranked_results = sorted(reranked_results, key=lambda x: x[1], reverse=True)
     logging.info("Переранжирование завершено.")
@@ -512,7 +404,7 @@ def search_movies(query, top_k=20):
             results = []
 
         # Переранжируем результаты
-        reranked_results = rerank_results(query, results)
+        reranked_results = rerank_results(query, results, top_k)
 
         output = ""
         for movie_id, score in reranked_results[:top_k]:
@@ -522,7 +414,7 @@ def search_movies(query, top_k=20):
                 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"
-                output += f"<p><strong>Релевантность (Voyage AI reranker score):</strong> {score:.4f}</p>\n"
+                output += f"<p><strong>Релевантность (Jina AI reranker score):</strong> {score:.4f}</p>\n"
                 output += "<hr>\n"
 
         search_time = time.time() - start_time