100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
"""
|
||
Mem0 记忆层初始化与健康检查。
|
||
启动时验证 Ollama 可达和模型可用，失败则抛出 RuntimeError 阻止后端启动。
|
||
"""
|
||
|
||
import logging
|
||
import httpx
|
||
from mem0 import Memory
|
||
|
||
from app.config import (
|
||
get_deepseek_api_key,
|
||
get_ollama_base_url,
|
||
get_mem0_embedding_model,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_memory: Memory | None = None
|
||
|
||
|
||
def _check_ollama_health() -> None:
|
||
"""检查 Ollama 服务是否可达，以及 embedding 模型是否已安装。"""
|
||
base_url = get_ollama_base_url()
|
||
model_name = get_mem0_embedding_model()
|
||
|
||
try:
|
||
resp = httpx.get(f"{base_url}/api/tags", timeout=5)
|
||
resp.raise_for_status()
|
||
except (httpx.ConnectError, httpx.TimeoutException, httpx.HTTPStatusError) as e:
|
||
raise RuntimeError(
|
||
f"无法连接 Ollama 服务（{base_url}）。"
|
||
f"请确认 Ollama 已启动：启动方式参见 https://ollama.com\n"
|
||
f"原始错误：{e}"
|
||
) from e
|
||
|
||
models = resp.json().get("models", [])
|
||
installed = [m.get("name", "").split(":")[0] for m in models]
|
||
if model_name not in installed:
|
||
raise RuntimeError(
|
||
f"Ollama 已运行，但未找到 embedding 模型 '{model_name}'。\n"
|
||
f"已安装的模型：{installed}\n"
|
||
f"请运行：ollama pull {model_name}"
|
||
)
|
||
|
||
logger.info("Ollama 健康检查通过：%s 模型可用", model_name)
|
||
|
||
|
||
def init_memory() -> Memory:
|
||
"""初始化 Mem0 Memory 单例。首次调用时执行健康检查。"""
|
||
global _memory
|
||
if _memory is not None:
|
||
return _memory
|
||
|
||
_check_ollama_health()
|
||
|
||
deepseek_key = get_deepseek_api_key()
|
||
if not deepseek_key:
|
||
raise RuntimeError(
|
||
"DEEPSEEK_API_KEY 未配置。Mem0 需要该 key 进行事实提取。\n"
|
||
"请在 .env 中设置 DEEPSEEK_API_KEY"
|
||
)
|
||
|
||
config = {
|
||
"llm": {
|
||
"provider": "deepseek",
|
||
"config": {
|
||
"model": "deepseek-chat",
|
||
"temperature": 0.1,
|
||
"max_tokens": 1500,
|
||
"api_key": deepseek_key,
|
||
},
|
||
},
|
||
"embedder": {
|
||
"provider": "ollama",
|
||
"config": {
|
||
"model": get_mem0_embedding_model(),
|
||
"ollama_base_url": get_ollama_base_url(),
|
||
},
|
||
},
|
||
"vector_store": {
|
||
"provider": "qdrant",
|
||
"config": {
|
||
"collection_name": "epeekit_memories",
|
||
"embedding_model_dims": 768,
|
||
"path": "./data/qdrant",
|
||
},
|
||
},
|
||
}
|
||
|
||
_memory = Memory.from_config(config)
|
||
logger.info("Mem0 记忆层初始化成功")
|
||
return _memory
|
||
|
||
|
||
def get_memory() -> Memory:
|
||
"""获取已初始化的 Memory 实例。未初始化时自动调用 init_memory()。"""
|
||
if _memory is None:
|
||
return init_memory()
|
||
return _memory
|