Spaces:
Running
Running
File size: 4,865 Bytes
a55ea57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """
File: /client.py
Created Date: Sunday October 19th 2025
Author: Christian Nonis <alch.infoemail@gmail.com>
-----
Last Modified: Sunday October 19th 2025 12:37:27 pm
Modified By: the developer formerly known as Christian Nonis at <alch.infoemail@gmail.com>
-----
"""
import json
import time
from typing import Optional
from redis import Redis
from redis.connection import ConnectionPool
from src.adapters.interfaces.cache import CacheClient
from src.config import config
TASK_RETENTION_SECONDS = 3600 * 24 * 7
class RedisClient(CacheClient):
"""
Redis client.
"""
def __init__(self):
pool = ConnectionPool(
host=config.redis.host,
port=config.redis.port,
password=config.redis.password,
socket_connect_timeout=5,
socket_timeout=5,
retry_on_timeout=True,
health_check_interval=30,
max_connections=50,
)
self.client = Redis(connection_pool=pool)
def _get_key(self, key: str, brain_id: str) -> str:
"""
Get the prefixed key for a given brain_id.
"""
return f"{brain_id}:{key}"
def get(self, key: str, brain_id: str) -> str:
"""
Get a value from the cache.
"""
prefixed_key = self._get_key(key, brain_id)
result = self.client.get(prefixed_key)
if result is None:
return None
if isinstance(result, bytes):
return result.decode("utf-8")
return result
def set(
self, key: str, value: str, brain_id: str, expires_in: Optional[int] = None
) -> bool:
"""
Set a value in the cache with an expiration time.
Task keys are stored WITHOUT TTL to prevent eviction by volatile-lru.
"""
prefixed_key = self._get_key(key, brain_id)
if key.startswith("task:"):
task_id = key.split(":")[-1]
self.client.hset(
f"{brain_id}:_tasks",
task_id,
json.dumps({"data": value, "created_at": time.time()}),
)
return True
return self.client.set(
prefixed_key, value, **({"ex": expires_in} if expires_in else {})
)
def delete(self, key: str, brain_id: str) -> bool:
"""
Delete a value from the cache.
"""
prefixed_key = self._get_key(key, brain_id)
if key.startswith("task:"):
task_id = key.split(":")[-1]
return self.client.hdel(f"{brain_id}:_tasks", task_id)
return self.client.delete(prefixed_key)
def clear_brain_data(self, brain_id: str) -> int:
"""Delete only keys scoped to one brain without flushing shared Redis."""
deleted = 0
chunk = []
for key in self.client.scan_iter(match=f"{brain_id}:*", count=500):
chunk.append(key)
if len(chunk) == 500:
deleted += int(self.client.delete(*chunk))
chunk = []
if chunk:
deleted += int(self.client.delete(*chunk))
return deleted
def clear_brain_registry_cache(self, brain_id: str) -> int:
return int(
self.client.delete(
f"system:brain:{brain_id}",
f"system:brainpat:{brain_id}",
)
)
def get_task_keys(self, brain_id: str) -> list[str]:
"""
Get all task keys for a given brain_id.
Also cleans up expired tasks (older than TASK_RETENTION_SECONDS).
"""
tasks_hash = f"{brain_id}:_tasks"
all_tasks = self.client.hgetall(tasks_hash)
valid_keys = []
now = time.time()
expired_keys = []
for task_id, task_data in all_tasks.items():
task_id_str = (
task_id.decode("utf-8") if isinstance(task_id, bytes) else task_id
)
try:
data = json.loads(task_data)
created_at = data.get("created_at", 0)
if now - created_at > TASK_RETENTION_SECONDS:
expired_keys.append(task_id_str)
else:
valid_keys.append(f"task:{task_id_str}")
except (json.JSONDecodeError, TypeError):
expired_keys.append(task_id_str)
if expired_keys:
self.client.hdel(tasks_hash, *expired_keys)
return valid_keys
def get_task(self, task_id: str, brain_id: str) -> Optional[str]:
"""
Get a specific task by ID.
"""
task_data = self.client.hget(f"{brain_id}:_tasks", task_id)
if task_data is None:
return None
try:
data = json.loads(task_data)
return data.get("data")
except (json.JSONDecodeError, TypeError):
return None
_redis_client = RedisClient()
|