agentscope-runtime 0.1.5b2__py3-none-any.whl → 0.2.0b1__py3-none-any.whl

agentscope_runtime/common/collections/in_memory_mapping.py

@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+from typing import Any
+
+from .base_mapping import Mapping
+
+
+class InMemoryMapping(Mapping):
+    def __init__(self):
+        self.store = {}
+
+    def set(self, key: str, value: Any):
+        self.store[key] = value
+
+    def get(self, key: str) -> Any:
+        return self.store.get(key)
+
+    def delete(self, key: str):
+        if key in self.store:
+            del self.store[key]
+
+    def scan(self, prefix: str = None):
+        if prefix is None:
+            yield from list(self.store.keys())
+        else:
+            yield from list(
+                key for key in self.store if key.startswith(prefix)
+            )

agentscope_runtime/common/collections/redis_mapping.py

@@ -0,0 +1,42 @@
+# -*- coding: utf-8 -*-
+import json
+
+from typing import Any
+
+from .base_mapping import Mapping
+
+
+class RedisMapping(Mapping):
+    def __init__(self, redis_client, prefix: str = ""):
+        self.client = redis_client
+        self.prefix = prefix.rstrip(":") + ":" if prefix else ""
+
+    def _get_full_key(self, key: str) -> str:
+        return f"{self.prefix}{key}"
+
+    def _strip_prefix(self, full_key: str) -> str:
+        if self.prefix and full_key.startswith(self.prefix):
+            return full_key[len(self.prefix) :]
+        return full_key
+
+    def set(self, key: str, value: Any):
+        self.client.set(self._get_full_key(key), json.dumps(value))
+
+    def get(self, key: str) -> Any:
+        value = self.client.get(self._get_full_key(key))
+        return json.loads(value) if value else None
+
+    def delete(self, key: str):
+        self.client.delete(self._get_full_key(key))
+
+    def scan(self, prefix: str = ""):
+        search_pattern = f"{self._get_full_key(prefix)}*"
+        cursor = 0
+        while cursor != 0:
+            cursor, keys = self.client.scan(
+                cursor=cursor,
+                match=search_pattern,
+            )
+            for key in keys:
+                decoded_key = key.decode("utf-8")
+                yield self._strip_prefix(decoded_key)