from elasticsearch import Elasticsearch from libs.text_processor import TextProcessor from typing import List, Dict, AsyncGenerator import hashlib import os from dotenv import load_dotenv load_dotenv() ELASTICSEARCH_USER = os.getenv("ELASTICSEARCH_USER") ELASTICSEARCH_PWD = os.getenv("ELASTICSEARCH_PWD") TITLE_INDEX = os.getenv("TITLE_INDEX") CHUNC_INDEX = os.getenv("CHUNC_INDEX") ELASTICSEARCH_HOST = os.getenv("ELASTICSEARCH_HOST") class ElasticsearchOperations: def __init__(self): self.processor = TextProcessor() self.es = Elasticsearch(hosts=[ELASTICSEARCH_HOST], verify_certs=False, http_auth=(ELASTICSEARCH_USER, ELASTICSEARCH_PWD)) self.index_name = "text_chunks" def add_document(self, index: str, doc_id:str, doc: any): self.es.index(index=index, id=doc_id, body=doc) def update_document(self, index: str, doc_id:str, doc: any): self.es.update(index=index, id=doc_id, body=doc) def get_document(self, index:str, doc_id:str): return self.es.get(index=index, id=doc_id)["_source"] def refresh_index(self, index: str): self.es.indices.refresh(index=index) def del_document(self, text: str, doc_id: str): chunks = self.processor.chunk_text(text) for i, chunk in enumerate(chunks): id = f"{doc_id}_chunk_{i}" query = { "query": { "match": { "id": id } } } response = self.es.search(index=self.index_name, body=query) print(response) if response.get('hits', {}).get('total', {}).get('value', 0) == 0: print(">>> document not found") continue else: print(">>> document was found") response = self.es.delete(index=self.index_name,id=id) print(">>> delete document",response) print(">>> finished delete document ") def index_document(self, text: str, title: str) ->None: md = hashlib.md5(title.encode()) doc_id = md.hexdigest() doc = { 'text': text, 'title': title } response = self.es.index(index='ins_expert_docs', id=doc_id, document=doc) print(response['result']) def get_doc_id(self, title:str) -> str: md = hashlib.md5(title.encode()) doc_id = md.hexdigest() return doc_id def delete_index(self, index: str): self.es.indices.delete(index=index, ignore=[400, 404]) def ingest_document(self, text: str, title: str) -> None: """将文本分块并存入Elasticsearch""" doc_id = self.get_doc_id(title) chunks = self.processor.chunk_text(text) documents = [] for i, chunk in enumerate(chunks): doc = { 'id': f"{doc_id}_chunk_{i}", 'text': f"{title}\n{title}\n{title}\n"+chunk, 'embedding': self.processor.generate_embeddings([chunk])[0].tolist() } print(f">>> insert id :{doc['id']}") documents.append(doc) if documents: body = [] for doc in documents: body.append({"index":{}}) body.append(doc) self.es.bulk(index=self.index_name, body=body) def ingest_text(self, text: str, doc_id: str) -> None: """将文本分块并存入Elasticsearch""" self.ingest_document(text, doc_id) def search_chunc_index(self, index:str, title: str, top_k:int = 3): query_embedding = self.processor.generate_embeddings([title])[0].tolist() script_query = { "script_score": { "query": {"match_all": {}}, "script": { "source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0", "params": {"query_vector": query_embedding} } } } response = self.es.search( index=index, body={ "size": top_k, "query": script_query, "_source": ["text"] } ) return [ { "id": hit["_id"], "title": hit["_source"]["title"], "text": hit["_source"]["text"], "score": hit["_score"] } for hit in response["hits"]["hits"] ] def search_title_index(self, index:str, title: str, top_k:int = 3): query_embedding = self.processor.generate_embeddings([title])[0].tolist() script_query = { "script_score": { "query": {"match_all": {}}, "script": { "source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0", "params": {"query_vector": query_embedding} } } } response = self.es.search( index=index, body={ "size": top_k, "query": script_query, "_source": ["text","title"] } ) return [ { "id": hit["_id"], "title": hit["_source"]["title"], "text": hit["_source"]["text"], "score": hit["_score"] } for hit in response["hits"]["hits"] ] def search_word_index(self, index:str, words: List[str], top_k: int = 1): search_request = [] for w in words: search_request.append( {"index": index}) search_request.append( { "query": { "match": { "word": w } }, "size": 5 } ) #response = self.es.search(index=index, # body={"size": top_k, "query": query, "_source":["articles","word"]}) response = self.es.msearch(body=search_request) result_list = [] for i, result in enumerate(response["responses"]): result_tmp = [ { "id": hit["_id"], "articles": hit["_source"]["articles"], "word": hit["_source"]["word"], "score": hit["_score"] } for hit in result["hits"]["hits"]] result_list = result_list + result_tmp return result_list def search_similar_texts(self, query: str, top_k: int = 5) -> List[Dict]: """查询与输入文本最相似的文本块""" query_embedding = self.processor.generate_embeddings([query])[0].tolist() script_query = { "script_score": { "query": {"match_all": {}}, "script": { "source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0", "params": {"query_vector": query_embedding} } } } response = self.es.search( index=self.index_name, body={ "size": top_k, "query": script_query, "_source": ["text"] } ) return [ { "id": hit["_id"], "title": hit["_source"]["text"].split("\n")[0], "text": hit["_source"]["text"], "score": hit["_score"] } for hit in response["hits"]["hits"] ] def update_document(self, datas: list): """ 更新文档 """ try: failed_list = [] for data in datas: # 更新文档（实际上是重新索引） response = self.index(index=self.index_name, id=data['doc_id'], body=data['doc']) if response.meta.status != 200: failed_list.append(data) # 刷新索引以立即应用更改（可选） self.es.indices.refresh(index=self.index_name) print(f"Elasticsearch update_doc finished! failed_list -> {failed_list}") except Exception as e: print(f"Elasticsearch update_doc error: {e}") def get_doc(self, doc_id): """ 获取文档数据 """ try: doc = self.es.get(index=self.index_name, id=doc_id)['_source'] return doc except Exception as e: print(f"Error retrieving document {doc_id}: {e}") return None