update

common/auth.py

@@ -14,7 +14,7 @@ ACCESS_TOKEN_EXPIRE_MINUTES = 30
 # Захардкоженные пользователи
 USERS = [
     {"username": "admin", "password": "admin123"},
-    {"username": "user", "password": "user123"},
+    {"username": "demo", "password": "sTrUPsORPA"},
 ]
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/auth/login/token")

components/dbo/models/llm_config.py

@@ -16,7 +16,7 @@ class LLMConfig(Base):
 
     __tablename__ = "llm_config"
 
-    is_default: Mapped[bool] = mapped_column(Boolean, is_default=False)
+    is_default: Mapped[bool] = mapped_column(Boolean, default=False)
     model: Mapped[String] = mapped_column(String)
     temperature: Mapped[float] = mapped_column(Float)
     top_p: Mapped[float] = mapped_column(Float)

components/dbo/models/llm_prompt.py

@@ -14,7 +14,7 @@ class LlmPrompt(Base):
 
     __tablename__ = "llm_prompt"
 
-    is_default: Mapped[bool] = mapped_column(Boolean, is_default=False)
+    is_default: Mapped[bool] = mapped_column(Boolean, default=False)
     name: Mapped[String] = mapped_column(String)
     text: Mapped[String] = mapped_column(String)
     type: Mapped[String] = mapped_column(String)

components/llm/common.py

@@ -73,6 +73,7 @@ class Message(BaseModel):
     role: str
     content: str
     searchResults: str
+    searchEntities: Optional[List[str]] = []
 
 class ChatRequest(BaseModel):
     history: List[Message]

routes/llm.py

@@ -1,9 +1,10 @@
 import json
 import logging
 import os
-from typing import Annotated, AsyncGenerator, Optional
+from typing import Annotated, AsyncGenerator, List, Optional
 from uuid import UUID
 
+from common import auth
 from components.services.dialogue import DialogueService
 from fastapi.responses import StreamingResponse
 
@@ -70,11 +71,12 @@ def insert_search_results_to_message(
     return False
 
 def try_insert_search_results(
-    chat_request: ChatRequest, search_results: str
+    chat_request: ChatRequest, search_results: str, entities: List[str]
 ) -> bool:
     for msg in reversed(chat_request.history):
         if msg.role == "user" and not msg.searchResults:
             msg.searchResults = search_results
+            msg.searchEntities = entities
             return True
     return False
 
@@ -117,8 +119,9 @@ async def sse_generator(request: ChatRequest, llm_api: DeepInfraApi, system_prom
     """
     Генератор для стриминга ответа LLM через SSE.
     """
+    qe_result = None
     try:
-        qe_result = await dialogue_service.get_qe_result(request.history)    
+        qe_result = await dialogue_service.get_qe_result(request.history) 
         
     except Exception as e:
         logger.error(f"Error in SSE chat stream while dialogue_service.get_qe_result: {str(e)}", stack_info=True)
@@ -134,13 +137,16 @@ async def sse_generator(request: ChatRequest, llm_api: DeepInfraApi, system_prom
             text_chunks = entity_service.build_text(chunks, scores)
             search_results_event = {
                 "event": "search_results",
-                "data": f"{text_chunks}"
+                "data": {
+                    "text": text_chunks, 
+                    "ids": chunk_ids.tolist()
+                }
             }
             yield f"data: {json.dumps(search_results_event, ensure_ascii=False)}\n\n"
 
             # new_message = f'<search-results>\n{text_chunks}\n</search-results>\n{last_query.content}'
             
-            try_insert_search_results(request, text_chunks)
+            try_insert_search_results(request, text_chunks, chunk_ids.tolist())
     except Exception as e:
         logger.error(f"Error in SSE chat stream while searching: {str(e)}", stack_info=True)
         yield "data: {\"event\": \"error\", \"data\":\""+str(e)+"\" }\n\n"
@@ -171,6 +177,7 @@ async def chat_stream(
     entity_service: Annotated[EntityService, Depends(DI.get_entity_service)],
     dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)],
     dialogue_service: Annotated[DialogueService, Depends(DI.get_dialogue_service)],
+    current_user: Annotated[any, Depends(auth.get_current_user)]
 ):
     try:
         p = llm_config_service.get_default()