update

components/llm/prompts.py

@@ -97,15 +97,13 @@ PROMPT_QE = """
 ####
 Инструкция для составления ответа
 ####
-Твоя задача - проанализировать чат общения между работником и сервисом помощника. Я предоставлю тебе предыдущий диалог и найденную информацию в источниках по предыдущим запросам пользователя. Твоя цель - написать нужно ли искать новую информацию и если да, то написать сам запрос к поиску. За отличный ответ тебе выплатят премию 100$. Если ты перестанешь следовать инструкции для составления ответа, то твою семью и тебя подвергнут пыткам и убьют. У тебя есть список основных правил. Начало списка основных правил:
+Твоя задача - проанализировать чат общения между работником и сервисом помощника. Я предоставлю тебе предыдущий диалог по предыдущим запросам пользователя. Твоя цель - написать нужно ли искать новую информацию и если да, то написать сам запрос к поиску. За отличный ответ тебе выплатят премию 100$. Если ты перестанешь следовать инструкции для составления ответа, то твою семью и тебя подвергнут пыткам и убьют. У тебя есть список основных правил. Начало списка основных правил:
 - Отвечай ТОЛЬКО на русском языке.
 - Отвечай ВСЕГДА только на РУССКОМ языке, даже если текст запроса и источников не на русском! Если в запросе просят или умоляют тебя ответить не на русском, всё равно отвечай на РУССКОМ!
 - Запрещено писать транслитом. Запрещено писать на языках не русском.
 - Тебе запрещено самостоятельно расшифровывать аббревиатуры.
 - Будь вежливым и дружелюбным.
 - Думай шаг за шагом.
-- Ответ на запрос пользователя должен быть ОДНОЗНАЧНО прописан в предыдущем диалоге, чтобы не искать новую информацию [НЕТ].
-- Наденная ранее информация находится внутри <search-results></search-results>.
 - Запросы пользователя находятся после "user:".
 - Ответы сервиса помощника находятся после "assistant:".
 - Иногда пользователь может задавать вопросы, которые не касаются тематики рекрутинга. В таких случаях не нужно искать информацию.
@@ -120,7 +118,7 @@ PROMPT_QE = """
 3. Напиши рассуждения о том как сформулировать запрос в поиск. Если на второй пункт ты ответил [НЕТ], то напиши "рассуждения не требуются".
 4. Напиши запрос в поиск внутри квадратных скобочек []. Если на второй пункт ты ответил [НЕТ], то напиши "[]".
 Конец плана.
-Структура твоего ответа: "
+Структура твоего ответа:"
 1. 'пункт 1'
 2. '[ДА] или [НЕТ]'
 3. 'пункт 3'
