106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
|
|
import redis.asyncio as aioredis
|
|
|
|
from .config import Settings
|
|
|
|
|
|
class SearchCache:
|
|
def __init__(self, settings: Settings):
|
|
self._url = settings.redis_url
|
|
self._ttl = settings.cache_ttl
|
|
self._client: aioredis.Redis | None = None
|
|
|
|
async def _conn(self) -> aioredis.Redis | None:
|
|
if not self._url:
|
|
return None
|
|
if self._client is None:
|
|
self._client = aioredis.from_url(
|
|
self._url,
|
|
decode_responses=True,
|
|
socket_connect_timeout=5,
|
|
socket_timeout=5,
|
|
)
|
|
return self._client
|
|
|
|
@staticmethod
|
|
def _search_key(query: str, top_k: int, max_chars: int) -> str:
|
|
raw = f"{query}|{top_k}|{max_chars}"
|
|
h = hashlib.sha256(raw.encode()).hexdigest()[:20]
|
|
return f"aisou:search:{h}"
|
|
|
|
async def get_search(self, query: str, top_k: int, max_chars: int) -> dict | None:
|
|
try:
|
|
c = await self._conn()
|
|
if c is None:
|
|
return None
|
|
raw = await c.get(self._search_key(query, top_k, max_chars))
|
|
return json.loads(raw) if raw else None
|
|
except Exception:
|
|
return None
|
|
|
|
async def set_search(self, query: str, top_k: int, max_chars: int, data: dict) -> None:
|
|
try:
|
|
c = await self._conn()
|
|
if c is None:
|
|
return
|
|
await c.setex(
|
|
self._search_key(query, top_k, max_chars),
|
|
self._ttl,
|
|
json.dumps(data, ensure_ascii=False),
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
async def get_job(self, job_id: str) -> dict | None:
|
|
try:
|
|
c = await self._conn()
|
|
if c is None:
|
|
return None
|
|
raw = await c.get(f"aisou:job:{job_id}")
|
|
return json.loads(raw) if raw else None
|
|
except Exception:
|
|
return None
|
|
|
|
async def set_job(self, job_id: str, data: dict, ttl: int = 7200) -> None:
|
|
try:
|
|
c = await self._conn()
|
|
if c is None:
|
|
return
|
|
await c.setex(
|
|
f"aisou:job:{job_id}",
|
|
ttl,
|
|
json.dumps(data, ensure_ascii=False),
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
async def push_history(self, entry: dict) -> None:
|
|
try:
|
|
c = await self._conn()
|
|
if c is None:
|
|
return
|
|
await c.lpush("aisou:history", json.dumps(entry, ensure_ascii=False))
|
|
await c.ltrim("aisou:history", 0, 99)
|
|
except Exception:
|
|
pass
|
|
|
|
async def get_history(self, limit: int = 10) -> list[dict]:
|
|
try:
|
|
c = await self._conn()
|
|
if c is None:
|
|
return []
|
|
items = await c.lrange("aisou:history", 0, limit - 1)
|
|
result: list[dict] = []
|
|
for item in items:
|
|
try:
|
|
result.append(json.loads(item))
|
|
except Exception:
|
|
pass
|
|
return result
|
|
except Exception:
|
|
return []
|