ailine/backend/rag_core/client.py

# rag_core/client.py
import os
from .config import QDRANT_URL, QDRANT_API_KEY, DOCSTORE_URI
from qdrant_client import QdrantClient, AsyncQdrantClient
from typing import Tuple
from langchain_core.stores import BaseStore
import logging

logger = logging.getLogger(__name__)


def create_qdrant_client(timeout: int = 300) -> QdrantClient:
    """
    创建并返回一个配置好的 Qdrant 客户端。

    Args:
        timeout: 请求超时时间（秒），默认 300 秒（索引构建需要较长超时）。

    Returns:
        配置好的 QdrantClient 实例。

    Raises:
        ValueError: 如果 QDRANT_URL 未配置。
    """
    if not QDRANT_URL:
        raise ValueError("Qdrant URL 未配置")

    client_kwargs = {
        "url": QDRANT_URL,
        "timeout": timeout,
    }
    if QDRANT_API_KEY:
        client_kwargs["api_key"] = QDRANT_API_KEY

    return QdrantClient(**client_kwargs)


def create_async_qdrant_client(timeout: int = 300) -> AsyncQdrantClient:
    """
    创建并返回一个配置好的 Qdrant 异步客户端。

    Args:
        timeout: 请求超时时间（秒），默认 300 秒。

    Returns:
        配置好的 AsyncQdrantClient 实例。

    Raises:
        ValueError: 如果 QDRANT_URL 未配置。
    """
    if not QDRANT_URL:
        raise ValueError("Qdrant URL 未配置")

    client_kwargs = {
        "url": QDRANT_URL,
        "timeout": timeout,
    }
    if QDRANT_API_KEY:
        client_kwargs["api_key"] = QDRANT_API_KEY

    return AsyncQdrantClient(**client_kwargs)


def get_docstore_uri() -> str:
    """获取 docstore 专用的数据库连接字符串（可与主库相同）"""
    return DOCSTORE_URI


def create_docstore(
    table_name: str = "parent_documents",
    pool_config: dict | None = None,
    max_concurrency: int | None = None
) -> Tuple[BaseStore, str]:
    """
    工厂函数，创建 PostgreSQL 文档存储。

    Args:
        table_name: PostgreSQL 表名（默认：parent_documents）
        pool_config: 连接池配置
        max_concurrency: 最大并发操作数，如果为 None 则不限制

    Returns:
        元组 (存储实例, 连接字符串)

    Raises:
        ImportError: 缺少必要的依赖

    Example:
        >>> # 创建 PostgreSQL 存储
        >>> store, conn = create_docstore(
        ...     table_name="parent_docs",
        ...     max_concurrency=10
        ... )
    """
    from .doc_store import PostgresDocStore
    
    conn_str = get_docstore_uri()
    store = PostgresDocStore(
        connection_string=conn_str,
        table_name=table_name,
        pool_config=pool_config,
        max_concurrency=max_concurrency
    )
    logger.info(f"PostgreSQL docstore 已创建: {table_name}")
    return store, conn_str