akyldash/backend/search/query.py

"""Run document searches against the local OpenSearch index."""

from __future__ import annotations

import argparse
import json
import re
import urllib.parse

from search.minjust_opensearch import APP_VERSION, request_json


LEGAL_FORM_ALIASES = {
    "ru": {
        "осоо": ("ооо", "общество с ограниченной ответственностью"),
        "ооо": ("осоо", "общество с ограниченной ответственностью"),
        "общество с ограниченной ответственностью": ("осоо", "ооо"),
        "общества с ограниченной ответственностью": ("осоо", "ооо"),
    },
    "ky": {
        "жчк": ("жоопкерчилиги чектелген коом",),
        "жоопкерчилиги чектелген коом": ("жчк",),
    },
}
COMPANY_REGISTRATION_DOCUMENTS = (("230044970", 3000), ("667", 2000), ("4", 1000))


def query_variants(language: str, query: str) -> list[str]:
    variants = [query]
    for alias, expansions in LEGAL_FORM_ALIASES[language].items():
        pattern = re.compile(rf"(?<!\w){re.escape(alias)}(?!\w)", re.IGNORECASE)
        if pattern.search(query):
            variants.extend(pattern.sub(expansion, query) for expansion in expansions)
    return list(dict.fromkeys(variants))


def company_registration_clauses(language: str, query: str) -> list[dict]:
    non_registration_context = {
        "ru": (r"\bсч[её]т\w*", r"\bбанк\w*", r"\bфилиал\w*", r"\bпредставительств\w*"),
        "ky": (r"\bбанк\w*", r"\bэсеп\w*", r"\bфилиал\w*", r"\bөкүлчүлүк\w*"),
    }[language]
    normalized_query = query.casefold()
    if any(re.search(pattern, normalized_query) for pattern in non_registration_context):
        return []

    forms, patterns = {
        "ru": (
            r"(?:осоо|ооо|обществ\w*\s+с\s+ограниченной\s+ответственностью)",
            (
                r"\bкак\s+(?:откры\w*|созда\w*|зарегистр\w*)\s+{form}\b",
                r"\b(?:порядок|процедура)\s+(?:открыти\w*|создани\w*|регистраци\w*)\s+{form}\b",
                r"\b(?:государственн\w+\s+)?регистраци\w*\s+{form}\b",
                r"\bсоздани\w*\s+{form}\b",
            ),
        ),
        "ky": (
            r"(?:жчк|жоопкерчилиги\s+чектелген\s+коом)",
            (
                r"\bжчк\s+ач\w*\s+тартиби\b",
                r"\bжчк\s+кантип\s+ач\w*\b",
                r"\b{form}\s+(?:ачуу|ачуунун|каттоо|каттоонун)\b",
                r"\b(?:ачуу|каттоо)\s+{form}\b",
            ),
        ),
    }[language]
    if not any(re.search(pattern.format(form=forms), normalized_query) for pattern in patterns):
        return []

    status = {"ru": "Действует", "ky": "Күчүндө"}[language]
    # ponytail: curated legal mapping; replace with a reviewed intent catalog when coverage expands.
    def clause(document_code: str, boost: int) -> dict:
        return {
            "constant_score": {
                "filter": {
                    "bool": {
                        "filter": [
                            {"term": {"document_code": document_code}},
                            {"term": {f"status_{language}": status}},
                        ]
                    }
                },
                "boost": boost,
            }
        }

    return [clause(code, boost) for code, boost in COMPANY_REGISTRATION_DOCUMENTS]


def build_search_query(language: str, query: str) -> dict:
    intent_clauses = company_registration_clauses(language, query)
    full_text = [
        {
            "multi_match": {
                "query": variant,
                "fields": [f"document_name_{language}^4", f"text_{language}"],
                "type": "cross_fields",
                "operator": "OR" if intent_clauses else "AND",
            }
        }
        for variant in query_variants(language, query)
    ]
    planned = {"should": [*full_text, *intent_clauses], "minimum_should_match": 1}
    return {"bool": planned}


def build_search_body(language: str, query: str, top_k: int) -> bytes:
    return json.dumps({
        "size": top_k,
        "track_total_hits": False,
        "_source": ["document_code"],
        "query": {"bool": {"filter": [{"term": {"language": language}}, {"term": {"is_current_edition": True}}], "must": build_search_query(language, query)}},
        "collapse": {"field": "document_code"},
    }, ensure_ascii=False).encode()


def search_documents(base_url: str, index: str, language: str, query: str, top_k: int) -> list[str]:
    url = f"{base_url.rstrip('/')}/{urllib.parse.quote(index, safe='')}/_search"
    response = request_json(url, "POST", build_search_body(language, query, top_k), "application/json")
    try:
        return [hit["_source"]["document_code"] for hit in response["hits"]["hits"]]
    except (KeyError, TypeError) as error:
        raise RuntimeError("OpenSearch search response is incomplete") from error


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("query")
    parser.add_argument("--language", choices=("ru", "ky"), required=True)
    parser.add_argument("--url", default="http://127.0.0.1:9200")
    parser.add_argument("--index", default="akyldash-fragments-v1")
    parser.add_argument("--top-k", type=int, default=10)
    parser.add_argument("--version", action="version", version=APP_VERSION)
    arguments = parser.parse_args()
    if arguments.top_k <= 0:
        raise SystemExit("--top-k must be greater than zero")
    print(json.dumps(search_documents(arguments.url, arguments.index, arguments.language, arguments.query, arguments.top_k), ensure_ascii=False))
    print(f"Akyldash Backend v{APP_VERSION} · Frontend — not created")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())