Simplify to memory system api

memory_system_api/service.py

@@ -0,0 +1,165 @@
+"""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 (
+    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 ingest_messages(self, request: MessageIngestRequest) -> MessageIngestResponse:
+        messages = self._messages_from_request(request)
+        if not messages:
+            raise ValueError("at least one message is required")
+
+        user_key = await self.openviking.ensure_user(request.user_id)
+        await self.openviking.ensure_session(user_key, request.session_id)
+
+        print("user_key:", user_key)  # Debugging line to check the user_key value
+
+        async def write_openviking() -> list[dict[str, Any]]:
+            results = []
+            for message in messages:
+                results.append(
+                    await self.openviking.append_message(user_key, 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, session_id: str) -> CommitResponse:
+        user_key = await self.openviking.ensure_user(user_id)
+
+        async def commit_openviking() -> dict[str, Any]:
+            return await self.openviking.commit_session(user_key, 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, session_id: str) -> ExtractResponse:
+        user_key = await self.openviking.ensure_user(user_id)
+        backends = {
+            "openviking": await self._capture(lambda: self.openviking.extract_session(user_key, 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 search(self, request: SearchRequest) -> SearchResponse:
+        user_key = await self.openviking.ensure_user(request.user_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)
+
+        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)
+        items = self._merge_search_items(backends)
+        return SearchResponse(status=self._aggregate_status(backends), items=items[: request.limit], backends=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
+            return BackendStatus(status="failed", error=str(exc))
+
+    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(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}