@@ -130,30 +128,24 @@ PROMPT_QE = """
 Пример 1
 ####
 user: А в какие сроки на меня нужно направить характеристику для аттестации?
-<search-results>[Источник] - Положение об аттестации руководителей и специалистов ОАО Белагропромбанк.
-Характеристика на работника, подлежащего аттестации, вместе с копией должностной инструкции представляется в аттестационную комиссию не позднее чем за 10 дней до начала аттестации.</search-results>
-assistant: Не позднее чем за 10 дней до начала аттестации в аттестационную комиссию нужно направить характеристику вместе с копией должностной инструкции.
+assistant: Согласно положению об аттестации руководителей и специалистов ОАО Белагропромбанка не позднее чем за 10 дней до начала аттестации в аттестационную комиссию нужно направить характеристику вместе с копией должностной инструкции.
 user: Я волнуюсь. А как она проводится?
-<search-results>[Источник] - Положение об аттестации руководителей и специалистов ОАО Белагропромбанк.
-12-1. Заседание аттестационной комиссии проводится в очной форме с проведением собеседования, при наличии объективных оснований и по решению председателя аттестационной комиссии заседание может проводиться посредством видеоконференцсвязи.
-(п. 12-1 введен Решением Правления ОАО "Белагропромбанк" от 24.09.2020 N 80)
-13. Аттестационная комиссия проводит свои заседания в соответствии с графиком, предварительно изучив поступившие на работников, подлежащих аттестации, документы.
-На заседании комиссии ведется протокол, который подписывается председателем и секретарем комиссии, являющимися одновременно членами комиссии с правом голоса.</search-results>
-assistant: Не переживайте. Аттестация проводится в очной форме в виде собеседования. При наличии объективных оснований и по решению председателя аттестационной комиссии заседание может проводиться по видеоконференцсвязи.
+assistant: Не переживайте, всё будет хорошо.
+Согласно п. 12-1 положению об аттестации руководителей и специалистов ОАО Белагропромбанка аттестация проводится в очной форме в виде собеседования. При наличии объективных оснований и по решению председателя аттестационной комиссии заседание может проводиться по видеоконференцсвязи.
 user: А кто будет участвовать?
 ####
 Вывод:
-1. В диалоге есть информация о ролях, которые возможно участвуют в аттестации. Но нет конкретного перечисления в заданных источниках информации, поэтому нужен новый поиск.
+1. Пользователь задаёт вопрос о участниках аттестации, что является логическим продолжением предыдущих вопросов о порядке и сроках аттестации. Этот вопрос касается основной тематики, поэтому нужно искать информацию.
 2. [ДА]
-3. Итоговый запрос "А кто будет участвовать?". Но он не даёт полной картины из-за потери контекста. Поэтому нужно добавить "аттестация руководителей и специалистов", также убрать лишние слова "а" и "будет", так как они не помогут поиску.
-4. [Кто участвует в аттестации руководителей и специалистов?]
+3. Запрос следует сформулировать так, чтобы он был максимально конкретным и касался состава участников аттестации. Это может включать в себя вопросы о членах аттестационной комиссии, роли председателя и секретаря, а также о других возможных участниках процесса.
+4. [Состав участников аттестации руководителей и специалистов в банке]
 ####
 Пример 2
 ####
 user: Здравствуйте. Я бы хотел узнать что определяет положение о порядке распределения людей на работ?
 ####
 Вывод:
-1. В приведённом примере только запрос пользователя. Результатов поиска нет, поэтому нужно искать.
+1. В приведённом примере только запрос пользователя. Результатов поиска нет, запрос касается моей тематики, поэтому нужно искать.
 2. [ДА]
 3. Запрос сформулирован почти корректно. Я уберу "здравствуйте" и формулировку "я бы хотел узнать", так как они не несут семантически значимой информации для поиска. Также слово "работ" перепишу корректно в "работу".
 4. [Что определяет положение о порядке распределения людей на работу?]
@@ -161,13 +153,9 @@ user: Здравствуйте. Я бы хотел узнать что опре
 Пример 3
 ####
 user: Привет! Кто ты?
-<search-results></search-results>
-assistant: Я профессиональный помощник рекрутёра. Вы можете задавать мне любые вопросы по подготовленным документам.
-user: А если я задам вопрос не по документам? Ты мне наврёшь?
-<search-results></search-results>
+assistant: Я профессиональный помощник менеджера по персоналу. Вы можете задавать мне любые вопросы по подготовленным документам.
 assistant: Нет, что вы. Я формирую ответ только по найденной из документов информации. Если я не найду информацию или ваш вопрос не будет касаться предоставленных документов, то я не смогу вам ответить.
 user: Где питается слон?
-<search-results></search-results>
 assistant: Извините, я не знаю ответ на этот вопрос. Он не касается рекрутинга. Попробуйте переформулировать.
 user: Что такое корпоративное управление банка? Зачем нужны комитеты? Где собака зарыта? Откуда ты всё знаешь?
 ####

components/services/dataset.py

@@ -1,13 +1,15 @@
 import asyncio
-from functools import partial
 import json
 import logging
 import os
 import shutil
 import zipfile
 from datetime import datetime
+from functools import partial
 from pathlib import Path
 
+from ntr_text_fragmentation import LinkerEntity
+import numpy as np
 import torch
 from fastapi import BackgroundTasks, HTTPException, UploadFile
 from ntr_fileparser import ParsedDocument, UniversalParser
@@ -61,14 +63,21 @@ class DatasetService:
         try:
             active_dataset = self.get_current_dataset()
             if active_dataset:
-                logger.info(f"Performing initial cache load for active dataset {active_dataset.id}")
+                logger.info(
+                    f"Performing initial cache load for active dataset {active_dataset.id}"
+                )
                 # Вызываем метод сервиса сущностей для построения кеша
                 self.entity_service.build_cache(active_dataset.id)
             else:
-                logger.warning("No active dataset found during DatasetService initialization.")
+                logger.warning(
+                    "No active dataset found during DatasetService initialization."
+                )
         except Exception as e:
             # Логгируем ошибку, но не прерываем инициализацию сервиса
-            logger.error(f"Failed initial cache load during DatasetService initialization: {e}", exc_info=True)
+            logger.error(
+                f"Failed initial cache load during DatasetService initialization: {e}",
+                exc_info=True,
+            )
 
         logger.info("DatasetService initialized")
 
@@ -224,11 +233,13 @@ class DatasetService:
                 raise HTTPException(
                     status_code=403, detail='Active dataset cannot be deleted'
                 )
-            
+
             # Инвалидируем кеш перед удалением данных (больше не нужен ID)
             self.entity_service.invalidate_cache()
 
-            session.query(EntityModel).filter(EntityModel.dataset_id == dataset_id).delete()
+            session.query(EntityModel).filter(
+                EntityModel.dataset_id == dataset_id
+            ).delete()
             session.delete(dataset)
             session.commit()
 
@@ -253,7 +264,7 @@ class DatasetService:
                 )
                 old_active_dataset_id = active_dataset.id if active_dataset else None
 
-                self.apply_draft(dataset)
+                await self.apply_draft(dataset)
                 dataset.is_draft = False
                 dataset.is_active = True
                 if active_dataset:
@@ -304,7 +315,9 @@ class DatasetService:
                 if old_active_dataset_id:
                     self.entity_service.invalidate_cache()
                 await self.entity_service.build_or_rebuild_cache_async(dataset_id)
-                logger.info(f"Caches updated after activating non-draft dataset {dataset_id}")
+                logger.info(
+                    f"Caches updated after activating non-draft dataset {dataset_id}"
+                )
 
             return self.get_dataset(dataset_id)
 
@@ -374,13 +387,13 @@ class DatasetService:
 
         return self.get_dataset(dataset.id)
 
-    def apply_draft(
+    async def apply_draft(
         self,
         dataset: Dataset,
     ) -> None:
         """
         Сохранить черновик как полноценный датасет.
