Refactor OpenViking Memory API and User Management

- Updated API authentication headers to use `X-API-Key` for both admin and user APIs.
- Modified the account creation process to directly create user-specific accounts without requiring an admin workspace.
- Enhanced user creation to return account-specific details, including `admin_user_id`.
- Introduced new endpoints for retrieving task status and user profiles, allowing for more flexible user data management.
- Updated search functionality to include additional parameters such as `level` and `score_threshold`.
- Improved the handling of user keys in the storage layer to associate them with specific accounts.
- Added tests to validate the new user account creation process and search functionalities, ensuring proper integration with the OpenViking service.
- Included new documentation to reflect changes in API usage and expected request/response formats.

memory_system_api/api.py

@@ -6,9 +6,11 @@ from fastapi import APIRouter, Depends, HTTPException, Query, status
 from .auth import verify_api_key
 from .schemas import (
     MessageIngestRequest,
+    ProfileRequest,
     SearchRequest,
     SessionContextRequest,
     SessionUserRequest,
+    TaskStatusRequest,
     UserCreateRequest,
 )
 from .service import MemorySystemService
@@ -123,6 +125,23 @@ async def get_openviking_task(
         raise user_auth_error(exc) from exc
 
 
+@router.post("/openviking/tasks/{task_id}")
+async def get_openviking_task_from_body(
+    task_id: str,
+    request: TaskStatusRequest,
+    service: MemorySystemService = Depends(get_service),
+):
+    try:
+        return await service.get_openviking_task(
+            request.user_id,
+            request.user_key,
+            task_id,
+            session_id=request.session_id,
+        )
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
+
+
 @router.post("/search")
 async def search(
     request: SearchRequest,
@@ -134,14 +153,34 @@ async def search(
         raise user_auth_error(exc) from exc
 
 
+@router.post("/users/{user_id}/profile")
+async def get_profile_from_body(
+    user_id: str,
+    request: ProfileRequest,
+    service: MemorySystemService = Depends(get_service),
+):
+    try:
+        return await service.get_profile(
+            user_id,
+            request.user_key,
+            query=request.query,
+            limit=request.limit,
+            level=request.level,
+        )
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
+
+
 @router.get("/users/{user_id}/profile")
 async def get_profile(
     user_id: str,
     user_key: str = Query(min_length=1),
+    query: str = Query(default="用户画像", min_length=1),
+    limit: int = Query(default=10, ge=1, le=100),
+    level: int = Query(default=2, ge=0),
     service: MemorySystemService = Depends(get_service),
 ):
     try:
-        service.openviking.credential_for_user(user_id, user_key)
+        return await service.get_profile(user_id, user_key, query=query, limit=limit, level=level)
     except PermissionError as exc:
         raise user_auth_error(exc) from exc
-    return await service.get_profile(user_id)

memory_system_api/clients.py

@@ -49,36 +49,24 @@ class OpenVikingMemorySystemClient:
         user_key = self._extract_user_key(data)
         if user_key:
             self.store.save_account_key(account_id, admin_user_id, user_key)
-            self.store.save_user_key(admin_user_id, user_key)
+            self.store.save_user_key(admin_user_id, user_key, account_id=account_id)
         return data
 
-    async def ensure_admin_workspace(self) -> None:
-        if self.store.get_account_key(ADMIN_ACCOUNT_ID):
-            return
-        await self.create_account(ADMIN_ACCOUNT_ID, ADMIN_USER_ID)
-
     async def create_user(self, user_id: str, role: str = "user") -> dict[str, Any]:
         existing = self.store.get_user_key(user_id)
+        account_id = self._account_id_for_user(user_id)
         if existing:
-            return {"status": "ok", "result": {"account_id": ADMIN_ACCOUNT_ID, "user_id": user_id, "user_key": existing}}
+            return {
+                "status": "ok",
+                "result": {
+                    "account_id": account_id,
+                    "admin_user_id": user_id,
+                    "user_id": user_id,
+                    "user_key": existing,
+                },
+            }
 
-        await self.ensure_admin_workspace()
-        if user_id == ADMIN_USER_ID:
-            user_key = self.store.get_user_key(user_id)
-            return {"status": "ok", "result": {"account_id": ADMIN_ACCOUNT_ID, "user_id": user_id, "user_key": user_key}}
-
-        async with self._client(self.root_key) as client:
-            response = await client.post(
-                f"/api/v1/admin/accounts/{ADMIN_ACCOUNT_ID}/users",
-                json={"user_id": user_id, "role": role},
-            )
-            response.raise_for_status()
-            data = response.json()
-
-        user_key = self._extract_user_key(data)
-        if user_key:
-            self.store.save_user_key(user_id, user_key)
-        return data
+        return await self.create_account(account_id, user_id)
 
     def credential_for_user(
         self,
@@ -98,7 +86,7 @@ class OpenVikingMemorySystemClient:
     ) -> OpenVikingCredential:
         return OpenVikingCredential(
             api_key=user_key,
-            account_id=ADMIN_ACCOUNT_ID,
+            account_id=self._account_id_for_user(user_id),
             user_id=user_id,
             agent_id=agent_id,
             user_key_auth=True,
@@ -164,22 +152,46 @@ class OpenVikingMemorySystemClient:
             return response.json()
 
     async def search(
-        self, credential: OpenVikingCredential | str, session_id: str | None, query: str, limit: int
+        self,
+        credential: OpenVikingCredential | str,
+        query: str,
+        limit: int,
+        level: int = 2,
+        score_threshold: float = 0.8,
+        target_uri: str = "viking://user/memories",
     ) -> dict[str, Any]:
-        payload: dict[str, Any] = {"query": query, "limit": limit}
-        if session_id:
-            payload["session_id"] = session_id
-        if isinstance(credential, OpenVikingCredential) and credential.user_id:
-            payload["target_uri"] = (
-                f"viking://user/{credential.user_id}/{session_id}"
-                if session_id
-                else f"viking://user/{credential.user_id}/memories/"
-            )
+        payload: dict[str, Any] = {
+            "query": query,
+            "target_uri": target_uri,
+            "limit": limit,
+            "level": level,
+            "score_threshold": score_threshold,
+        }
         async with self._credential_client(credential) as client:
             response = await client.post("/api/v1/search/search", json=payload)
             response.raise_for_status()
             return response.json()
 
+    async def search_profile_memories(
+        self,
+        credential: OpenVikingCredential | str,
+        query: str,
+        limit: int,
+        level: int,
+    ) -> dict[str, Any]:
+        async with self._credential_client(credential) as client:
+            response = await client.post(
+                "/api/v1/search/search",
+                json={
+                    "query": query,
+                    "limit": limit,
+                    "level": level,
+                    "target_uri": "viking://user/memories",
+                },
+            )
+            response.raise_for_status()
+            return response.json()
+
     async def get_session_context(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
         async with self._credential_client(credential) as client:
             response = await client.get(f"/api/v1/sessions/{session_id}/context")
@@ -201,11 +213,7 @@ class OpenVikingMemorySystemClient:
         return self._client(credential.api_key, headers)
 
     def _client(self, api_key: str, extra_headers: dict[str, str] | None = None) -> httpx.AsyncClient:
-        headers = {"Content-Type": "application/json"}
-        if api_key == self.root_key:
-            headers["X-API-Key"] = api_key
-        else:
-            headers["Authorization"] = f"Bearer {api_key}"
+        headers = {"Content-Type": "application/json", "X-API-Key": api_key}
         if extra_headers:
             headers.update(extra_headers)
         return httpx.AsyncClient(
@@ -226,6 +234,9 @@ class OpenVikingMemorySystemClient:
             "admin_user_id": admin_user_id,
         }
 
+    def _account_id_for_user(self, user_id: str) -> str:
+        return f"{user_id}_account"
+
     def _save_session(self, credential: OpenVikingCredential | str, session_id: str) -> None:
         if isinstance(credential, OpenVikingCredential) and credential.user_id:
             self.store.save_session(credential.user_id, session_id)

memory_system_api/schemas.py

@@ -24,6 +24,10 @@ class SessionUserRequest(BaseModel):
     user_key: str = Field(min_length=1)
 
 
+class TaskStatusRequest(SessionUserRequest):
+    session_id: str | None = Field(default=None, min_length=1)
+
+
 class SearchRequest(BaseModel):
     user_id: str = Field(min_length=1)
     user_key: str = Field(min_length=1)
@@ -31,6 +35,9 @@ class SearchRequest(BaseModel):
     query: str = Field(min_length=1)
     use_llm: bool = False
     limit: int = Field(default=10, ge=1, le=100)
+    level: int = Field(default=2, ge=0)
+    score_threshold: float = Field(default=0.8, ge=0, le=1)
+    target_uri: str = Field(default="viking://user/memories", min_length=1)
 
 
 class SessionContextRequest(BaseModel):
@@ -40,6 +47,13 @@ class SessionContextRequest(BaseModel):
     limit: int = Field(default=10, ge=1, le=100)
 
 
+class ProfileRequest(BaseModel):
+    user_key: str = Field(min_length=1)
+    query: str = Field(default="用户画像", min_length=1)
+    limit: int = Field(default=10, ge=1, le=100)
+    level: int = Field(default=2, ge=0)
+
+
 class BackendStatus(BaseModel):
     status: OperationStatus
     result: Any = None
@@ -88,4 +102,5 @@ class SessionContextResponse(BaseModel):
 class ProfileResponse(BaseModel):
     status: OperationStatus
     profile: Any = None
+    items: list[dict[str, Any]] = Field(default_factory=list)
     backends: dict[str, BackendStatus]

memory_system_api/service.py

@@ -103,9 +103,14 @@ class MemorySystemService:
         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)
+            return await self.openviking.search(
+                credential,
+                request.query,
+                request.limit,
+                request.level,
+                request.score_threshold,
+                request.target_uri,
+            )
 
         async def search_everos() -> dict[str, Any]:
             return await self.everos.search(
@@ -160,10 +165,32 @@ class MemorySystemService:
             backends=self._compact_session_context_backends(backends),
         )
 
-    async def get_profile(self, user_id: str) -> ProfileResponse:
-        backends = {"everos": await self._capture(lambda: self.everos.get_profile(user_id))}
+    async def get_profile(
+        self,
+        user_id: str,
+        user_key: str,
+        query: str = "用户画像",
+        limit: int = 10,
+        level: int = 2,
+    ) -> ProfileResponse:
+        credential = self.openviking.credential_for_user(user_id, user_key)
+        backends = await self._run_backends(
+            everos=lambda: self.everos.get_profile(user_id),
+            openviking=lambda: self.openviking.search_profile_memories(credential, query, limit, level),
+        )
+        backends = self._remove_vectors_from_backends(backends)
         profile = backends["everos"].result if backends["everos"].status == "success" else None
-        return ProfileResponse(status=self._aggregate_status(backends), profile=profile, backends=backends)
+        items = (
+            self._items_from_backend_result("openviking", backends["openviking"].result)[:limit]
+            if backends["openviking"].status == "success"
+            else []
+        )
+        return ProfileResponse(
+            status=self._aggregate_status(backends),
+            profile=profile,
+            items=items,
+            backends=self._compact_profile_backends(backends),
+        )
 
     async def health(self) -> dict[str, Any]:
         backends = await self._run_backends(openviking=self.openviking.health, everos=self.everos.health)
@@ -323,6 +350,31 @@ class MemorySystemService:
             return {key: value for key, value in compact.items() if value is not None}
         return self._compact_backend_result(backend_name, result)
 
+    def _compact_profile_backends(self, backends: dict[str, BackendStatus]) -> dict[str, BackendStatus]:
+        return {
+            name: backend.model_copy(update={"result": self._compact_profile_backend_result(name, backend.result)})
+            for name, backend in backends.items()
+        }
+
+    def _compact_profile_backend_result(self, backend_name: str, result: Any) -> Any:
+        if backend_name == "openviking":
+            return self._compact_backend_result("openviking", result)
+        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] = {}
+            for key in ("total_count", "count"):
+                if key in data:
+                    compact[key] = data[key]
+            compact["counts"] = {
+                key: len(data.get(key) or [])
+                for key in ("episodes", "profiles", "agent_cases", "agent_skills")
+                if isinstance(data.get(key), list)
+            }
+            return compact
+        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)})

memory_system_api/store.py

@@ -65,7 +65,7 @@ class OpenVikingUserKeyStore:
                 ).fetchone()
         return str(row[0]) if row else None
 
-    def save_user_key(self, user_id: str, user_key: str) -> None:
+    def save_user_key(self, user_id: str, user_key: str, account_id: str = ADMIN_ACCOUNT_ID) -> None:
         now = datetime.now(timezone.utc).isoformat()
         with self._connect() as conn:
             conn.execute(
@@ -76,7 +76,7 @@ class OpenVikingUserKeyStore:
                     user_key = excluded.user_key,
                     updated_at = excluded.updated_at
                 """,
-                (user_id, ADMIN_ACCOUNT_ID, user_key, now, now),
+                (user_id, account_id, user_key, now, now),
             )
 
     def user_key_matches(self, user_id: str, user_key: str) -> bool: