Merge pull request #1191 from writer/vlad/AB-683-pt2

feat: expose API for KV access

Author: UladzislauK-Writer

src/writer/keyvalue_storage.py

@@ -1,7 +1,7 @@
 import logging
 import os
 from functools import partial
-from typing import Any, Dict, Literal, Optional, Protocol
+from typing import Any, Dict, List, Literal, Optional, Protocol
 
 import httpx
 
@@ -42,6 +42,9 @@ def _get_agent_ids(self):
 
     def get(self, key: str, type_: Literal["data", "secret"]) -> Dict[str, Any]:
         return self._request(partial(self._client.get, url=f"{self.api_url}/agent_{type_}/{key}")).json()
+
+    def get_data_keys(self) -> List[str]:
+        return self._request(partial(self._client.get, url=f"{self.api_url}/agent_data")).json()["keys"]
 
     def save(self, key: str, data: Any) -> Dict[str, Any]:
         try:

src/writer/serve.py

@@ -47,6 +47,7 @@
     AppProcessServerResponse,
     AutogenRequestBody,
     ComponentUpdateRequestPayload,
+    DeleteDataRequestBody,
     EventResponsePayload,
     HashRequestPayload,
     HashRequestResponsePayload,
@@ -55,6 +56,8 @@
     InitResponseBodyRun,
     InitSessionRequestPayload,
     InitSessionResponsePayload,
+    RetrieveDataRequestBody,
+    RetrieveDataResponseBody,
     ServeMode,
     StateEnquiryResponsePayload,
     WriterEvent,
@@ -290,6 +293,38 @@ async def autogen(requestBody: AutogenRequestBody, request: Request):
             agent_token_header
             )
 
+    @app.post("/api/data/retrieve")
+    async def retrieve_data(requestBody: RetrieveDataRequestBody) -> RetrieveDataResponseBody:
+        from writer.keyvalue_storage import writer_kv_storage
+
+        all_keys = writer_kv_storage.get_data_keys()
+
+        keys_to_fetch = []
+        for key in all_keys:
+            if key in requestBody.skip_keys:
+                continue
+            if requestBody.key_contains and requestBody.key_contains not in key:
+                continue
+            keys_to_fetch.append(key)
+        
+        async def fetch_value(key: str):
+            return key, await asyncio.to_thread(writer_kv_storage.get, key, "data")
+
+        kv_pairs = await asyncio.gather(*(fetch_value(key) for key in keys_to_fetch))
+
+        return RetrieveDataResponseBody(result={k: v["data"] for k, v in kv_pairs})
+
+    @app.post("/api/data/delete")
+    async def delete_data(requestBody: DeleteDataRequestBody) -> None:
+        from writer.keyvalue_storage import writer_kv_storage
+
+        async def delete_key(key: str):
+            return key, await asyncio.to_thread(writer_kv_storage.delete, key)
+
+        await asyncio.gather(*(delete_key(key) for key in requestBody.keys))
+
+        return None
+
     @app.post("/api/init")
     async def init(
         initBody: InitRequestBody, request: Request, response: Response

src/writer/ss_types.py

@@ -79,6 +79,19 @@ class AutogenRequestBody(BaseModel):
     description: str
 
 
+class RetrieveDataRequestBody(BaseModel):
+    skip_keys: List[str] = []
+    key_contains: Optional[str] = None
+
+
+class RetrieveDataResponseBody(BaseModel):
+    result: Dict[str, Any]
+
+
+class DeleteDataRequestBody(BaseModel):
+    keys: List[str]
+
+
 class InitRequestBody(BaseModel):
     proposedSessionId: Optional[str] = None