tomtan
269661afff
* 新增 MemoryGatewayConfig 和 MemoryConfig dataclass，用于配置管理。 * 实现 MemoryGatewayUserCredential 和 MemoryGatewayCredentialStore，用于处理用户凭据。 * 创建 MemoryGatewayService，用于管理与 Memory Gateway 的交互。 * 开发用于记忆设置的 JSON 配置文件。 * 增强单元测试，覆盖新功能，包括凭据存储和服务行为。 * 更新 entrypoint 和实例创建脚本，以初始化 Memory Gateway 用户存储。
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
"""Per-instance credential storage for Memory Gateway users."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class MemoryGatewayUserCredential:
|
|
user_id: str
|
|
user_key: str = field(repr=False)
|
|
|
|
|
|
class MemoryGatewayCredentialStore:
|
|
"""Persist Beaver username -> Gateway credential mappings."""
|
|
|
|
def __init__(self, path: str | Path) -> None:
|
|
self.path = Path(path)
|
|
|
|
def get(self, username: str) -> MemoryGatewayUserCredential | None:
|
|
users = self._load_users()
|
|
payload = users.get(username)
|
|
if not isinstance(payload, dict):
|
|
return None
|
|
user_id = str(payload.get("userId") or "").strip()
|
|
user_key = str(payload.get("userKey") or "").strip()
|
|
if not user_id or not user_key:
|
|
return None
|
|
return MemoryGatewayUserCredential(user_id=user_id, user_key=user_key)
|
|
|
|
def save(self, username: str, credential: MemoryGatewayUserCredential) -> None:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
users = self._load_users()
|
|
users[username] = {
|
|
"userId": credential.user_id,
|
|
"userKey": credential.user_key,
|
|
}
|
|
payload = {"users": dict(sorted(users.items()))}
|
|
fd, tmp_name = tempfile.mkstemp(
|
|
prefix=f".{self.path.name}.",
|
|
suffix=".tmp",
|
|
dir=str(self.path.parent),
|
|
)
|
|
tmp_path = Path(tmp_name)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
|
json.dump(payload, handle, ensure_ascii=False, indent=2)
|
|
handle.write("\n")
|
|
os.chmod(tmp_path, 0o600)
|
|
os.replace(tmp_path, self.path)
|
|
os.chmod(self.path, 0o600)
|
|
finally:
|
|
if tmp_path.exists():
|
|
tmp_path.unlink()
|
|
|
|
def _load_users(self) -> dict[str, Any]:
|
|
if not self.path.exists():
|
|
return {}
|
|
data = json.loads(self.path.read_text(encoding="utf-8"))
|
|
if not isinstance(data, dict):
|
|
return {}
|
|
users = data.get("users")
|
|
return users if isinstance(users, dict) else {}
|
|
|
|
|
|
def default_memory_gateway_users_path() -> Path:
|
|
raw = os.getenv("BEAVER_MEMORY_GATEWAY_USERS_PATH")
|
|
if raw:
|
|
return Path(raw)
|
|
return Path.home() / ".beaver" / "memory_gateway_users.json"
|