update

common/auth.py

@@ -11,10 +11,14 @@ SECRET_KEY = os.environ.get("JWT_SECRET", "ooooooh_thats_my_super_secret_key")
 ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_MINUTES = 1440
 
+class User(BaseModel):
+    username: str
+    password: str
+    
 # Захардкоженные пользователи
 USERS = [
-    {"username": "admin", "password": "admin123"},
-    {"username": "demo", "password": "sTrUPsORPA"},
+    User(username="admin", password="admin123"),
+    User(username="demo", password="sTrUPsORPA")
 ]
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/auth/login/token")
@@ -39,7 +43,7 @@ async def get_current_user(token: str = Depends(oauth2_scheme)):
         username: str = payload.get("sub")
         if username is None:
             raise HTTPException(status_code=401, detail="Invalid token")
-        user = next((u for u in USERS if u["username"] == username), None)
+        user = next((u for u in USERS if u.username == username), None)
         if user is None:
             raise HTTPException(status_code=401, detail="User not found")
         return user

components/llm/deepinfra_api.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from typing import AsyncGenerator, Optional, List
 import httpx
@@ -256,13 +257,17 @@ class DeepInfraApi(LlmApi):
                 logging.error(f"Request failed: status code {response.status_code}")
                 logging.error(response.text)
                 
-    async def predict_chat_stream(self, request: ChatRequest, system_prompt, params: LlmPredictParams) -> str:
+    async def predict_chat_stream(self, request: ChatRequest, system_prompt, params: LlmPredictParams, max_retries: int = 3, retry_delay: float = 0.5) -> str:
         """
         Выполняет запрос к API с поддержкой потокового вывода (SSE) и возвращает результат.
-
+        
         Args:
-            prompt (str): Входной текст для предсказания.
-
+            request (ChatRequest): Запрос чата
+            system_prompt: Системный промпт
+            params (LlmPredictParams): Параметры предсказания
+            max_retries (int): Максимальное количество попыток переподключения (по умолчанию 3)
+            retry_delay (float): Задержка между попытками в секундах (по умолчанию 0.5)
+        
         Returns:
             str: Сгенерированный текст.
         """
@@ -270,32 +275,41 @@ class DeepInfraApi(LlmApi):
             request = self.create_chat_request(request, system_prompt, params)
             request["stream"] = True
             
-            print(super().create_headers())
-            async with client.stream("POST", f"{self.params.url}/v1/openai/chat/completions", json=request, headers=super().create_headers()) as response:
-                if response.status_code != 200:
-                    # Если ошибка, читаем ответ для получения подробностей
-                    error_content = await response.aread()
-                    raise Exception(f"API error: {error_content.decode('utf-8')}")
-
-                # Для хранения результата
-                generated_text = ""
-
-                # Асинхронное чтение построчно
-                async for line in response.aiter_lines():
-                    if line.startswith("data: "):  # SSE-сообщения начинаются с "data: "
-                        try:
-                            # Парсим JSON из строки
-                            data = json.loads(line[len("data: "):].strip())
-                            if data == "[DONE]":  # Конец потока
-                                break
-                            if "choices" in data and data["choices"]:
-                                # Получаем текст из текущего токена
-                                token_value = data["choices"][0].get("delta", {}).get("content", "")
-                                generated_text += token_value
-                        except json.JSONDecodeError:
-                            continue  # Игнорируем строки, которые не удается декодировать
+            for attempt in range(max_retries + 1):
+                try:
+                    async with client.stream("POST", f"{self.params.url}/v1/openai/chat/completions", 
+                                        json=request, 
+                                        headers=super().create_headers()) as response:
+                        
+                        if response.status_code != 200:
+                            error_content = await response.aread()
+                            raise Exception(f"API error: {error_content.decode('utf-8')}")
+
+                        generated_text = ""
+                        
+                        async for line in response.aiter_lines():
+                            if line.startswith("data: "):
+                                try:
+                                    data = json.loads(line[len("data: "):].strip())
+                                    if data == "[DONE]":
+                                        break
+                                    if "choices" in data and data["choices"]:
+                                        token_value = data["choices"][0].get("delta", {}).get("content", "")
+                                        generated_text += token_value
+                                except json.JSONDecodeError:
+                                    continue
+                        
+                        return generated_text.strip()
+                        
+                except Exception as e:
+                    if attempt < max_retries:
+                        # Ждем перед следующей попыткой, если это не последняя попытка
+                        await asyncio.sleep(retry_delay)
+                        continue
+                    else:
+                        # Если исчерпаны все попытки, пробрасываем исключение
+                        raise Exception(f"predict_chat_stream failed after {max_retries} retries: {str(e)}")
 
