Update app.py

app.py

@@ -51,6 +51,9 @@ if JINA_API_KEY is None:
     raise ValueError("JINA_API_KEY environment variable not set.")
 JINA_RERANKER_MODEL = "jina-reranker-v2-base-multilingual"
 
+# Jina AI Dashboard API
+JINA_DASHBOARD_API_URL = 'https://embeddings-dashboard-api.jina.ai/api/v1/api_key/user'
+
 # Имена таблиц
 embeddings_table = "movie_embeddings"
 query_cache_table = "query_cache"
@@ -170,6 +173,23 @@ def get_movie_data_from_db(conn, movie_ids):
         logging.error(f"Ошибка при получении данных фильмов из БД: {e}")
     return movie_data_dict
 
+def get_jina_ai_balance():
+    """Получает остаток баланса Jina AI."""
+    try:
+        headers = {
+            'Content-Type': 'application/json'
+        }
+        params = {
+            'api_key': JINA_API_KEY
+        }
+        response = requests.get(JINA_DASHBOARD_API_URL, headers=headers, params=params)
+        response.raise_for_status()
+        data = response.json()
+        return data['wallet']['total_balance']
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Ошибка при запросе к API баланса Jina AI: {e}")
+        return None
+
 def rerank_with_api(query, results, top_k, rerank_top_k=None):
     """Переранжирует результаты с помощью Jina AI Reranker API."""
     logging.info(f"Начало переранжирования для запроса: '{query}'")
@@ -195,10 +215,10 @@ def rerank_with_api(query, results, top_k, rerank_top_k=None):
     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'])}")
+    logging.info(f"Отправка данных на реранжировку (documents count): {len(data['documents'])}, top_n: {data['top_n']}")
 
     try:
         response = requests.post(JINA_API_URL, headers=headers, json=data)
@@ -216,11 +236,11 @@ def rerank_with_api(query, results, top_k, rerank_top_k=None):
             logging.warning("Ответ от API не содержит ключа 'results'.")
 
         logging.info("Переранжирование завершено.")
-        return reranked_results, True
+        return reranked_results, True, data["top_n"]
 
     except requests.exceptions.RequestException as e:
         logging.error(f"Ошибка при запросе к API реранжировщика: {e}")
-        return results, False
+        return results, False, data["top_n"]
 
 def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None):
     """Внутренняя функция для поиска фильмов по запросу (используется и в Gradio, и в API)."""
@@ -250,6 +270,9 @@ def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None
                 logging.error(f"Ошибка при сохранении эмбеддинга запроса: {e}")
                 conn.rollback()
 
+        # Определяем количество фильмов для запроса из БД
+        db_limit = rerank_top_k or top_k * 2
+
         # Используем косинусное расстояние для поиска
         try:
             with conn.cursor() as cur:
@@ -263,7 +286,7 @@ def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None
                     FROM "{embeddings_table}" m, query_embedding
                     ORDER BY similarity DESC
                     LIMIT %s
-                """, (query_crc32, int(top_k * 2)))
+                """, (query_crc32, int(db_limit)))
 
                 results = cur.fetchall()
             logging.info(f"Найдено {len(results)} предварительных результатов поиска.")
@@ -274,11 +297,12 @@ def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None
             conn.close()
 
         # Переранжируем результаты с помощью API
-        reranked_results, rerank_success = rerank_with_api(query, results, top_k, rerank_top_k)
-        
+        reranked_results, rerank_success, reranked_count = rerank_with_api(query, results, top_k, rerank_top_k)
+
         if not rerank_success:
             logging.warning("Переранжировка не удалась, используются сырые результаты.")
-            reranked_results = results[:top_k] # Используем срез для ограничения количества результатов
+            reranked_results = results[:top_k]  # Используем срез для ограничения количества результатов
+            reranked_count = 0
         else:
             reranked_results = reranked_results[:top_k]
 
@@ -317,7 +341,8 @@ def search_movies_internal(query: str, top_k: int = 25, rerank_top_k: int = None
                     "year": movie_data['year'],
                     "genres": [genre['name'] for genre in movie_data['genres']],
                     "description": movie_data.get('description', ''),
-                    "relevance_score": score if rerank_success else (movie_data_dict.get(movie_id, (None, None))[1] if movie_data_dict.get(movie_id, (None, None)) is not None else 0.0) # Сохраняем similarity, если нет реранжировки
+                    "relevance_score": score if rerank_success else (
+                        movie_data_dict.get(movie_id, (None, None))[1] if movie_data_dict.get(movie_id,(None, None)) is not None else 0.0)  # Сохраняем similarity, если нет реранжировки
                 })
             else:
                 logging.warning(f"Данные для фильма с ID {movie_id} не найдены в БД.")
@@ -325,19 +350,30 @@ 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()
+
         return {
             "status": "success",
             "results": formatted_results,
             "search_time": search_time,
             "total_movies": total_movies,
-            "searched_movies": searched_movies
+            "searched_movies": searched_movies,
+            "returned_movies": len(formatted_results),  # Количество возвращенных фильмов
+            "reranked_movies": reranked_count,  # Количество фильмов, обработанных реранкером
+            "jina_balance": jina_balance  # Остаток баланса Jina AI
         }, search_time
 
     except Exception as e:
         logging.error(f"Ошибка при выполнении поиска: {e}")
         return {
             "status": "error",
-            "message": str(e)
+            "message": str(e),
+            "search_time": 0,
+            "total_movies": 0,
+            "searched_movies": 0,
+            "returned_movies": 0,
+            "reranked_movies": 0,
+            "jina_balance": None
         }, 0
 
 @app.get("/search/", response_model=dict)
@@ -603,4 +639,4 @@ async def root():
 
 # Запускаем FastAPI
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0")