update

common/decorators.py

@@ -0,0 +1,8 @@
+
+def singleton(cls):
+    instances = {}
+    def get_instance(*args, **kwargs):
+        if cls not in instances:
+            instances[cls] = cls(*args, **kwargs)
+        return instances[cls]
+    return get_instance

components/embedding_extraction.py

@@ -6,23 +6,27 @@ import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
 from transformers import (AutoModel, AutoTokenizer, BatchEncoding,
-                          XLMRobertaModel)
+                          XLMRobertaModel, PreTrainedTokenizer, PreTrainedTokenizerFast)
 from transformers.modeling_outputs import \
     BaseModelOutputWithPoolingAndCrossAttentions as EncoderOutput
 
-logger = logging.getLogger(__name__)
+from common.decorators import singleton
 
+logger = logging.getLogger(__name__)
 
+@singleton
 class EmbeddingExtractor:
     """Класс обрабатывает текст вопроса и возвращает embedding"""
 
     def __init__(
         self,
-        model_id: str,
+        model_id: str | None,
         device: str | torch.device | None = None,
         batch_size: int = 1,
         do_normalization: bool = True,
         max_len: int = 510,
+        model: XLMRobertaModel = None,
+        tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast = None
     ):
         """
         Класс, соединяющий в себе модель, токенизатор и параметры векторизации.
@@ -33,6 +37,8 @@ class EmbeddingExtractor:
             batch_size: Размер батча (по умолчанию - 1).
             do_normalization: Нормировать ли вектора (по умолчанию - True).
             max_len: Максимальная длина текста в токенах (по умолчанию - 510).
+            model: Экземпляр загруженной модели.
+            tokenizer: Экземпляр загруженного токенизатора.
         """
         if device is None:
             device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -40,11 +46,19 @@ class EmbeddingExtractor:
             device = torch.device(device)
 
         self.device = device
+        
         # Инициализация модели
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=True)
-        self.model: XLMRobertaModel = AutoModel.from_pretrained(model_id, local_files_only=True).to(
-            self.device
-        )
+        if model is not None and tokenizer is not None:
+            self.tokenizer = tokenizer
+            self.model = model
+        elif model_id is not None:
+            print('EmbeddingExtractor: model loading '+model_id+' to '+str(self.device))
+            self.tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=True)
+            self.model: XLMRobertaModel = AutoModel.from_pretrained(model_id, local_files_only=True).to(
+                self.device
+            )
+            
+            print('EmbeddingExtractor: model loaded')
         self.model.eval()
         self.model.share_memory()
 

components/llm/deepinfra_api.py

@@ -1,5 +1,5 @@
 import json
-from typing import Optional, List
+from typing import AsyncGenerator, Optional, List
 import httpx
 import logging
 from transformers import AutoTokenizer
@@ -286,7 +286,6 @@ class DeepInfraApi(LlmApi):
                         try:
                             # Парсим JSON из строки
                             data = json.loads(line[len("data: "):].strip())
-                            print(data)
                             if data == "[DONE]":  # Конец потока
                                 break
                             if "choices" in data and data["choices"]:
@@ -298,6 +297,47 @@ class DeepInfraApi(LlmApi):
 
             return generated_text.strip()
         
