adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'feature/rag2' into release/v0.3.1

Browse Source 
* feature/rag2:
  [modify] fix
  [modify] Optimize ES connections and add rerank security checks
This commit is contained in:
Mark
2026-04-21 13:44:39 +08:00
parent 74157adb12 3b359df02f
commit aecb0f6497
4 changed files with 85 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
api/app/controllers/chunk_controller.py
View File

@@ -457,7 +457,7 @@ async def retrieve_chunks(
if doc.metadata["doc_id"] not in seen_ids:
seen_ids.add(doc.metadata["doc_id"])
unique_rs.append(doc)
rs = vector_service.rerank(query=retrieve_data.query, docs=unique_rs, top_k=retrieve_data.top_k)
rs = vector_service.rerank(query=retrieve_data.query, docs=unique_rs, top_k=retrieve_data.top_k) if unique_rs else []
if retrieve_data.retrieve_type == chunk_schema.RetrieveType.Graph:
kb_ids = [str(kb_id) for kb_id in private_kb_ids]
workspace_ids = [str(workspace_id) for workspace_id in private_workspace_ids]

2
api/app/core/rag/nlp/search.py
View File

@@ -113,7 +113,7 @@ def knowledge_retrieval(
continue
# Use the specified reranker for re-ranking
if reranker_id:
if reranker_id and all_results:
try:
all_results = rerank(db=db, reranker_id=reranker_id, query=query, docs=all_results, top_k=reranker_top_k)
except Exception as rerank_error:

4
api/app/core/rag/utils/es_conn.py
View File

@@ -68,9 +68,9 @@ class ESConnection(DocStoreConnection):
client_config = {
"hosts": [hosts],
"basic_auth": (os.getenv("ELASTICSEARCH_USERNAME", "elastic"), os.getenv("ELASTICSEARCH_PASSWORD", "elastic")),
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 100000)),
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 30)),
"retry_on_timeout": os.getenv("ELASTICSEARCH_RETRY_ON_TIMEOUT", True) == "true",
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 10000)),
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 3)),
}
# Only add SSL settings if using HTTPS

193
api/app/core/rag/vdb/elasticsearch/elasticsearch_vector.py
View File

