209 lines
6.6 KiB
Python
209 lines
6.6 KiB
Python
"""
|
||
Elasticsearch配置文件
|
||
包含云端ES连接配置和相关常量
|
||
"""
|
||
|
||
import os
|
||
from typing import Optional
|
||
|
||
class ElasticsearchConfig:
|
||
"""Elasticsearch配置类"""
|
||
|
||
# ES连接配置（优先从环境变量读取）
|
||
ES_HOST = os.getenv("ELASTICSEARCH_HOST", "http://101.200.154.78:9200")
|
||
ES_USERNAME = os.getenv("ELASTICSEARCH_USERNAME", "elastic")
|
||
ES_PASSWORD = os.getenv("ELASTICSEARCH_PASSWORD", "Abcd123456")
|
||
|
||
# 备用SSH隧道方式（如果端口未开放时使用）
|
||
# ES_HOST = "http://localhost:9200" # 通过SSH隧道映射到本地
|
||
|
||
# 索引配置
|
||
INDEX_PREFIX = "hipporag"
|
||
EDGES_INDEX_SUFFIX = "edges"
|
||
PASSAGES_INDEX_SUFFIX = "passages"
|
||
|
||
# 向量配置
|
||
VECTOR_DIMENSION = 1024 # text-embedding-v3的维度
|
||
SIMILARITY_METRIC = "cosine"
|
||
|
||
# 批处理配置
|
||
DEFAULT_BATCH_SIZE = 100
|
||
DEFAULT_TEXT_BATCH_SIZE = 40
|
||
DEFAULT_EDGE_BATCH_SIZE = 256
|
||
|
||
# 搜索配置
|
||
DEFAULT_TOP_K = 10
|
||
DEFAULT_NUM_CANDIDATES_MULTIPLIER = 2
|
||
|
||
@classmethod
|
||
def get_edges_index_name(cls, keyword: str) -> str:
|
||
"""获取边索引名称"""
|
||
return f"{cls.INDEX_PREFIX}_{keyword}_{cls.EDGES_INDEX_SUFFIX}"
|
||
|
||
@classmethod
|
||
def get_passages_index_name(cls, keyword: str) -> str:
|
||
"""获取文本段落索引名称"""
|
||
return f"{cls.INDEX_PREFIX}_{keyword}_{cls.PASSAGES_INDEX_SUFFIX}"
|
||
|
||
@classmethod
|
||
def get_es_config(cls) -> dict:
|
||
"""获取ES连接配置"""
|
||
return {
|
||
"hosts": [cls.ES_HOST],
|
||
"basic_auth": (cls.ES_USERNAME, cls.ES_PASSWORD),
|
||
"verify_certs": False, # 如果使用自签名证书
|
||
"ssl_show_warn": False,
|
||
"request_timeout": 300, # 请求超时时间，增加至5分钟
|
||
"retry_on_timeout": True,
|
||
"max_retries": 5, # 重试次数，增加至5次
|
||
"sniff_on_start": False, # 禁用节点嗅探，避免网络问题
|
||
"sniff_on_connection_fail": False, # 禁用连接失败时的嗅探
|
||
"sniff_timeout": 10,
|
||
"http_compress": True, # 启用HTTP压缩
|
||
"http_auth": (cls.ES_USERNAME, cls.ES_PASSWORD),
|
||
"timeout": 300, # 连接超时
|
||
"connection_class": None, # 使用默认连接类
|
||
"selector_class": None, # 使用默认选择器
|
||
"dead_timeout": 60, # 死连接超时时间
|
||
"retry_on_status": {502, 503, 504, 408, 429}, # 重试的HTTP状态码
|
||
"maxsize": 25, # 连接池大小
|
||
}
|
||
|
||
|
||
class SSHTunnelConfig:
|
||
"""SSH隧道配置类"""
|
||
|
||
# SSH连接配置（需要向管理员获取）
|
||
SSH_HOST = "101.200.154.78" # SSH服务器地址
|
||
SSH_PORT = 22 # SSH端口
|
||
SSH_USERNAME = "your_username" # 需要替换为实际用户名
|
||
SSH_KEY_PATH = "path/to/your/private_key.pem" # 需要替换为实际私钥路径
|
||
# SSH_PASSWORD = "your_password" # 如果使用密码认证
|
||
|
||
# 隧道映射配置
|
||
LOCAL_PORT = 9200 # 本地端口
|
||
REMOTE_HOST = "localhost" # 服务器内部ES地址
|
||
REMOTE_PORT = 9200 # 服务器内部ES端口
|
||
|
||
@classmethod
|
||
def get_ssh_config(cls) -> dict:
|
||
"""获取SSH隧道配置"""
|
||
config = {
|
||
'ssh_host': cls.SSH_HOST,
|
||
'ssh_port': cls.SSH_PORT,
|
||
'ssh_username': cls.SSH_USERNAME,
|
||
'local_port': cls.LOCAL_PORT,
|
||
'remote_host': cls.REMOTE_HOST,
|
||
'remote_port': cls.REMOTE_PORT
|
||
}
|
||
|
||
# 添加认证信息
|
||
if hasattr(cls, 'SSH_KEY_PATH') and cls.SSH_KEY_PATH != "path/to/your/private_key.pem":
|
||
config['ssh_key_path'] = cls.SSH_KEY_PATH
|
||
elif hasattr(cls, 'SSH_PASSWORD'):
|
||
config['ssh_password'] = cls.SSH_PASSWORD
|
||
|
||
return config
|
||
|
||
|
||
class HippoRAGConfig:
|
||
"""HippoRAG检索配置"""
|
||
|
||
# 推理配置
|
||
TOPK_EDGES = 10
|
||
TOPK_NODES = 30
|
||
WEIGHT_ADJUST = 0.05
|
||
PPR_ALPHA = 0.85
|
||
PPR_MAX_ITER = 100
|
||
PPR_TOL = 1e-6
|
||
|
||
# 检索模式
|
||
RETRIEVAL_MODE = "query2edge" # query2edge, query2node, ner2node
|
||
|
||
# OneAPI配置（从环境变量读取）
|
||
@classmethod
|
||
def get_oneapi_config(cls) -> dict:
|
||
"""获取OneAPI配置"""
|
||
return {
|
||
"api_key": os.getenv('ONEAPI_KEY'),
|
||
"base_url": os.getenv('ONEAPI_BASE_URL'),
|
||
"model_embed": os.getenv('ONEAPI_MODEL_EMBED', 'text-embedding-v3'),
|
||
"model_gen": os.getenv('ONEAPI_MODEL_GEN', 'qwen2-7b-instruct')
|
||
}
|
||
|
||
|
||
# 索引映射模板（兼容版本）
|
||
EDGES_INDEX_MAPPING = {
|
||
"mappings": {
|
||
"properties": {
|
||
"edge_index": {"type": "integer"},
|
||
"head_node_id": {"type": "keyword"},
|
||
"tail_node_id": {"type": "keyword"},
|
||
"head_entity": {"type": "keyword"},
|
||
"relation": {"type": "keyword"},
|
||
"tail_entity": {"type": "keyword"},
|
||
"triple_text": {"type": "text"},
|
||
"embedding": {
|
||
"type": "dense_vector",
|
||
"dims": 1024,
|
||
"index": True,
|
||
"similarity": "cosine"
|
||
},
|
||
"created_at": {"type": "date"},
|
||
"keyword": {"type": "keyword"}
|
||
}
|
||
},
|
||
"settings": {
|
||
"number_of_shards": 1,
|
||
"number_of_replicas": 0
|
||
}
|
||
}
|
||
|
||
PASSAGES_INDEX_MAPPING = {
|
||
"mappings": {
|
||
"properties": {
|
||
"passage_id": {"type": "keyword"},
|
||
"content": {"type": "text"},
|
||
"file_id": {"type": "keyword"},
|
||
"evidence": {"type": "keyword"},
|
||
"embedding": {
|
||
"type": "dense_vector",
|
||
"dims": 1024,
|
||
"index": True,
|
||
"similarity": "cosine"
|
||
},
|
||
"created_at": {"type": "date"},
|
||
"keyword": {"type": "keyword"}
|
||
}
|
||
},
|
||
"settings": {
|
||
"number_of_shards": 1,
|
||
"number_of_replicas": 0
|
||
}
|
||
}
|
||
|
||
# 新增：Node节点ES索引mapping
|
||
NODES_INDEX_MAPPING = {
|
||
"mappings": {
|
||
"properties": {
|
||
"node_id": {"type": "keyword"},
|
||
"name": {"type": "text"},
|
||
"type": {"type": "keyword"},
|
||
"concepts": {"type": "text"},
|
||
"synsets": {"type": "text"},
|
||
"embedding": {
|
||
"type": "dense_vector",
|
||
"dims": 1024,
|
||
"index": True,
|
||
"similarity": "cosine"
|
||
},
|
||
"created_at": {"type": "date"},
|
||
"keyword": {"type": "keyword"}
|
||
}
|
||
},
|
||
"settings": {
|
||
"number_of_shards": 1,
|
||
"number_of_replicas": 0 # 改回0，避免黄色状态
|
||
}
|
||
}
|