dsProject/dsRag/Test/T5_SentenceSave.py

import logging
from datetime import datetime
from elasticsearch import Elasticsearch

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('../sentence_save.log'),
        logging.StreamHandler()
    ]
)

from Config.Config import ES_CONFIG
from Test.T2_Txt2Vec import text_to_embedding

# 初始化ES连接
es = Elasticsearch(
    hosts=ES_CONFIG["hosts"],
    basic_auth=ES_CONFIG["basic_auth"],
    verify_certs=ES_CONFIG["verify_certs"],
    ssl_show_warn=ES_CONFIG["ssl_show_warn"]
)

def check_text_exists(text, index_name="raw_texts"):
    """检查文本是否已存在"""
    try:
        # 添加长度验证
        if len(text) > 256:
            text = text[:256]
            
        query = {
            "query": {
                "bool": {
                    "must": [
                        {"term": {"text.keyword": text}},
                        {"match": {"text": text}}
                    ]
                }
            }
        }
        res = es.search(index=index_name, body=query, size=1)
        return res['hits']['hits'][0]['_id'] if res['hits']['hits'] else None
    except Exception as e:
        print(f"检查失败: {str(e)}")
        return None

def save_vector(text, index_name="knowledge_base"):
    """将文本向量化后保存到ES"""
    try:
        logging.info(f"开始向量化文本: {text[:50]}...")
        vector = text_to_embedding(text)
        
        doc = {
            "text": text,
            "vector": vector,
            "timestamp": datetime.now()
        }
        
        res = es.index(index=index_name, document=doc)
        logging.info(f"成功保存向量化文档到{index_name}, ID: {res['_id']}, 文本长度: {len(text)}")
        return res['_id']
    except Exception as e:
        logging.error(f"向量保存失败: {str(e)}", exc_info=True)
        raise

def save_raw_text(text, index_name="raw_texts"):
    """保存原始文本到ES"""
    try:
        # 准备文档
        doc = {
            "text": text,
            "timestamp": datetime.now()
        }
        
        # 保存到ES
        res = es.index(index=index_name, document=doc)
        print(f"原始文本已保存，ID: {res['_id']}")
        return res['_id']
    except Exception as e:
        print(f"保存失败: {str(e)}")
        raise

def save_text_with_check(text):
    """带检查的保存流程"""
    # 检查文本是否已存在
    existing_id = check_text_exists(text)
    if existing_id:
        print(f"文本已存在，ID: {existing_id}")
        return existing_id
    
    # 保存向量
    vector_id = save_vector(text)
    
    # 保存原始文本
    raw_id = save_raw_text(text)
    
    return {"vector_id": vector_id, "raw_id": raw_id}

# 使用示例
if __name__ == "__main__":
    # 示例文本
    sample_text = "如何更换支付宝绑定银行卡"
    
    # 带检查的保存
    save_text_with_check(sample_text)