Files
ailine/rag_indexer/config.py
root 8db63e7a8d 重构:添加模型服务模块,支持嵌入和重排服务的自动降级
新增功能:
- 创建 app/model_services 模块,提供统一的模型服务获取接口
- 实现 BaseServiceProvider 基类和 FallbackServiceChain 降级链
- 实现 get_embedding_service():优先本地 llama.cpp,降级到智谱 API
- 实现 get_rerank_service():优先本地 llama.cpp,降级到智谱 API
- 支持单例管理,确保全局只有一个服务实例

修改内容:
- 更新 app/config.py,添加智谱 API 相关配置
- 修改 rag_core/vector_store.py:支持接受外部传入的 embeddings
- 修改 rag_core/retriever_factory.py:支持接受外部传入的 embeddings
- 修改 app/agent/rag_initializer.py:使用 get_embedding_service()
- 修改 app/rag/pipeline.py:使用 get_rerank_service()
- 修改 app/memory/mem0_client.py:智能判断可用服务配置 mem0
- 修改 rag_indexer/index_builder.py:支持使用新服务,保持向后兼容
- 修改 rag_indexer/config.py:添加智谱配置

环境变量:
- ZHIPUAI_API_KEY:智谱 API 密钥(必选)
- ZHIPU_EMBEDDING_MODEL:可选,默认 embedding-3
- ZHIPU_RERANK_MODEL:可选,默认 rerank-2
- ZHIPU_API_BASE:可选,默认 https://open.bigmodel.cn/api/paas/v4
2026-04-24 22:52:36 +08:00

85 lines
2.8 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
RAG Indexer 配置管理模块
集中管理所有环境变量配置项,避免散落在各个文件中
所有配置直接从环境变量读取,无默认值,避免配置混乱
需要类型转换的配置在此处理
"""
import os
from dotenv import load_dotenv
load_dotenv()
# ========== 辅助函数:类型转换 ==========
def _get_str(key: str) -> str | None:
"""获取字符串配置"""
return os.getenv(key)
def _get_int(key: str) -> int | None:
"""获取整数配置,自动转换"""
value = os.getenv(key)
if value is not None:
try:
return int(value)
except (ValueError, TypeError):
pass
return None
def _get_list_str(key: str, default: list[str] | None = None) -> list[str]:
"""获取字符串列表配置,从逗号分隔的字符串解析"""
value = os.getenv(key)
if value is not None:
return [item.strip() for item in value.split(",") if item.strip()]
return default or []
# ========== 第三方 API 密钥 ==========
ZHIPUAI_API_KEY = _get_str("ZHIPUAI_API_KEY")
# ========== 智谱 API 配置 ==========
# 嵌入模型:根据 https://docs.bigmodel.cn/cn/guide/start/model-overview
# 可选embedding-2、embedding-3
ZHIPU_EMBEDDING_MODEL = _get_str("ZHIPU_EMBEDDING_MODEL") or "embedding-3"
ZHIPU_API_BASE = _get_str("ZHIPU_API_BASE") or "https://open.bigmodel.cn/api/paas/v4"
# ========== 向量数据库配置URL + API 密钥 配对) ==========
QDRANT_URL = _get_str("QDRANT_URL")
QDRANT_API_KEY = _get_str("QDRANT_API_KEY")
# ========== 嵌入服务配置URL + API 密钥 配对) ==========
LLAMACPP_EMBEDDING_URL = _get_str("LLAMACPP_EMBEDDING_URL")
LLAMACPP_API_KEY = _get_str("LLAMACPP_API_KEY")
# ========== 文档存储配置(分离配置 + 完整 URI ==========
# 分离配置(优先使用)
DB_HOST = _get_str("DB_HOST")
DB_PORT = _get_int("DB_PORT")
DB_USER = _get_str("DB_USER")
DB_PASSWORD = _get_str("DB_PASSWORD")
DB_NAME = _get_str("DB_NAME")
# 完整连接字符串(直接从环境变量读取)
DB_URI = _get_str("DB_URI")
if not DB_URI and all([DB_HOST, DB_PORT, DB_USER, DB_PASSWORD, DB_NAME]):
DB_URI = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}?sslmode=disable"
# 文档存储 URI直接从环境变量读取默认同 DB_URI
DOCSTORE_URI = _get_str("DOCSTORE_URI") or DB_URI
# ========== 文档加载器配置unstructured 库) ==========
# OCR 语言列表(逗号分隔,如 "chi_sim,eng"
RAG_OCR_LANGUAGES = _get_list_str("RAG_OCR_LANGUAGES", ["chi_sim", "eng"])
# 文档主语言列表(逗号分隔,如 "zh"
RAG_DOC_LANGUAGES = _get_list_str("RAG_DOC_LANGUAGES", ["zh"])
# ========== 索引器专用配置 ==========
# 默认索引存储路径
INDEX_STORAGE_PATH = _get_str("INDEX_STORAGE_PATH")