Update app.py

app.py

@@ -13,7 +13,7 @@ from sklearn.preprocessing import normalize
 from concurrent.futures import ThreadPoolExecutor
 import requests
 from fastapi import FastAPI, HTTPException, Query
-from typing import List
+from typing import List, Optional
 import uvicorn
 from starlette.requests import Request
 from starlette.responses import HTMLResponse, JSONResponse
@@ -46,7 +46,7 @@ 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"
@@ -173,14 +173,14 @@ def get_movie_data_from_db(conn, movie_ids):
         logging.error(f"Ошибка при получении данных фильмов из БД: {e}")
     return movie_data_dict
 
-def get_jina_ai_balance():
+def get_jina_ai_balance(api_key: str):
     """Получает остаток баланса Jina AI."""
     try:
         headers = {
             'Content-Type': 'application/json'
         }
         params = {
-            'api_key': JINA_API_KEY
+            'api_key': api_key
         }
         response = requests.get(JINA_DASHBOARD_API_URL, headers=headers, params=params)
         response.raise_for_status()
@@ -190,10 +190,15 @@ def get_jina_ai_balance():
         logging.error(f"Ошибка при запросе к API баланса Jina AI: {e}")
         return None
 
-def rerank_with_api(query, results, top_k, rerank_top_k=None):
+def rerank_with_api(query, results, top_k, rerank_top_k=None, api_key=None):
     """Переранжирует результаты с помощью Jina AI Reranker API."""
     logging.info(f"Начало переранжирования для запроса: '{query}'")
 
+    # Если rerank_top_k равен 0, не используем реранкер
+    if rerank_top_k == 0:
+        logging.info("Переранжирование отключено (rerank_top_k = 0).")
+        return results, False, 0
+
     # Получаем данные фильмов из БД
     conn = get_db_connection()
     movie_ids = [movie_id for movie_id, _ in results]
@@ -208,14 +213,16 @@ def rerank_with_api(query, results, top_k, rerank_top_k=None):
         else:
             logging.warning(f"Данные для фильма с ID {movie_id} не найдены в БД.")
 
+    reranked_count = min(rerank_top_k or top_k*2, len(documents))
+
     headers = {
         'Content-Type': 'application/json',
-        'Authorization': f'Bearer {JINA_API_KEY}'
+        'Authorization': f'Bearer {api_key or JINA_API_KEY}'
     }
     data = {
         "model": JINA_RERANKER_MODEL,
         "query": query,
-        "top_n": rerank_top_k or top_k * 2,
+        "top_n": rerank_top_k or top_k*2,
         "documents": documents
     }
     logging.info(f"Отправка данных на реранжировку (documents count): {len(data['documents'])}, top_n: {data['top_n']}")
@@ -236,13 +243,13 @@ def rerank_with_api(query, results, top_k, rerank_top_k=None):
             logging.warning("Ответ от API не содержит ключа 'results'.")
 
         logging.info("Переранжирование завершено.")
-        return reranked_results, True, data["top_n"]
+        return reranked_results, True, reranked_count
 
     except requests.exceptions.RequestException as e:
         logging.error(f"Ошибка при запросе к API реранжировщика: {e}")
-        return results, False, data["top_n"]
+        return results, False, reranked_count
 
-def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None):
+def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: Optional[int] = None, jina_api_key: Optional[str] = None):
     """Внутренняя функция для поиска фильмов по запросу (используется и в Gradio, и в API)."""
     start_time = time.time()
 
@@ -295,14 +302,18 @@ def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None
             results = []
         finally:
             conn.close()
-
-        # Переранжируем результаты с помощью API
-        reranked_results, rerank_success, reranked_count = rerank_with_api(query, results, top_k, rerank_top_k)
+        
+        # Используем реранкер только если rerank_top_k не равен 0
+        if rerank_top_k != 0:
+            reranked_results, rerank_success, reranked_count = rerank_with_api(query, results, top_k, rerank_top_k, jina_api_key)
+        else:
+            reranked_results = results
+            rerank_success = False
+            reranked_count = 0
 
         if not rerank_success:
             logging.warning("Переранжировка не удалась, ис��ользуются сырые результаты.")
             reranked_results = results[:top_k]  # Используем срез для ограничения количества результатов
-            reranked_count = 0
         else:
             reranked_results = reranked_results[:top_k]
 
@@ -350,7 +361,7 @@ def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None
         search_time = time.time() - start_time
         logging.info(f"Поиск выполнен за {search_time:.2f} секунд.")
 
-        jina_balance = get_jina_ai_balance()
+        jina_balance = get_jina_ai_balance(jina_api_key or JINA_API_KEY)
 
         return {
             "status": "success",
@@ -377,12 +388,35 @@ def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None
         }, 0
 
 @app.get("/search/", response_model=dict)
-async def api_search_movies(query: str = Query(..., description="Поисковый запрос"), 
+async def api_search_movies(query: str = Query(..., description="Поисковый запрос"),
                             top_k: int = Query(25, description="Количество возвращаемых результатов"),
-                            rerank_top_k: int = Query(None, description="Количество фильмов для передачи в реранкер (если не указано, то top_k*2)")):
-    """API endpoint для поиска фильмов."""
+                            rerank_top_k: Optional[int] = Query(None, description="Количество фильмов для передачи в реранкер (если не указано, то top_k*2)"),
+                            jina_api_key: Optional[str] = Query(None, description="API ключ Jina AI (если не указан, используется значение из переменной окружения JINA_API_KEY)")):
+    """
+    API endpoint для поиска фильмов.
+
+    Parameters
+    ----------
+    query : str
+        Поисковый запрос.
+    top_k : int, optional
+        Количество возвращаемых результатов, по умолчанию 25.
+    rerank_top_k : Optional[int], optional
+        Количество фильмов для передачи в реранкер. 
+        Если 0 - реранкер не используется.
+        Если не указано, то используется top_k*2.
+        По умолчанию None.
+    jina_api_key : Optional[str], optional
+        API ключ Jina AI. Если не указан, используется значение из переменной окружения JINA_API_KEY. 
+        По умолчанию None.
+
+    Returns
+    -------
+    dict
+        Словарь с результатами поиска.
+    """
     try:
-        results, _ = search_movies_internal(query, top_k, rerank_top_k)
+        results, _ = search_movies_internal(query, top_k, rerank_top_k, jina_api_key)
         return results
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))