es_client.py 1.54 KB
"""Elasticsearch向量数据库客户端实现."""

from typing import Any, Optional

from config import settings
from db.base import VectorDBClient


class ElasticsearchClient(VectorDBClient):
    """Elasticsearch向量数据库客户端.

    使用Elasticsearch的dense_vector类型存储向量,支持向量相似度搜索。
    """

    def __init__(self, es_url: Optional[str] = None):
        """初始化ES客户端.

        Args:
            es_url: Elasticsearch连接URL,默认从配置读取
        """
        self.es_url = es_url
        self._client: Optional[Any] = None

    async def connect(self) -> None:
        """建立ES连接."""
        from elasticsearch import AsyncElasticsearch

        self._client = AsyncElasticsearch([self.es_url])
        # 验证连接
        await self._client.ping()

    async def disconnect(self) -> None:
        """断开ES连接."""
        if self._client:
            await self._client.close()
            self._client = None

    async def health_check(self) -> bool:
        """检查ES连接健康状态."""
        if not self._client:
            return False
        try:
            return await self._client.ping()
        except Exception:
            return False
_es_client: Optional[ElasticsearchClient] = None


def get_es_client() -> ElasticsearchClient:
    """获取全局ES客户端实例(单例模式).

    Returns:
        ElasticsearchClient: ES客户端实例
    """
    global _es_client
    if _es_client is None:
        _es_client = ElasticsearchClient()
    return _es_client