feat(memory): initialize optional gateway layer
This commit is contained in:
@ -3,6 +3,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import logging
|
||||||
import os
|
import os
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@ -17,6 +18,7 @@ from beaver.memory.curated.store import MemoryStore
|
|||||||
from beaver.memory.runs import RunMemoryStore
|
from beaver.memory.runs import RunMemoryStore
|
||||||
from beaver.memory.skills import SkillLearningStore
|
from beaver.memory.skills import SkillLearningStore
|
||||||
from beaver.services.memory_service import MemoryService
|
from beaver.services.memory_service import MemoryService
|
||||||
|
from beaver.services.memory_gateway_service import MemoryGatewayService
|
||||||
from beaver.skills.drafts import DraftService
|
from beaver.skills.drafts import DraftService
|
||||||
from beaver.skills.learning import EvidenceSelector, SkillDraftSynthesizer, SkillLearningPipelineService, SkillLearningService
|
from beaver.skills.learning import EvidenceSelector, SkillDraftSynthesizer, SkillLearningPipelineService, SkillLearningService
|
||||||
from beaver.skills.learning.safety import SkillDraftSafetyChecker
|
from beaver.skills.learning.safety import SkillDraftSafetyChecker
|
||||||
@ -59,6 +61,8 @@ from beaver.tools.builtins import (
|
|||||||
WriteFileTool,
|
WriteFileTool,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
class EngineLoadResult:
|
class EngineLoadResult:
|
||||||
@ -80,6 +84,7 @@ class EngineLoadResult:
|
|||||||
session_manager: SessionManager | None = None
|
session_manager: SessionManager | None = None
|
||||||
curated_memory_store: MemoryStore | None = None
|
curated_memory_store: MemoryStore | None = None
|
||||||
memory_service: MemoryService | None = None
|
memory_service: MemoryService | None = None
|
||||||
|
memory_gateway_service: MemoryGatewayService | None = None
|
||||||
run_memory_store: RunMemoryStore | None = None
|
run_memory_store: RunMemoryStore | None = None
|
||||||
skill_learning_store: SkillLearningStore | None = None
|
skill_learning_store: SkillLearningStore | None = None
|
||||||
tool_registry: ToolRegistry | None = None
|
tool_registry: ToolRegistry | None = None
|
||||||
@ -155,6 +160,7 @@ class EngineLoader:
|
|||||||
session_manager: SessionManager | None = None,
|
session_manager: SessionManager | None = None,
|
||||||
curated_memory_store: MemoryStore | None = None,
|
curated_memory_store: MemoryStore | None = None,
|
||||||
memory_service: MemoryService | None = None,
|
memory_service: MemoryService | None = None,
|
||||||
|
memory_gateway_service: MemoryGatewayService | None = None,
|
||||||
run_memory_store: RunMemoryStore | None = None,
|
run_memory_store: RunMemoryStore | None = None,
|
||||||
skill_learning_store: SkillLearningStore | None = None,
|
skill_learning_store: SkillLearningStore | None = None,
|
||||||
tool_registry: ToolRegistry | None = None,
|
tool_registry: ToolRegistry | None = None,
|
||||||
@ -180,6 +186,7 @@ class EngineLoader:
|
|||||||
self._session_manager = session_manager
|
self._session_manager = session_manager
|
||||||
self._curated_memory_store = curated_memory_store
|
self._curated_memory_store = curated_memory_store
|
||||||
self._memory_service = memory_service
|
self._memory_service = memory_service
|
||||||
|
self._memory_gateway_service = memory_gateway_service
|
||||||
self._run_memory_store = run_memory_store
|
self._run_memory_store = run_memory_store
|
||||||
self._skill_learning_store = skill_learning_store
|
self._skill_learning_store = skill_learning_store
|
||||||
self._tool_registry = tool_registry
|
self._tool_registry = tool_registry
|
||||||
@ -208,6 +215,7 @@ class EngineLoader:
|
|||||||
curated_memory_store = self._curated_memory_store or MemoryStore(curated_root)
|
curated_memory_store = self._curated_memory_store or MemoryStore(curated_root)
|
||||||
memory_service = self._memory_service or MemoryService(curated_root, store=curated_memory_store)
|
memory_service = self._memory_service or MemoryService(curated_root, store=curated_memory_store)
|
||||||
memory_service.initialize()
|
memory_service.initialize()
|
||||||
|
memory_gateway_service = self._resolve_memory_gateway_service()
|
||||||
run_memory_store = self._run_memory_store or RunMemoryStore(workspace / "memory" / "runs")
|
run_memory_store = self._run_memory_store or RunMemoryStore(workspace / "memory" / "runs")
|
||||||
skill_learning_store = self._skill_learning_store or SkillLearningStore(workspace / "memory" / "skills")
|
skill_learning_store = self._skill_learning_store or SkillLearningStore(workspace / "memory" / "skills")
|
||||||
|
|
||||||
@ -298,11 +306,12 @@ class EngineLoader:
|
|||||||
config=self.config,
|
config=self.config,
|
||||||
tools=[spec.name for spec in tool_registry.list_specs()],
|
tools=[spec.name for spec in tool_registry.list_specs()],
|
||||||
skills=[record.name for record in skills_loader.list_skills(filter_unavailable=False)],
|
skills=[record.name for record in skills_loader.list_skills(filter_unavailable=False)],
|
||||||
memory_stores=["curated"],
|
memory_stores=["curated", *(["memory_gateway"] if memory_gateway_service is not None else [])],
|
||||||
permissions=[],
|
permissions=[],
|
||||||
session_manager=session_manager,
|
session_manager=session_manager,
|
||||||
curated_memory_store=memory_service.get_store(),
|
curated_memory_store=memory_service.get_store(),
|
||||||
memory_service=memory_service,
|
memory_service=memory_service,
|
||||||
|
memory_gateway_service=memory_gateway_service,
|
||||||
run_memory_store=run_memory_store,
|
run_memory_store=run_memory_store,
|
||||||
skill_learning_store=skill_learning_store,
|
skill_learning_store=skill_learning_store,
|
||||||
tool_registry=tool_registry,
|
tool_registry=tool_registry,
|
||||||
@ -328,6 +337,23 @@ class EngineLoader:
|
|||||||
result.register_closeable("mcp_manager", lambda: _close_mcp_manager(mcp_manager))
|
result.register_closeable("mcp_manager", lambda: _close_mcp_manager(mcp_manager))
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
def _resolve_memory_gateway_service(self) -> MemoryGatewayService | None:
|
||||||
|
memory_config = self.config.memory
|
||||||
|
if memory_config.mode == "curated":
|
||||||
|
return None
|
||||||
|
|
||||||
|
gateway_config = memory_config.gateway
|
||||||
|
if memory_config.explicit and not gateway_config.is_configured:
|
||||||
|
raise ValueError(
|
||||||
|
"Explicit hybrid memory requires complete Memory Gateway configuration"
|
||||||
|
)
|
||||||
|
if not gateway_config.is_configured:
|
||||||
|
logger.warning(
|
||||||
|
"Memory Gateway is not configured; continuing with curated memory only"
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
return self._memory_gateway_service or MemoryGatewayService(gateway_config)
|
||||||
|
|
||||||
|
|
||||||
def _close_mcp_manager(manager: MCPConnectionManager) -> None:
|
def _close_mcp_manager(manager: MCPConnectionManager) -> None:
|
||||||
try:
|
try:
|
||||||
|
|||||||
@ -0,0 +1,84 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from beaver.engine import EngineLoader
|
||||||
|
from beaver.foundation.config import BeaverConfig, MemoryConfig, MemoryGatewayConfig
|
||||||
|
|
||||||
|
|
||||||
|
def test_loader_keeps_curated_memory_in_explicit_curated_mode(tmp_path) -> None:
|
||||||
|
config = BeaverConfig(memory=MemoryConfig(mode="curated", explicit=True))
|
||||||
|
|
||||||
|
loaded = EngineLoader(workspace=tmp_path, config=config).load()
|
||||||
|
|
||||||
|
try:
|
||||||
|
assert loaded.memory_gateway_service is None
|
||||||
|
assert loaded.curated_memory_store is not None
|
||||||
|
assert loaded.memory_service is not None
|
||||||
|
assert "memory" in loaded.tools
|
||||||
|
assert loaded.memory_stores == ["curated"]
|
||||||
|
finally:
|
||||||
|
loaded.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_loader_adds_gateway_service_without_disabling_curated_memory(tmp_path) -> None:
|
||||||
|
gateway_config = MemoryGatewayConfig(
|
||||||
|
base_url="http://gateway.test",
|
||||||
|
user_id="gateway-user",
|
||||||
|
user_key="uk_secret",
|
||||||
|
)
|
||||||
|
config = BeaverConfig(
|
||||||
|
memory=MemoryConfig(mode="hybrid", explicit=True, gateway=gateway_config)
|
||||||
|
)
|
||||||
|
fake_gateway_service = object()
|
||||||
|
|
||||||
|
loaded = EngineLoader(
|
||||||
|
workspace=tmp_path,
|
||||||
|
config=config,
|
||||||
|
memory_gateway_service=fake_gateway_service,
|
||||||
|
).load()
|
||||||
|
|
||||||
|
try:
|
||||||
|
assert loaded.memory_gateway_service is fake_gateway_service
|
||||||
|
assert loaded.curated_memory_store is not None
|
||||||
|
assert loaded.memory_service is not None
|
||||||
|
assert "memory" in loaded.tools
|
||||||
|
assert loaded.memory_stores == ["curated", "memory_gateway"]
|
||||||
|
finally:
|
||||||
|
loaded.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_loader_implicit_hybrid_without_credentials_warns_and_degrades(
|
||||||
|
tmp_path,
|
||||||
|
caplog,
|
||||||
|
) -> None:
|
||||||
|
config = BeaverConfig(memory=MemoryConfig(mode="hybrid", explicit=False))
|
||||||
|
|
||||||
|
with caplog.at_level(logging.WARNING):
|
||||||
|
loaded = EngineLoader(workspace=tmp_path, config=config).load()
|
||||||
|
|
||||||
|
try:
|
||||||
|
assert loaded.memory_gateway_service is None
|
||||||
|
assert loaded.curated_memory_store is not None
|
||||||
|
assert "memory" in loaded.tools
|
||||||
|
assert "continuing with curated memory only" in caplog.text
|
||||||
|
finally:
|
||||||
|
loaded.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_loader_explicit_hybrid_without_credentials_fails_without_secret(tmp_path) -> None:
|
||||||
|
config = BeaverConfig(
|
||||||
|
memory=MemoryConfig(
|
||||||
|
mode="hybrid",
|
||||||
|
explicit=True,
|
||||||
|
gateway=MemoryGatewayConfig(user_key="uk_super_secret"),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
EngineLoader(workspace=tmp_path, config=config).load()
|
||||||
|
|
||||||
|
assert "Memory Gateway" in str(exc_info.value)
|
||||||
|
assert "uk_super_secret" not in str(exc_info.value)
|
||||||
Reference in New Issue
Block a user