ailine/rag_core/store/factory.py

"""
文档存储工厂 - 创建不同类型的存储实例。

提供统一的接口来创建本地文件存储或 PostgreSQL 存储。
"""

import os
import logging
from typing import Optional, Tuple

from langchain_core.stores import BaseStore
from rag_core.store.postgres import PostgresDocStore
from dotenv import load_dotenv
load_dotenv()

logger = logging.getLogger(__name__)        

# 默认连接字符串（从环境变量读取）
DEFAULT_DB_URI = os.getenv(
    "DB_URI",
    "postgresql://postgres:huang1998@ai-postgres:5432/langgraph_db?sslmode=disable"
)


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


def create_docstore(
    store_type: str = "postgres",
    connection_string: Optional[str] = None,
    table_name: str = "parent_documents",
    pool_config: Optional[dict] = None,
    max_concurrency: Optional[int] = None
) -> Tuple[BaseStore, Optional[str]]:
    """
    工厂函数，创建 PostgreSQL 文档存储。
    
    Args:
        store_type: 存储类型，目前仅支持 "postgres"（默认）
        connection_string: PostgreSQL 连接字符串
        table_name: PostgreSQL 表名（默认：parent_documents）
        pool_config: 连接池配置
        max_concurrency: 最大并发操作数，如果为 None 则不限制
    
    Returns:
        元组 (存储实例, 连接字符串)
    
    Raises:
        ValueError: 不支持的存储类型
        ImportError: 缺少必要的依赖
    
    Example:
        >>> # 创建 PostgreSQL 存储
        >>> store, conn = create_docstore(
        ...     connection_string="postgresql://user:pass@host:5432/db",
        ...     table_name="parent_docs",
        ...     max_concurrency=10
        ... )
    """
    store_type = store_type.lower()
    
    if store_type == "postgres":
        conn_str = connection_string or get_docstore_uri()
        store = PostgresDocStore(
            connection_string=conn_str,
            table_name=table_name,
            pool_config=pool_config,
            max_concurrency=max_concurrency
        )
        return store, conn_str
    
    else:
        raise ValueError(f"不支持的存储类型: {store_type}。目前仅支持: postgres")