File size: 2,756 Bytes
16959be
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import copy

import json
import requests
from requests.auth import HTTPBasicAuth
from models.law_component import LawComponent

base_query = {
    "query": {
        "bool": {
            "should": [
                {
                    "match": {
                        "text": {
                            "query": None,
                            "boost": 1.0
                        }
                    }
                },
                {
                    "match": {
                        "chapterTitle": {
                            "query": None,
                            "boost": 1.0
                        }
                    }
                },
                {
                    "match_phrase": {
                        "text": {
                            "query": None,
                            "boost": 1.0
                        }
                    }
                },
                {
                    "match_phrase": {
                        "chapterTitle": {
                            "query": None,
                            "boost": 1.0
                        }
                    }
                }
            ],
            "minimum_should_match": 1
        }
    }
}

class ESRetriever:

    def __init__(self, es_host, es_index_name, es_username="", es_password=""):
        self.es_host = es_host
        self.es_index_name = es_index_name
        self.es_username = es_username
        self.es_password = es_password

        if (es_username != "" and es_password != ""):
            self.auth = HTTPBasicAuth(es_username, es_password)
        else:
            self.auth = None

    # Returns LawComponent
    def retrieve(self, query_text: str):
        query = copy.deepcopy(base_query)
        query['query']['bool']['should'][0]['match']['text']['query'] = query_text
        query['query']['bool']['should'][1]['match']['chapterTitle']['query'] = query_text
        query['query']['bool']['should'][2]['match_phrase']['text']['query'] = query_text
        query['query']['bool']['should'][3]['match_phrase']['chapterTitle']['query'] = query_text

        # try:
        response = requests.get(
            self.es_host + self.es_index_name + '/_search',
            headers={'Content-Type': 'application/json'},
            data=json.dumps(query),
            auth=self.auth
        )

        if response.ok:
            results = response.json()["hits"]["hits"]
            retrieval_results = []
            for result in results:
                lc = LawComponent.from_uri(result["_source"]["uri"])
                lc.set_text(result["_source"]["text"])
                retrieval_results.append(lc)
            return retrieval_results

        #
        # response.content