+    async def get_predict_chat_generator(self, request: ChatRequest, system_prompt: str, 
+                                     params: LlmPredictParams) -> AsyncGenerator[str, None]:
+        """
+        Выполняет потоковый запрос к API и возвращает токены по мере их генерации.
+
+        Args:
+            request (ChatRequest): История чата.
+            system_prompt (str): Системный промпт.
+            params (LlmPredictParams): Параметры предсказания.
+
+        Yields:
+            str: Токены ответа LLM.
+        """
+        params
+        async with httpx.AsyncClient() as client:
+            request_data = self.create_chat_request(request, system_prompt, params)
+            request_data["stream"] = True
+
+            async with client.stream(
+                "POST",
+                f"{self.params.url}/v1/openai/chat/completions",
+                json=request_data,
+                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')}")
+
+                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", "")
+                                if token_value:
+                                    yield token_value
+                        except json.JSONDecodeError:
+                            continue
+        
     async def predict(self, prompt: str, system_prompt: str) -> str:
         """
         Выполняет запрос к API и возвращает результат.

routes/llm.py

@@ -1,8 +1,11 @@
+import json
 import logging
 import os
-from typing import Annotated, Optional
+from typing import Annotated, AsyncGenerator, Optional
 from uuid import UUID
 
+from fastapi.responses import StreamingResponse
+
 from components.services.dataset import DatasetService
 from components.services.entity import EntityService
 from fastapi import APIRouter, Depends, HTTPException
@@ -42,6 +45,97 @@ llm_params = LlmParams(
 # TODO: унести в DI
 llm_api = DeepInfraApi(params=llm_params)
 
+# TODO: Вынести
+def get_last_user_message(chat_request: ChatRequest) -> Optional[Message]:
+    return next(
+        (
+            msg
+            for msg in reversed(chat_request.history)
+            if msg.role == "user"
+            and (msg.searchResults is None or not msg.searchResults)
+        ),
+        None,
+    )
+
+def insert_search_results_to_message(
+    chat_request: ChatRequest, new_content: str
+) -> bool:
+    for msg in reversed(chat_request.history):
+        if msg.role == "user" and (
+            msg.searchResults is None or not msg.searchResults
+        ):
+            msg.content = new_content
+            return True
+    return False
+        
+async def sse_generator(request: ChatRequest, llm_api: DeepInfraApi, system_prompt: str, 
+                        predict_params: LlmPredictParams,
+                        dataset_service: DatasetService, 
+                        entity_service: EntityService) -> AsyncGenerator[str, None]:
+    """
+    Генератор для стриминга ответа LLM через SSE.
+    """
+    # Обработка поиска
+    last_query = get_last_user_message(request)
+    if last_query:
+        dataset = dataset_service.get_current_dataset()
+        if dataset is None:
+            raise HTTPException(status_code=400, detail="Dataset not found")
+        _, scores, chunk_ids = entity_service.search_similar(last_query.content, dataset.id)
+        chunks = entity_service.chunk_repository.get_chunks_by_ids(chunk_ids)
+        text_chunks = entity_service.build_text(chunks, scores)
+        search_results_event = {
+            "event": "search_results",
+            "data": f"\n<search-results>\n{text_chunks}\n</search-results>"
+        }
+        yield f"data: {json.dumps(search_results_event, ensure_ascii=False)}\n\n"
+
+        new_message = f'{last_query.content}\n<search-results>\n{text_chunks}\n</search-results>'
+        insert_search_results_to_message(request, new_message)
+            
+    # Стриминг токенов ответа
+    async for token in llm_api.get_predict_chat_generator(request, system_prompt, predict_params):
+        token_event = {"event": "token", "data": token}
+        logger.info(f"Streaming token: {token}")
+        yield f"data: {json.dumps(token_event, ensure_ascii=False)}\n\n"
+
+    # Финальное событие
+    yield "data: {\"event\": \"done\"}\n\n"
+    
+
+@router.post("/chat/stream")
+async def chat_stream(
+    request: ChatRequest,
+    config: Annotated[Configuration, Depends(DI.get_config)],
+    llm_api: Annotated[DeepInfraApi, Depends(DI.get_llm_service)],
+    prompt_service: Annotated[LlmPromptService, Depends(DI.get_llm_prompt_service)],
+    llm_config_service: Annotated[LLMConfigService, Depends(DI.get_llm_config_service)],
+    entity_service: Annotated[EntityService, Depends(DI.get_entity_service)],
+    dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)],
+):
+    try:
+        p = llm_config_service.get_default()
+        system_prompt = prompt_service.get_default()
+
+        predict_params = LlmPredictParams(
+            temperature=p.temperature,
+            top_p=p.top_p,
+            min_p=p.min_p,
+            seed=p.seed,
+            frequency_penalty=p.frequency_penalty,
+            presence_penalty=p.presence_penalty,
+            n_predict=p.n_predict,
+            stop=[],
+        )
+
+        return StreamingResponse(
+            sse_generator(request, llm_api, system_prompt.text, predict_params, dataset_service, entity_service),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}
+        )
+    except Exception as e:
+        logger.error(f"Error in SSE chat stream: {str(e)}", stack_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
 
 @router.post("/chat")
 async def chat(
@@ -68,29 +162,6 @@ async def chat(
             stop=[],
         )
 
-        # TODO: Вынести
-        def get_last_user_message(chat_request: ChatRequest) -> Optional[Message]:
-            return next(
-                (
-                    msg
-                    for msg in reversed(chat_request.history)
-                    if msg.role == "user"
-                    and (msg.searchResults is None or not msg.searchResults)
-                ),
-                None,
-            )
-
-        def insert_search_results_to_message(
-            chat_request: ChatRequest, new_content: str
-        ) -> bool:
-            for msg in reversed(chat_request.history):
-                if msg.role == "user" and (
-                    msg.searchResults is None or not msg.searchResults
-                ):
-                    msg.content = new_content
-                    return True
-            return False
-
         last_query = get_last_user_message(request)
         search_result = None
         
@@ -126,4 +197,4 @@ async def chat(
         logger.error(
             f"Error processing LLM request: {str(e)}", stack_info=True, stacklevel=10
         )
-        return {"error": str(e)}
+        return {"error": str(e)}