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