refactor: 重构RAG核心组件,简化代码结构和测试文件
Some checks failed
构建并部署 AI Agent 服务 / deploy (push) Failing after 6m53s
Some checks failed
构建并部署 AI Agent 服务 / deploy (push) Failing after 6m53s
This commit is contained in:
@@ -1,7 +1,12 @@
|
||||
# rag_core/client.py
|
||||
import os
|
||||
from .config import QDRANT_URL, QDRANT_API_KEY
|
||||
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:
|
||||
@@ -54,3 +59,47 @@ def create_async_qdrant_client(timeout: int = 300) -> AsyncQdrantClient:
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user