test build

backend/app/core/fulltext.py

@@ -0,0 +1,135 @@
+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)