@@ -1,25 +1,22 @@
import os
import logging
from typing import Any, cast
import threading
from typing import Any
from urllib.parse import urlparse
import uuid
import requests
from elasticsearch import Elasticsearch, helpers
from elasticsearch.helpers import BulkIndexError
from packaging.version import parse as parse_version
from pydantic import BaseModel, model_validator
from abc import ABC
# langchain-community
# langchain-xinference
# from langchain_community.embeddings import XinferenceEmbeddings
# from langchain_xinference import XinferenceRerank
from langchain_core.documents import Document
from app.core.models.base import RedBearModelConfig
from app.core.models import RedBearLLM, RedBearRerank
from app.core.models import RedBearRerank
from app.core.models.embedding import RedBearEmbeddings
from app.models.models_model import ModelConfig, ModelApiKey
from app.services.model_service import ModelConfigService
from app.models.models_model import ModelApiKey
from app.models.knowledge_model import Knowledge
from app.core.rag.vdb.field import Field
@@ -29,37 +26,9 @@ from app.core.rag.models.chunk import DocumentChunk
logger = logging.getLogger(__name__)
class ElasticSearchConfig(BaseModel):
# Regular Elasticsearch config
host: str | None = None
port: int | None = None
username: str | None = None
password: str | None = None
# Common config
ca_certs: str | None = None
verify_certs: bool = False
request_timeout: int = 100000
retry_on_timeout: bool = True
max_retries: int = 10000
@model_validator(mode="before")
@classmethod
def validate_config(cls, values: dict):
# Regular Elasticsearch validation
if not values.get("host"):
raise ValueError("config HOST is required for regular Elasticsearch")
if not values.get("port"):
raise ValueError("config PORT is required for regular Elasticsearch")
if not values.get("username"):
raise ValueError("config USERNAME is required for regular Elasticsearch")
if not values.get("password"):
raise ValueError("config PASSWORD is required for regular Elasticsearch")
return values
class ElasticSearchVector(BaseVector):
def __init__(self, index_name: str, config: ElasticSearchConfig, embedding_config: ModelApiKey, reranker_config: ModelApiKey):
def __init__(self, index_name: str, client: Elasticsearch,
embedding_config: ModelApiKey, reranker_config: ModelApiKey):
super().__init__(index_name.lower())
# 初始化 Embedding 模型（自动支持火山引擎多模态）
@@ -77,58 +46,8 @@ class ElasticSearchVector(BaseVector):
api_key=reranker_config.api_key,
base_url=reranker_config.api_base
))
self._client = self._init_client(config)
self._version = self._get_version()
self._check_version()
def _init_client(self, config: ElasticSearchConfig) -> Elasticsearch:
"""
Initialize Elasticsearch client for regular Elasticsearch.
"""
try:
# Regular Elasticsearch configuration
parsed_url = urlparse(config.host or "")
if parsed_url.scheme in {"http", "https"}:
hosts = f"{config.host}:{config.port}"
use_https = parsed_url.scheme == "https"
else:
hosts = f"https://{config.host}:{config.port}"
use_https = False
client_config = {
"hosts": [hosts],
"basic_auth": (config.username, config.password),
"request_timeout": config.request_timeout,
"retry_on_timeout": config.retry_on_timeout,
"max_retries": config.max_retries,
}
# Only add SSL settings if using HTTPS
if use_https:
client_config["verify_certs"] = config.verify_certs
if config.ca_certs:
client_config["ca_certs"] = config.ca_certs
client = Elasticsearch(**client_config)
# Test connection
if not client.ping():
raise ConnectionError("Failed to connect to Elasticsearch")
except requests.ConnectionError as e:
raise ConnectionError(f"Vector database connection error: {str(e)}")
except Exception as e:
raise ConnectionError(f"Elasticsearch client initialization failed: {str(e)}")
return client
def _get_version(self) -> str:
info = self._client.info()
return cast(str, info["version"]["number"])
def _check_version(self):
if parse_version(self._version) < parse_version("8.0.0"):
raise ValueError("Elasticsearch vector database version must be greater than 8.0.0")
# 使用外部传入的共享客户端
self._client = client
def get_type(self) -> str:
return "elasticsearch"
@@ -745,29 +664,79 @@ class ElasticSearchVector(BaseVector):
class ElasticSearchVectorFactory:
@staticmethod
def init_vector(knowledge: Knowledge) -> ElasticSearchVector:
"""ES 向量服务工厂 - 单例共享连接"""
_client: Elasticsearch | None = None
_lock = threading.Lock()
_version_checked = False
@classmethod
def _get_shared_client(cls) -> Elasticsearch:
"""获取共享的 ES 客户端（线程安全的懒加载单例）"""
if cls._client is not None:
return cls._client
with cls._lock:
# 双重检查，防止并发时重复创建
if cls._client is not None:
return cls._client
try:
parsed_url = urlparse(os.getenv("ELASTICSEARCH_HOST", "127.0.0.1") or "")
if parsed_url.scheme in {"http", "https"}:
hosts = f'{os.getenv("ELASTICSEARCH_HOST")}:{os.getenv("ELASTICSEARCH_PORT", 9200)}'
use_https = parsed_url.scheme == "https"
else:
hosts = f'https://{os.getenv("ELASTICSEARCH_HOST", "127.0.0.1")}:{os.getenv("ELASTICSEARCH_PORT", 9200)}'
use_https = False
client_config = {
"hosts": [hosts],
"basic_auth": (
os.getenv("ELASTICSEARCH_USERNAME", "elastic"),
os.getenv("ELASTICSEARCH_PASSWORD", "elastic"),
),
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 30)),
"retry_on_timeout": True,
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 3)),
"connections_per_node": int(os.getenv("ELASTICSEARCH_CONNECTIONS_PER_NODE", 10)),
}
if use_https:
client_config["verify_certs"] = os.getenv("ELASTICSEARCH_VERIFY_CERTS", "false") == "true"
ca_certs = os.getenv("ELASTICSEARCH_CA_CERTS")
if ca_certs:
client_config["ca_certs"] = str(ca_certs)
client = Elasticsearch(**client_config)
if not client.ping():
raise ConnectionError("Failed to connect to Elasticsearch")
# 版本检查只做一次
if not cls._version_checked:
info = client.info()
version = info["version"]["number"]
if parse_version(version) < parse_version("8.0.0"):
raise ValueError(f"Elasticsearch version must be >= 8.0.0, got {version}")
cls._version_checked = True
logger.info(f"Elasticsearch shared client initialized, version: {version}")
cls._client = client
except requests.ConnectionError as e:
raise ConnectionError(f"Vector database connection error: {str(e)}")
except Exception as e:
raise ConnectionError(f"Elasticsearch client initialization failed: {str(e)}")
return cls._client
@classmethod
def init_vector(cls, knowledge: Knowledge) -> ElasticSearchVector:
"""创建向量服务实例（共享 ES 连接）"""
client = cls._get_shared_client()
collection_name = f"Vector_index_{knowledge.id}_Node"
# Use regular Elasticsearch with config values
config_dict = {
"host": os.getenv("ELASTICSEARCH_HOST", "127.0.0.1"),
"port": os.getenv("ELASTICSEARCH_PORT", 9200),
"username": os.getenv("ELASTICSEARCH_USERNAME", "elastic"),
"password": os.getenv("ELASTICSEARCH_PASSWORD", "elastic"),
}
# Common configuration
config_dict.update(
{
"ca_certs": str(os.getenv("ELASTICSEARCH_CA_CERTS")) if os.getenv("ELASTICSEARCH_CA_CERTS") else None,
"verify_certs": os.getenv("ELASTICSEARCH_VERIFY_CERTS", False) == "true",
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 100000)),
"retry_on_timeout": os.getenv("ELASTICSEARCH_RETRY_ON_TIMEOUT", True) == "true",
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 10000)),
}
)
if knowledge.embedding is None:
raise ValueError(f"embedding_id config error: {str(knowledge.embedding_id)}")
if knowledge.reranker is None:
@@ -775,9 +744,9 @@ class ElasticSearchVectorFactory:
return ElasticSearchVector(
index_name=collection_name,
config=ElasticSearchConfig(**config_dict),
client=client,
embedding_config=knowledge.embedding.api_keys[0],
reranker_config=knowledge.reranker.api_keys[0]
reranker_config=knowledge.reranker.api_keys[0],
)