91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
"""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
|
||
|
||
|
||
def company_registration_clauses(language: str, query: str) -> list[dict]:
|
||
patterns = {
|
||
"ru": (r"\bкак\s+откры\w*\s+осоо\b", r"\b(порядок|процедура)\s+откры\w*\s+осоо\b", r"\bкак\s+зарегистр\w*\s+осоо\b"),
|
||
"ky": (r"\bжчк\s+ач\w*\s+тартиби\b", r"\bжчк\s+кантип\s+ач\w*\b"),
|
||
}[language]
|
||
if not any(re.search(pattern, query.casefold()) 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("230044970", 2000), clause("667", 1000)]
|
||
|
||
|
||
def build_search_body(language: str, query: str, top_k: int) -> bytes:
|
||
full_text = {
|
||
"multi_match": {
|
||
"query": query,
|
||
"fields": [f"document_name_{language}", f"text_{language}"],
|
||
"type": "cross_fields",
|
||
}
|
||
}
|
||
clauses = company_registration_clauses(language, query)
|
||
bool_query = {"filter": {"term": {"language": language}}}
|
||
if clauses:
|
||
bool_query.update({"should": [full_text, *clauses], "minimum_should_match": 1})
|
||
else:
|
||
bool_query["must"] = full_text
|
||
return json.dumps({
|
||
"size": top_k,
|
||
"track_total_hits": False,
|
||
"_source": ["document_code"],
|
||
"query": {"bool": bool_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())
|