136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
from elasticsearch import Elasticsearch
|
|
from typing import List, Tuple, Dict, Any
|
|
from app.config import settings
|
|
|
|
class FulltextSearch:
|
|
def __init__(self):
|
|
self.es = Elasticsearch([f"http://{settings.ELASTICSEARCH_HOST}:{settings.ELASTICSEARCH_PORT}"])
|
|
self.index_name = "documents"
|
|
self._create_index()
|
|
|
|
def _create_index(self):
|
|
if not self.es.indices.exists(index=self.index_name):
|
|
self.es.indices.create(
|
|
index=self.index_name,
|
|
body={
|
|
"settings": {
|
|
"analysis": {
|
|
"analyzer": {
|
|
"ru_en": {
|
|
"type": "custom",
|
|
"tokenizer": "standard",
|
|
"filter": ["lowercase", "ru_stop", "ru_stem"]
|
|
}
|
|
},
|
|
"filter": {
|
|
"ru_stop": {
|
|
"type": "stop",
|
|
"stopwords": "_russian_"
|
|
},
|
|
"ru_stem": {
|
|
"type": "stemmer",
|
|
"language": "russian"
|
|
}
|
|
}
|
|
}
|
|
},
|
|
"mappings": {
|
|
"properties": {
|
|
"doc_id": {"type": "keyword"},
|
|
"text": {"type": "text", "analyzer": "ru_en"},
|
|
"title": {"type": "text", "analyzer": "ru_en"},
|
|
"source": {"type": "keyword"},
|
|
"lang": {"type": "keyword"},
|
|
"url": {"type": "keyword"}
|
|
}
|
|
}
|
|
}
|
|
)
|
|
|
|
def index_document(self, doc_id: str, title: str, text: str, source: str, lang: str, url: str = ""):
|
|
self.es.index(
|
|
index=self.index_name,
|
|
id=doc_id,
|
|
body={
|
|
"doc_id": doc_id,
|
|
"title": title,
|
|
"text": text,
|
|
"source": source,
|
|
"lang": lang,
|
|
"url": url
|
|
}
|
|
)
|
|
|
|
def search(self, query: str, top_k: int = 10, lang: str = None) -> List[Dict[str, Any]]:
|
|
must = {
|
|
"multi_match": {
|
|
"query": query,
|
|
"fields": ["text^2", "title^3"],
|
|
"type": "best_fields",
|
|
"analyzer": "ru_en",
|
|
}
|
|
}
|
|
|
|
if lang:
|
|
search_body = {
|
|
"size": top_k,
|
|
"query": {
|
|
"bool": {
|
|
"must": must,
|
|
"filter": {"term": {"lang": lang}},
|
|
}
|
|
},
|
|
}
|
|
else:
|
|
search_body = {"size": top_k, "query": must}
|
|
|
|
results = self.es.search(index=self.index_name, body=search_body)
|
|
|
|
formatted_results = []
|
|
for hit in results['hits']['hits']:
|
|
source = hit['_source']
|
|
formatted_results.append({
|
|
"doc_id": source['doc_id'],
|
|
"title": source['title'],
|
|
"source": source['source'],
|
|
"url": source.get('url', ''),
|
|
"score": hit['_score']
|
|
})
|
|
|
|
return formatted_results
|
|
|
|
def search_phrase(self, phrase: str, top_k: int = 10) -> List[Dict[str, Any]]:
|
|
results = self.es.search(
|
|
index=self.index_name,
|
|
body={
|
|
"size": top_k,
|
|
"query": {
|
|
"match_phrase": {
|
|
"text": phrase
|
|
}
|
|
}
|
|
}
|
|
)
|
|
|
|
formatted_results = []
|
|
for hit in results['hits']['hits']:
|
|
source = hit['_source']
|
|
formatted_results.append({
|
|
"doc_id": source['doc_id'],
|
|
"title": source['title'],
|
|
"source": source['source'],
|
|
"url": source.get('url', ''),
|
|
"score": hit['_score']
|
|
})
|
|
|
|
return formatted_results
|
|
|
|
def bulk_index(self, documents: List[Dict[str, Any]]):
|
|
actions = []
|
|
for doc in documents:
|
|
actions.append({"index": {"_index": self.index_name, "_id": doc['doc_id']}})
|
|
actions.append(doc)
|
|
|
|
from elasticsearch.helpers import bulk
|
|
bulk(self.es, actions)
|