79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
|
"""Vector backend configuration.
|
||
|
|
|
||
|
|
Default is ``disabled`` for lightweight local startup. Install optional local
|
||
|
|
RAG dependencies and set RAG_VECTOR_BACKEND=chroma to restore local vectors.
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import importlib.util
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
from ..config import DATA_DIR
|
||
|
|
|
||
|
|
|
||
|
|
def settings():
|
||
|
|
path = DATA_DIR / 'rag-backend.json'
|
||
|
|
try:
|
||
|
|
config = json.loads(path.read_text(encoding='utf-8'))
|
||
|
|
if not isinstance(config, dict):
|
||
|
|
raise ValueError('知识索引配置必须是对象')
|
||
|
|
except FileNotFoundError:
|
||
|
|
config = {}
|
||
|
|
value = os.environ.get('RAG_VECTOR_BACKEND', config.get('backend', 'disabled'))
|
||
|
|
backend_name = str(value or 'disabled').strip().lower()
|
||
|
|
if backend_name not in {'disabled', 'chroma', 'qdrant'}:
|
||
|
|
raise ValueError('知识索引后端配置无效')
|
||
|
|
return {**config, 'backend': backend_name}
|
||
|
|
|
||
|
|
|
||
|
|
def configured_backend() -> str:
|
||
|
|
return settings()['backend']
|
||
|
|
|
||
|
|
|
||
|
|
def local_chroma_available() -> bool:
|
||
|
|
return importlib.util.find_spec('chromadb') is not None
|
||
|
|
|
||
|
|
|
||
|
|
def effective_backend() -> str:
|
||
|
|
"""Return the backend that can actually run in this installation."""
|
||
|
|
value = configured_backend()
|
||
|
|
if value == 'chroma' and not local_chroma_available():
|
||
|
|
return 'disabled'
|
||
|
|
return value
|
||
|
|
|
||
|
|
|
||
|
|
def is_disabled() -> bool:
|
||
|
|
return effective_backend() == 'disabled'
|
||
|
|
|
||
|
|
|
||
|
|
def is_qdrant():
|
||
|
|
return effective_backend() == 'qdrant'
|
||
|
|
|
||
|
|
|
||
|
|
def qdrant_settings():
|
||
|
|
config = settings()
|
||
|
|
url = str(os.environ.get('QDRANT_URL', config.get('url', ''))).rstrip('/')
|
||
|
|
if not url.startswith(('http://', 'https://')):
|
||
|
|
raise ValueError('请配置 Qdrant 服务地址')
|
||
|
|
key = os.environ.get('QDRANT_API_KEY', '')
|
||
|
|
if not key and config.get('api_key_file'):
|
||
|
|
key = Path(config['api_key_file']).read_text(encoding='utf-8').strip()
|
||
|
|
if not key:
|
||
|
|
raise ValueError('请配置 Qdrant API 密钥')
|
||
|
|
prefix = str(config.get('collection_prefix', 'yawan_'))
|
||
|
|
if not prefix or not all(c.isalnum() or c in '_-' for c in prefix):
|
||
|
|
raise ValueError('Qdrant 知识库命名空间无效')
|
||
|
|
return url, key, prefix
|
||
|
|
|
||
|
|
|
||
|
|
def namespace():
|
||
|
|
value = effective_backend()
|
||
|
|
if value == 'qdrant':
|
||
|
|
config = settings()
|
||
|
|
return json.dumps(['qdrant', str(os.environ.get('QDRANT_URL', config.get('url', ''))).rstrip('/'),
|
||
|
|
str(config.get('collection_prefix', 'yawan_'))], separators=(',', ':'))
|
||
|
|
if value == 'disabled':
|
||
|
|
return 'disabled'
|
||
|
|
return 'chroma' if local_chroma_available() else 'chroma-missing'
|