tomtan/memory-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refine memory system user-key flow and search output

Browse Source
This commit is contained in:
tomtan
2026-05-22 16:30:42 +08:00
parent 92632553ab
commit d73f59f38d
16 changed files with 1888 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
memory_system_api/service.py
View File

@ -6,6 +6,7 @@ from typing import Any, Awaitable, Callable
from .clients import EverOSMemorySystemClient, OpenVikingMemorySystemClient
from .schemas import (
AccountResponse,
BackendStatus,
CommitResponse,
ExtractResponse,
@ -22,12 +23,21 @@ class MemorySystemService:
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 = await self.openviking.ensure_user(request.user_id)
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]]:
@ -53,11 +63,11 @@ class MemorySystemService:
backends=backends,
)
async def commit_session(self, user_id: str, session_id: str) -> CommitResponse:
user_key = await self.openviking.ensure_user(user_id)
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(user_key, session_id)
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)
@ -65,25 +75,35 @@ class MemorySystemService:
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, session_id: str) -> ExtractResponse:
user_key = await self.openviking.ensure_user(user_id)
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(user_key, session_id)),
"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, task_id: str) -> dict[str, Any]:
user_key = await self.openviking.ensure_user(user_id)
return await self.openviking.get_task(user_key, task_id)
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:
user_key = await self.openviking.ensure_user(request.user_id)
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(user_key, request.session_id, request.query, request.limit)
return await self.openviking.find(user_key, request.user_id, request.query, request.limit)
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(
@ -97,7 +117,12 @@ class MemorySystemService:
backends = await self._run_backends(openviking=search_openviking, everos=search_everos)
backends = self._remove_vectors_from_backends(backends)
items = self._merge_search_items(backends)
return SearchResponse(status=self._aggregate_status(backends), items=items[: request.limit], backends=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))}
@ -157,7 +182,11 @@ class MemorySystemService:
for key in ("memories", "resources", "episodes", "profiles", "raw_messages"):
values = data.get(key)
if isinstance(values, list):
raw_items.extend(item for item in values if isinstance(item, dict))
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]:
@ -165,6 +194,72 @@ class MemorySystemService:
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)})