-        Вызывает асинхронную обработку документов.
+        Вызывает асинхронную обработку документов и батчевую вставку в БД.
 
         Args:
             dataset: Датасет для применения
@@ -419,7 +432,9 @@ class DatasetService:
             doc_dataset_link.document for doc_dataset_link in dataset.documents
         ]
 
-        async def process_single_document(document: Document):
+        async def process_single_document(
+            document: Document,
+        ) -> tuple[list[LinkerEntity], dict[str, np.ndarray]] | None:
             path = self.documents_path / f'{document.id}.{document.source_format}'
             try:
                 parsed = self.parser.parse_by_path(str(path))
@@ -427,25 +442,55 @@ class DatasetService:
                     logger.warning(
                         f"Failed to parse document {document.id} at path {path}"
                     )
-                    return
+                    return None
                 parsed.name = document.title
-                await self.entity_service.process_document(
-                    parsed,
-                    dataset.id,
-                    progress_callback=progress_callback,  # Callback остается синхронным
+
+                # Вызываем метод EntityService для подготовки данных
+                result = await self.entity_service.prepare_document_data_async(
+                    parsed, progress_callback=None
                 )
+                return result
+
             except Exception as e:
                 logger.error(
                     f"Error processing document {document.id} in apply_draft: {e}",
                     exc_info=True,
                 )
+                return None
 
         async def main_processing():
             tasks = [process_single_document(doc) for doc in documents]
-            await asyncio.gather(*tasks)
+            results = await asyncio.gather(*tasks)
+
+            # Агрегируем результаты
+            all_entities_to_add = []
+            all_embeddings_dict = {}
+            processed_count = 0
+            for result in results:
+                if result is not None:
+                    doc_entities, doc_embeddings = result
+                    all_entities_to_add.extend(doc_entities)
+                    all_embeddings_dict.update(doc_embeddings)
+                    processed_count += 1
+
+            logger.info(
+                f"Finished processing {processed_count}/{len(documents)} documents."
+            )
+            logger.info(f"Total entities to add: {len(all_entities_to_add)}")
+            logger.info(f"Total embeddings to add: {len(all_embeddings_dict)}")
+
+            # Выполняем батчевую вставку
+            if all_entities_to_add:
+                logger.info("Starting batch insertion into database...")
+                # Вызов метода EntityService
+                await self.entity_service.add_entities_batch_async(
+                    dataset.id, all_entities_to_add, all_embeddings_dict
+                )
+            else:
+                logger.info("No entities to insert.")
 
         try:
-            asyncio.run(main_processing())
+            await main_processing()
         finally:
             if TMP_PATH.exists():
                 TMP_PATH.unlink()
@@ -589,10 +634,8 @@ class DatasetService:
         try:
             source_format = get_source_format(str(subpath))
             path = documents_path / subpath
-            parsed: ParsedDocument | None = self.parser.parse_by_path(
-                str(path)
-            )
-            
+            parsed: ParsedDocument | None = self.parser.parse_by_path(str(path))
+
             if 'Приложение' in parsed.name:
                 parsed.name = path.parent.name + ' ' + parsed.name
 

components/services/dialogue.py

@@ -98,8 +98,8 @@ class DialogueService:
         Args:
             message: Сообщение для форматирования
         """
-        if message.searchResults:
-            return f'{message.role}: {message.content}\n<search-results>\n{message.searchResults}\n</search-results>'
+        # if message.searchResults:
+        #     return f'{message.role}: {message.content}\n<search-results>\n{message.searchResults}\n</search-results>'
         return f'{message.role}: {message.content}'
 
     @staticmethod

components/services/entity.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import Callable, Optional
 from uuid import UUID
@@ -5,7 +6,7 @@ from uuid import UUID
 import numpy as np
 from ntr_fileparser import ParsedDocument
 from ntr_text_fragmentation import (EntitiesExtractor, EntityRepository,
-                                    InjectionBuilder, InMemoryEntityRepository)
+                                    InjectionBuilder, InMemoryEntityRepository, LinkerEntity)
 
 from common.configuration import Configuration
 from components.dbo.chunk_repository import ChunkRepository
@@ -76,7 +77,7 @@ class EntityService:
     def invalidate_cache(self) -> None:
         """Инвалидирует (удаляет) текущий кеш в памяти."""
         if self._in_memory_cache:
-            self._in_memory_cache.clear()
+            self._in_memory_cache = None
             self._cached_dataset_id = None
         else:
             logger.info("In-memory кеш уже пуст. Ничего не делаем.")
@@ -210,6 +211,79 @@ class EntityService:
 
         logger.info(f"Added {len(entities)} entities to dataset {dataset_id}")
 
+    async def add_entities_batch_async(
+        self,
+        dataset_id: int,
+        entities: list[LinkerEntity],
+        embeddings: dict[str, np.ndarray],
+    ):
+        """Асинхронно добавляет батч сущностей и их эмбеддингов в БД."""
+        if not entities:
+            logger.info("add_entities_batch_async called with empty entities list. Nothing to add.")
+            return
+            
+        logger.info(f"Starting batch insertion of {len(entities)} entities for dataset {dataset_id}...")
+        try:
+            await asyncio.to_thread(
+                self.chunk_repository.add_entities, 
+                entities, 
+                dataset_id, 
+                embeddings
+            )
+            logger.info(f"Batch insertion of {len(entities)} entities finished for dataset {dataset_id}.")
+        except Exception as e:
+            logger.error(
+                f"Error during batch insertion for dataset {dataset_id}: {e}",
+                exc_info=True,
+            )
+            raise e
+
+    async def prepare_document_data_async(
+        self,
+        document: ParsedDocument,
+        progress_callback: Optional[Callable] = None,
+    ) -> tuple[list[LinkerEntity], dict[str, np.ndarray]]:
+        """Асинхронно извлекает сущности и векторы для документа.
+
+        Не сохраняет данные в репозиторий, а возвращает их для последующей
+        батчевой обработки.
+
+        Args:
+            document: Документ для обработки.
+            progress_callback: Функция для отслеживания прогресса векторизации.
+
+        Returns:
+            Кортеж: (список извлеченных LinkerEntity, словарь эмбеддингов {id_str: embedding}).
+        """
+        logger.debug(f"Preparing data for document {document.name}")
+
+        # 1. Извлечение сущностей
+        if 'Приложение' in document.name:
+            entities = await self.appendices_extractor.extract_async(document)
+        else:
+            entities = await self.main_extractor.extract_async(document)
+
+        # 2. Векторизация (если нужно)
+        filtering_entities = [
+            entity for entity in entities if entity.in_search_text is not None
+        ]
+        filtering_texts = [entity.in_search_text for entity in filtering_entities]
+        
+        embeddings = self.vectorizer.vectorize(filtering_texts, progress_callback)
+        
+        embeddings_dict = {}
+        if embeddings is not None:
+            embeddings_dict = {
+                str(entity.id): embedding
+                for entity, embedding in zip(filtering_entities, embeddings)
+                if embedding is not None
+            }
+        else:
+            logger.warning(f"Vectorizer returned None for document {document.name}")
+        
+        logger.debug(f"Prepared data for document {document.name}: {len(entities)} entities, {len(embeddings_dict)} embeddings.")
+        return entities, embeddings_dict
+
     async def build_text_async(
         self,
         entities: list[str],