""" 文档存储工厂 - 创建不同类型的存储实例。 提供统一的接口来创建本地文件存储或 PostgreSQL 存储。 """ import os import logging from typing import Optional, Tuple from langchain_core.stores import BaseStore from .postgres import PostgresDocStore 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")