RAG数据库生成

rag_indexer/store/factory.py

@@ -0,0 +1,73 @@
+"""
+文档存储工厂 - 创建不同类型的存储实例。
+
+提供统一的接口来创建本地文件存储或 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")