components/elastic/create_index_elastic_group.py

import json
import logging
import time
from pathlib import Path

from elasticsearch import Elasticsearch
from tqdm import tqdm


def create_index_elastic_group(
    path: str,
    logger: logging.Logger | None = None,
):
    if logger is None:
        logger = logging.getLogger(__name__)

    # Подключение к Elasticsearch
    es = Elasticsearch(hosts='localhost:9200')
    INDEX_NAME = 'group_search_elastic_nn'

    # Удаление старого индекса, если он существует
    if es.indices.exists(index=INDEX_NAME):
        es.indices.delete(index=INDEX_NAME)

    mapping = {
        "mappings": {
            "properties": {
                "group_name_nn": {"type": "text", "analyzer": "standard"},
                "group_composition_nn": {
                    "type": "nested",
                    "properties": {
                        "person_name_nn": {"type": "text", "analyzer": "standard"},
                        "position_in_group_nn": {
                            "type": "text",
                            "analyzer": "standard",
                        },
                    },
                },
            }
        }
    }

    # Создание индекса с указанным маппингом
    es.indices.create(index=INDEX_NAME, body=mapping)

    for ind, path in tqdm(enumerate(Path(path).iterdir())):
        # Открываем файл и читаем его содержимое
        with open(path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        # Индексирование документа в Elasticsearch
        es.index(index=INDEX_NAME, id=ind + 1, body=data)

        # Подсчет количества документов в индексе
        count_response = es.count(index=INDEX_NAME)
        logger.info(
            f"{ind}, Total documents in '{INDEX_NAME}': {count_response['count']}"
        )
        time.sleep(0.5)

    if es.indices.exists(index=INDEX_NAME):
        logger.info(f"Index '{INDEX_NAME}' exists.")

    # Подсчет количества документов в индексе
    count_response = es.count(index=INDEX_NAME)
    logger.info(f"Total documents in '{INDEX_NAME}': {count_response['count']}")

    query = "Какие действия являются первоочередными в момент обнаружения происшествия?"

    # Поиск документов, где поле "person_full_name" содержит определенное значение "Александров Д.В."
    # query_ = {
    #     "query": {
    #         "function_score": {
    #             "query": {
    #                 "multi_match": {
    #                     "query": f"{query}",
    #                     "fields": ["group_name"],
    #                     "fuzziness": "AUTO",
    #                     "analyzer": "standard"
    #                 }
    #             },
    #             "functions": [
    #                 {
    #                     "filter": {
    #                         "multi_match": {
    #                             "query": "персонального состава Персональный состав Комитета ПАО ГМК Норильский никель Рабочей группы",
    #                             "fields": ["group_name"],
    #                             "operator": "or"
    #                         }
    #                     },
    #                     "weight": 0.9  #// Понижает вес документов с этими словами
    #                 }
    #             ],
    #             "boost_mode": "multiply"  # // Умножает вес документов с фильтром на указанный коэффициент
    #         }
    #     }
    # }
    query_ = {
        "query": {
            "bool": {
                "should": [
                    {
                        "multi_match": {
                            "query": f"{query}",
                            "fields": ["group_name"],
                            "fuzziness": "AUTO",
                            "analyzer": "standard",
                        }
                    },
                    {
                        "multi_match": {
                            "query": "персонального состава Персональный состав Комитета ПАО ГМК Норильский никель Рабочей группы",
                            "fields": ["group_name"],
                            "operator": "or",
                            "boost": 0.1,
                        }
                    },
                ]
            }
        }
    }

    # Выполнение поиска в Elasticsearch
    response = es.search(index=INDEX_NAME, body=query_, size=2)
    logger.info(f"Number of hits: {response['hits']['total']['value']}")

    # Вывод результата поиска
    for hit in response['hits']['hits']:
        logger.info(hit['_source'])


if __name__ == '__main__':
    path = '/mnt/ntr_work/project/nmd800/data/group_card'
    create_index_elastic_group(path)