-            return generated_text.strip()
         
     async def get_predict_chat_generator(self, request: ChatRequest, system_prompt: str, 
                                      params: LlmPredictParams) -> AsyncGenerator[str, None]:

components/services/dialogue.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import re
-from typing import List
+from typing import List, Optional, Tuple
 
 from pydantic import BaseModel
 
@@ -19,6 +19,7 @@ logger = logging.getLogger(__name__)
 class QEResult(BaseModel):
     use_search: bool
     search_query: str | None
+    debug_message: Optional[str | None] = ""
 
 
 class DialogueService:
@@ -71,6 +72,7 @@ class DialogueService:
             return QEResult(
                 use_search=from_chat is not None,
                 search_query=from_chat.content if from_chat else None,
+                debug_message=response
             )
             
     def get_qe_result_from_chat(self, history: List[Message]) -> QEResult:
@@ -129,7 +131,8 @@ class DialogueService:
         else:
             raise ValueError("Первая часть текста должна содержать 'ДА' или 'НЕТ'.")
 
-        return QEResult(use_search=bool_var, search_query=second_part)
+        return QEResult(use_search=bool_var, search_query=second_part,
+                debug_message=input_text)
 
     def _get_search_query(self, history: List[Message]) -> Message | None:
         """

routes/auth.py

@@ -1,5 +1,5 @@
-from typing import Optional
-from fastapi import APIRouter, Body, Form, HTTPException
+from typing import Annotated, Optional
+from fastapi import APIRouter, Body, Depends, Form, HTTPException
 from datetime import timedelta
 import common.auth as auth
 
@@ -7,8 +7,8 @@ router = APIRouter(prefix="/auth", tags=["Auth"])
 
 def authenticate_user(username: str, password: str):
     """Проверяет, существует ли пользователь и правильный ли пароль."""
-    user = next((u for u in auth.USERS if u["username"] == username), None)
-    if not user or user["password"] != password:
+    user = next((u for u in auth.USERS if u.username == username), None)
+    if not user or user.password != password:
         raise HTTPException(status_code=401, detail="Неверный логин или пароль")
     return user
 
@@ -20,7 +20,7 @@ def generate_access_token(username: str):
 async def login_common(username: str, password: str):
     """Общий метод аутентификации."""
     user = authenticate_user(username, password)
-    access_token = generate_access_token(user["username"])
+    access_token = generate_access_token(user.username)
     return {"access_token": access_token, "token_type": "bearer"}
 
 @router.post("/login", summary="Авторизация через JSON")
@@ -31,4 +31,8 @@ async def login_json(request: auth.LoginRequest = Body(...)):
 @router.post("/login/token", summary="Авторизация через Form-Data")
 async def login_form(username: str = Form(...), password: str = Form(...)):
     """Принимает Form-Data (x-www-form-urlencoded)."""
-    return await login_common(username, password)
+    return await login_common(username, password)
+
+@router.post("/checktoken", summary="Проверяет, аутентифицирован ли пользователь")
+async def check_token(current_user: Annotated[auth.User, Depends(auth.get_current_user)]):
+    return {"current_user": current_user.username}

routes/llm.py

@@ -123,7 +123,13 @@ async def sse_generator(request: ChatRequest, llm_api: DeepInfraApi, system_prom
     """
     try:
         qe_result = await dialogue_service.get_qe_result(request.history) 
-        
+        qe_event = {
+            "event": "debug",
+            "data": {
+                "text": qe_result.debug_message
+            }
+        }
+        yield f"data: {json.dumps(qe_event, ensure_ascii=False)}\n\n"
     except Exception as e:
         logger.error(f"Error in SSE chat stream while dialogue_service.get_qe_result: {str(e)}", stack_info=True)
         yield "data: {\"event\": \"error\", \"data\":\""+str(e)+"\" }\n\n"