Add elasticsearch engine

Add engine/elastic/elasric.py file with ElasticEngine realization

app/models.py

@@ -4,7 +4,7 @@
 from flask_login import UserMixin
 
 from app import db, login
-from markov import NgrammTextGenerator
+from engine.markov import NgrammTextGenerator
 from .utils import set_model, get_model
 
 

app/routes.py

@@ -10,7 +10,7 @@
 from .models import (
     User, Document, MarkovModel)
 from .utils import (
-    get_model, get_text_corpus_from_file, get_text_corpus_from_postgres, get_msg_stack)
+    get_model, get_text_corpus_from_file, get_text_corpus_from_postgres, get_msg_stack, get_elastic_engine)
 from .forms import (
     LoginForm, RegistrationForm, DocumentForm, ModelForm)
 

app/utils.py

@@ -1,10 +1,18 @@
 import re
 from typing import Generator, Optional, List
 
-from markov import NgrammTextGenerator
-from markov.utils import PostgresStorage
+from config import Config
+from engine.markov import NgrammTextGenerator
+from engine.markov import PostgresStorage
+from engine.elastic import ElasticEngine
 
 __model: Optional[NgrammTextGenerator] = None
+__es: ElasticEngine = ElasticEngine(host=Config.ELASTIC_HOST)
+
+
+def get_elastic_engine() -> ElasticEngine:
+    global __es
+    return __es
 
 
 def get_model() -> NgrammTextGenerator:

config.py

@@ -14,3 +14,4 @@ class Config:
     #                           'sqlite:///' + os.path.join(BASE_DIR, 'users.db')
     SQLALCHEMY_TRACK_MODIFICATIONS = False
     MODELS_ROOT = pathlib.Path(__file__).parent / 'markov' / 'models'
+    ELASTIC_HOST = os.environ.get('ELASTIC_HOST', '0.0.0.0:9200')

engine/elastic/__init__.py

@@ -0,0 +1 @@
+from .elastic import ElasticEngine

engine/elastic/elastic.py

@@ -0,0 +1,72 @@
+from typing import List, Dict, Any
+
+from elasticsearch import Elasticsearch
+
+
+class ElasticEngine:
+    es: Elasticsearch
+    doc_type: str = "sentence"
+
+    def __init__(self, host: str, **kwargs):
+        self.es = Elasticsearch(host=host, **kwargs)
+
+    def add_index(self, name: str, number_of_shards: int = 1, number_of_replicas: int = 2):
+        self.es.indices.create(index=name, body={
+            "settings": {
+                "index": {
+                    "number_of_shards": number_of_shards,
+                    "number_of_replicas": number_of_replicas,
+                    "analysis": {
+                        "analyzer": {
+                            "t9_analyzer": {
+                                "type": "custom",
+                                "tokenizer": "standard",
+                                "filter": [
+                                    "lowercase",
+                                    "custom_edge_ngram"
+                                ]
+                            }
+                        },
+                        "filter": {
+                            "custom_edge_ngram": {
+                                "type": "edge_ngram",
+                                "min_gram": 2,
+                                "max_gram": 10
+                            }
+                        }
+                    }
+                }
+            },
+            "mappings": {
+                self.doc_type: {
+                    "properties": {
+                        "text": {
+                            "type": "text",
+                            "analyzer": "t9_analyzer",
+                            "search_analyzer": "standard"
+                        }
+                    }
+                }
+            }
+        })
+
+    def add_doc(self, index_name: str, text: str):
+        self.es.index(index=index_name, doc_type=self.doc_type, body={
+            'text': text
+        })
+
+    def get(self, index_name: str, phrase: str, count: int = 10) -> List[str]:
+        return [
+            doc['_source']['text']
+            for doc in self.es.search(
+                index=index_name,
+                doc_type=self.doc_type,
+                body={
+                    "query": {
+                        "match_phrase": {
+                            "text": phrase
+                        }
+                    },
+                    "size": count
+                })['hits']['hits']
+        ]

markov/__init__.py → engine/markov/__init__.py

@@ -1 +1,2 @@
 from .ngrams_text_generator import NgrammTextGenerator
+from .utils import PostgresStorage

markov/legacy/chain_storage.py → engine/markov/legacy/chain_storage.py

@@ -1,7 +1,7 @@
 import operator
 from typing import Iterable, Generator
 
-from markov.utils import PostgresStorage
+from engine.markov import PostgresStorage
 
 
 def accumulate(iterable: Iterable, func=operator.add) -> Generator:

markov/legacy/encoder_storage.py → engine/markov/legacy/encoder_storage.py

@@ -1,5 +1,5 @@
-from markov.utils.postgres import PostgresStorage
-from markov.utils.encoder import WordsEncoder
+from engine.markov import PostgresStorage
+from engine.markov import WordsEncoder
 
 
 class EncoderStorage(PostgresStorage):

markov/legacy/habr_storage.py → engine/markov/legacy/habr_storage.py

@@ -1,5 +1,5 @@
 from typing import Generator
-from markov.utils.postgres import PostgresStorage
+from engine.markov import PostgresStorage
 
 
 class HabrStorage(PostgresStorage):

markov/legacy/text_generator.py → engine/markov/legacy/text_generator.py

@@ -5,7 +5,7 @@
 import nltk
 
 from .chain_storage import ChainStorage
-from markov.utils import EncoderStorage, WordsEncoder, TextProcessor
+from engine.markov import EncoderStorage, WordsEncoder, TextProcessor
 
 
 class TextGenerator: