275 lines
12 KiB
Python
275 lines
12 KiB
Python
"""Orchestration for the lightweight Memory System API."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any, Awaitable, Callable
|
|
|
|
from .clients import EverOSMemorySystemClient, OpenVikingMemorySystemClient
|
|
from .schemas import (
|
|
AccountResponse,
|
|
BackendStatus,
|
|
CommitResponse,
|
|
ExtractResponse,
|
|
MessageIngestRequest,
|
|
MessageIngestResponse,
|
|
ProfileResponse,
|
|
SearchRequest,
|
|
SearchResponse,
|
|
)
|
|
|
|
|
|
class MemorySystemService:
|
|
def __init__(self, openviking: Any | None = None, everos: Any | None = None) -> None:
|
|
self.openviking = openviking or OpenVikingMemorySystemClient()
|
|
self.everos = everos or EverOSMemorySystemClient()
|
|
|
|
async def create_user(self, user_id: str) -> AccountResponse:
|
|
backends = {"openviking": await self._capture(lambda: self.openviking.create_user(user_id))}
|
|
account = backends["openviking"].result if backends["openviking"].status == "success" else None
|
|
return AccountResponse(status=self._aggregate_status(backends), account=account, backends=backends)
|
|
|
|
async def ingest_messages(self, request: MessageIngestRequest) -> MessageIngestResponse:
|
|
messages = self._messages_from_request(request)
|
|
if not messages:
|
|
raise ValueError("at least one message is required")
|
|
|
|
credential = self.openviking.credential_for_user(
|
|
request.user_id,
|
|
request.user_key,
|
|
agent_id=request.session_id,
|
|
)
|
|
await self.openviking.ensure_session(credential, request.session_id)
|
|
|
|
async def write_openviking() -> list[dict[str, Any]]:
|
|
results = []
|
|
for message in messages:
|
|
results.append(
|
|
await self.openviking.append_message(credential, request.session_id, message["role"], message["content"])
|
|
)
|
|
return results
|
|
|
|
async def write_everos() -> list[dict[str, Any]]:
|
|
results = []
|
|
for message in messages:
|
|
results.append(
|
|
await self.everos.append_message(request.user_id, request.session_id, message["role"], message["content"])
|
|
)
|
|
return results
|
|
|
|
backends = await self._run_backends(openviking=write_openviking, everos=write_everos)
|
|
return MessageIngestResponse(
|
|
status=self._aggregate_status(backends),
|
|
message_count=len(messages),
|
|
backends=backends,
|
|
)
|
|
|
|
async def commit_session(self, user_id: str, user_key: str, session_id: str) -> CommitResponse:
|
|
credential = self.openviking.credential_for_user(user_id, user_key, agent_id=session_id)
|
|
|
|
async def commit_openviking() -> dict[str, Any]:
|
|
return await self.openviking.commit_session(credential, session_id)
|
|
|
|
async def flush_everos() -> dict[str, Any]:
|
|
return await self.everos.flush(user_id, session_id)
|
|
|
|
backends = await self._run_backends(openviking=commit_openviking, everos=flush_everos)
|
|
return CommitResponse(status=self._aggregate_status(backends), backends=backends)
|
|
|
|
async def extract_session(self, user_id: str, user_key: str, session_id: str) -> ExtractResponse:
|
|
credential = self.openviking.credential_for_user(user_id, user_key, agent_id=session_id)
|
|
backends = {
|
|
"openviking": await self._capture(lambda: self.openviking.extract_session(credential, session_id)),
|
|
}
|
|
return ExtractResponse(status=self._aggregate_status(backends), backends=backends)
|
|
|
|
async def get_openviking_task(
|
|
self,
|
|
user_id: str,
|
|
user_key: str,
|
|
task_id: str,
|
|
session_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
credential = self.openviking.credential_for_user(user_id, user_key, agent_id=session_id)
|
|
return await self.openviking.get_task(credential, task_id)
|
|
|
|
async def search(self, request: SearchRequest) -> SearchResponse:
|
|
credential = self.openviking.credential_for_user(
|
|
request.user_id,
|
|
request.user_key,
|
|
agent_id=request.session_id,
|
|
)
|
|
everos_method = "agentic" if request.use_llm else "hybrid"
|
|
|
|
async def search_openviking() -> dict[str, Any]:
|
|
if request.use_llm:
|
|
return await self.openviking.search(credential, request.session_id, request.query, request.limit)
|
|
return await self.openviking.find(credential, request.query, request.limit)
|
|
|
|
async def search_everos() -> dict[str, Any]:
|
|
return await self.everos.search(
|
|
request.user_id,
|
|
request.session_id,
|
|
request.query,
|
|
everos_method,
|
|
request.limit,
|
|
)
|
|
|
|
backends = await self._run_backends(openviking=search_openviking, everos=search_everos)
|
|
backends = self._remove_vectors_from_backends(backends)
|
|
items = self._merge_search_items(backends)
|
|
compact_backends = self._compact_search_backends(backends)
|
|
return SearchResponse(
|
|
status=self._aggregate_status(backends),
|
|
items=items[: request.limit],
|
|
backends=compact_backends,
|
|
)
|
|
|
|
async def get_profile(self, user_id: str) -> ProfileResponse:
|
|
backends = {"everos": await self._capture(lambda: self.everos.get_profile(user_id))}
|
|
profile = backends["everos"].result if backends["everos"].status == "success" else None
|
|
return ProfileResponse(status=self._aggregate_status(backends), profile=profile, backends=backends)
|
|
|
|
async def health(self) -> dict[str, Any]:
|
|
backends = await self._run_backends(openviking=self.openviking.health, everos=self.everos.health)
|
|
return {"status": self._aggregate_status(backends), "backends": backends}
|
|
|
|
def _messages_from_request(self, request: MessageIngestRequest) -> list[dict[str, str]]:
|
|
messages = []
|
|
if request.user_message:
|
|
messages.append({"role": "user", "content": request.user_message})
|
|
if request.assistant_message:
|
|
messages.append({"role": "assistant", "content": request.assistant_message})
|
|
return messages
|
|
|
|
async def _run_backends(self, **calls: Callable[[], Awaitable[Any]]) -> dict[str, BackendStatus]:
|
|
names = list(calls)
|
|
results = await asyncio.gather(*(self._capture(calls[name]) for name in names))
|
|
return dict(zip(names, results))
|
|
|
|
async def _capture(self, call: Callable[[], Awaitable[Any]]) -> BackendStatus:
|
|
try:
|
|
return BackendStatus(status="success", result=await call())
|
|
except Exception as exc: # noqa: BLE001
|
|
message = str(exc)
|
|
error = f"{type(exc).__name__}: {message}" if message else type(exc).__name__
|
|
return BackendStatus(status="failed", error=error)
|
|
|
|
def _aggregate_status(self, backends: dict[str, BackendStatus]) -> str:
|
|
statuses = {backend.status for backend in backends.values()}
|
|
if statuses == {"success"}:
|
|
return "success"
|
|
if "success" in statuses:
|
|
return "partial_success"
|
|
return "failed"
|
|
|
|
def _merge_search_items(self, backends: dict[str, BackendStatus]) -> list[dict[str, Any]]:
|
|
items: list[dict[str, Any]] = []
|
|
for backend_name, backend in backends.items():
|
|
if backend.status != "success":
|
|
continue
|
|
items.extend(self._items_from_backend_result(backend_name, backend.result))
|
|
return items
|
|
|
|
def _items_from_backend_result(self, backend_name: str, result: Any) -> list[dict[str, Any]]:
|
|
if isinstance(result, dict) and isinstance(result.get("items"), list):
|
|
return [self._with_backend(backend_name, item) for item in result["items"] if isinstance(item, dict)]
|
|
data = result.get("data") if isinstance(result, dict) and isinstance(result.get("data"), dict) else result
|
|
if not isinstance(data, dict):
|
|
return []
|
|
if isinstance(data.get("result"), dict):
|
|
data = data["result"]
|
|
raw_items: list[dict[str, Any]] = []
|
|
for key in ("memories", "resources", "episodes", "profiles", "raw_messages"):
|
|
values = data.get(key)
|
|
if isinstance(values, list):
|
|
raw_items.extend(
|
|
self._compact_search_item(backend_name, key, item)
|
|
for item in values
|
|
if isinstance(item, dict)
|
|
)
|
|
return [self._with_backend(backend_name, item) for item in raw_items]
|
|
|
|
def _with_backend(self, backend_name: str, item: dict[str, Any]) -> dict[str, Any]:
|
|
if "source_backend" in item:
|
|
return item
|
|
return {"source_backend": backend_name, **item}
|
|
|
|
def _compact_search_item(self, backend_name: str, collection: str, item: dict[str, Any]) -> dict[str, Any]:
|
|
if backend_name == "everos":
|
|
fields = (
|
|
"id",
|
|
"user_id",
|
|
"session_id",
|
|
"timestamp",
|
|
"summary",
|
|
"score",
|
|
)
|
|
compact = {"memory_type": self._singular_memory_type(collection)}
|
|
compact.update({field: item[field] for field in fields if field in item and item[field] is not None})
|
|
return compact
|
|
return item
|
|
|
|
def _singular_memory_type(self, collection: str) -> str:
|
|
names = {
|
|
"memories": "memory",
|
|
"resources": "resource",
|
|
"episodes": "episode",
|
|
"profiles": "profile",
|
|
"raw_messages": "raw_message",
|
|
}
|
|
return names.get(collection, collection)
|
|
|
|
def _compact_search_backends(self, backends: dict[str, BackendStatus]) -> dict[str, BackendStatus]:
|
|
return {
|
|
name: backend.model_copy(update={"result": self._compact_backend_result(name, backend.result)})
|
|
for name, backend in backends.items()
|
|
}
|
|
|
|
def _compact_backend_result(self, backend_name: str, result: Any) -> Any:
|
|
if backend_name == "everos":
|
|
data = result.get("data") if isinstance(result, dict) and isinstance(result.get("data"), dict) else result
|
|
if not isinstance(data, dict):
|
|
return result
|
|
compact: dict[str, Any] = {
|
|
"counts": {
|
|
key: len(data.get(key) or [])
|
|
for key in ("episodes", "profiles", "raw_messages")
|
|
if isinstance(data.get(key), list)
|
|
}
|
|
}
|
|
if "query" in data:
|
|
compact["query"] = data["query"]
|
|
return compact
|
|
|
|
if backend_name == "openviking":
|
|
data = result.get("result") if isinstance(result, dict) and isinstance(result.get("result"), dict) else result
|
|
if not isinstance(data, dict):
|
|
return result
|
|
compact = {
|
|
"status": result.get("status") if isinstance(result, dict) else None,
|
|
"total": data.get("total"),
|
|
"counts": {
|
|
key: len(data.get(key) or [])
|
|
for key in ("memories", "resources", "skills")
|
|
if isinstance(data.get(key), list)
|
|
},
|
|
}
|
|
if "query_plan" in data:
|
|
compact["query_plan"] = data["query_plan"]
|
|
return {key: value for key, value in compact.items() if value is not None}
|
|
|
|
return result
|
|
|
|
def _remove_vectors_from_backends(self, backends: dict[str, BackendStatus]) -> dict[str, BackendStatus]:
|
|
return {
|
|
name: backend.model_copy(update={"result": self._remove_vectors(backend.result)})
|
|
for name, backend in backends.items()
|
|
}
|
|
|
|
def _remove_vectors(self, value: Any) -> Any:
|
|
if isinstance(value, dict):
|
|
return {key: self._remove_vectors(item) for key, item in value.items() if key != "vector"}
|
|
if isinstance(value, list):
|
|
return [self._remove_vectors(item) for item in value]
|
|
return value
|