agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/cache/redis_cache.py

@@ -0,0 +1,385 @@
+import os
+import asyncio
+import logging
+import weakref
+import pickle
+import time
+from typing import Dict, Any, Optional, List
+from uuid import uuid4
+from dataclasses import dataclass
+import redis.asyncio as redis
+from agent_api_server.configs import global_config
+from redis.exceptions import RedisError, WatchError
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class ThreadState:
+    thread_id: str
+    graph_name: str
+    status: str
+    created_at: float = time.time()
+    last_accessed: float = time.time()
+    metadata: Dict[str, Any] = None
+
+class AsyncRedisThreadStorage:
+    _instances = weakref.WeakSet()
+    _shutdown_lock = asyncio.Lock()
+    _worker_instances = weakref.WeakValueDictionary()
+
+    def __init__(self, max_retries: int = 3, retry_delay: float = 1.0):
+        self.worker_pid = os.getpid()
+        logger.info(
+            f"Initializing Redis storage for worker {self.worker_pid} "
+            f"(max_retries={max_retries}, retry_delay={retry_delay:.1f}s)"
+        )
+
+        self.redis_url = self._get_redis_url()
+        self.redis: Optional[redis.Redis] = None
+        self.max_retries = max(max_retries, 1)
+        self.retry_delay = max(retry_delay, 0.1)
+        self._is_initialized = False
+        self._lock = asyncio.Lock()
+
+        AsyncRedisThreadStorage._instances.add(self)
+        AsyncRedisThreadStorage._worker_instances[self.worker_pid] = self
+        logger.debug(f"New Redis storage instance registered for worker {self.worker_pid}")
+
+    @classmethod
+    def get_worker_instance(cls) -> 'AsyncRedisThreadStorage':
+        """Get or create instance for current worker process"""
+        worker_pid = os.getpid()
+        if worker_pid not in cls._worker_instances:
+            instance = cls()
+            cls._worker_instances[worker_pid] = instance
+            logger.debug(f"Created new Redis storage instance for worker {worker_pid}")
+        return cls._worker_instances[worker_pid]
+
+    @classmethod
+    async def close_worker_instance(cls):
+        """Close instance for current worker process"""
+        worker_pid = os.getpid()
+        if worker_pid in cls._worker_instances:
+            instance = cls._worker_instances.pop(worker_pid)
+            await instance.close()
+
+    @staticmethod
+    def _get_redis_url() -> str:
+        """获取安全的Redis连接字符串（隐藏密码）"""
+        redis_url = global_config.REDIS_URL
+        safe_redis_url = redis_url.split('@')[0] + '@[REDACTED]' if '@' in redis_url else redis_url
+        logger.debug("Using Redis URL: %s", safe_redis_url)
+        return redis_url
+
+    @classmethod
+    async def close_all(cls):
+        logger.info("Initiating shutdown of all Redis storage instances")
+
+        async with cls._shutdown_lock:
+            instances = list(cls._instances)
+            if not instances:
+                logger.debug("No active Redis instances to close")
+                return
+
+            logger.debug("Closing %d active Redis instance(s)", len(instances))
+            for instance in instances:
+                try:
+                    await instance.close()
+                    logger.debug("Redis instance closed successfully")
+                except Exception as e:
+                    logger.warning(
+                        "Error closing Redis instance: %s",
+                        str(e),
+                        exc_info=logger.isEnabledFor(logging.DEBUG)
+                    )
+
+    async def initialize(self) -> None:
+        """Initialize Redis connection"""
+        async with self._lock:
+            if self._is_initialized:
+                logger.debug("Redis already initialized, skipping")
+                return
+
+            logger.info("Starting Redis connection initialization")
+
+            try:
+                self.redis = redis.Redis.from_url(
+                    self.redis_url,
+                )
+
+                await self.redis.ping()
+                self._is_initialized = True
+                logger.info("Redis storage initialized successfully")
+
+            except Exception as e:
+                logger.error("Redis initialization failed", exc_info=True)
+                await self._safe_close()
+                raise RuntimeError("Redis storage initialization failed") from e
+
+    @staticmethod
+    def _redis_key(thread_id: str) -> str:
+        return f"langgraph:thread:{thread_id}"
+
+    async def _save_to_redis(self, state: ThreadState) -> bool:
+        try:
+            key = self._redis_key(state.thread_id)
+            data = {
+                "graph_name": state.graph_name,
+                "status": state.status,
+                "created_at": state.created_at,
+                "last_accessed": state.last_accessed,
+                "metadata": pickle.dumps(state.metadata) if state.metadata else None,
+                "worker_pid": str(self.worker_pid)
+            }
+
+            async with self.redis.pipeline(transaction=True) as pipe:
+                await pipe.hset(key, mapping={k: v for k, v in data.items() if v is not None})
+                await pipe.execute()
+            return True
+        except RedisError as e:
+            logger.error(f"Failed to save thread {state.thread_id} to Redis: {str(e)}")
+            return False
+
+    async def _load_from_redis(self, thread_id: str) -> Optional[ThreadState]:
+        try:
+            key = self._redis_key(thread_id)
+            data = await self.redis.hgetall(key)
+            if not data:
+                return None
+
+            metadata = pickle.loads(data[b'metadata']) if data.get(b'metadata') else None
+
+            return ThreadState(
+                thread_id=thread_id,
+                graph_name=data[b'graph_name'].decode(),
+                status=data[b'status'].decode(),
+                created_at=float(data[b'created_at']),
+                last_accessed=float(data[b'last_accessed']),
+                metadata=metadata
+            )
+        except RedisError as e:
+            logger.error(f"Failed to load thread {thread_id} from Redis: {str(e)}")
+            return None
+        except Exception as e:
+            logger.error(f"Failed to deserialize thread {thread_id}: {str(e)}")
+            return None
+
+    async def create_thread_with_id(self, thread_id, graph_name: str) -> ThreadState:
+        async with self._lock:
+            state = ThreadState(
+                thread_id=thread_id,
+                graph_name=graph_name,
+                status="created"
+            )
+
+            if not await self._save_to_redis(state):
+                raise RuntimeError("Failed to persist thread state")
+
+            logger.debug(f"Created thread {thread_id} for graph {graph_name}")
+            return state
+
+    async def create_thread(self, graph_name: str) -> ThreadState:
+        async with self._lock:
+            thread_id = str(uuid4())
+            state = ThreadState(
+                thread_id=thread_id,
+                graph_name=graph_name,
+                status="created"
+            )
+
+            if not await self._save_to_redis(state):
+                raise RuntimeError("Failed to persist thread state")
+
+            logger.debug(f"Created thread {thread_id} for graph {graph_name}")
+            return state
+
+    async def get_thread(self, thread_id: str) -> Optional[ThreadState]:
+        state = await self._load_from_redis(thread_id)
+        if state:
+            # Update last_accessed time
+            state.last_accessed = time.time()
+            await self._save_to_redis(state)
+            return state
+        return None
+
+    async def update_thread(self, thread_id: str, **updates) -> bool:
+        async with self._lock:
+            state = await self._load_from_redis(thread_id)
+            if not state:
+                return False
+
+            for k, v in updates.items():
+                setattr(state, k, v)
+            state.last_accessed = time.time()
+
+            return await self._save_to_redis(state)
+
+    async def delete_thread(self, thread_id: str) -> bool:
+        """异步删除线程"""
+        async with self._lock:
+            try:
+                await self.redis.delete(self._redis_key(thread_id))
+                logger.debug(f"Deleted thread {thread_id}")
+                return True
+            except RedisError as e:
+                logger.error(f"Failed to delete thread {thread_id} from Redis: {str(e)}")
+                return False
+
+    async def _scan_redis_keys(self, pattern: str = "langgraph:thread:*") -> List[str]:
+        """异步扫描Redis键"""
+        try:
+            cursor = '0'
+            keys = []
+            while cursor != 0:
+                cursor, partial_keys = await self.redis.scan(
+                    cursor=cursor,
+                    match=pattern,
+                    count=1000
+                )
+                keys.extend(partial_keys)
+            return [key.decode() for key in keys]
+        except RedisError as e:
+            logger.error(f"Failed to scan Redis keys: {str(e)}")
+            return []
+
+    async def list_threads(self) -> Dict[str, ThreadState]:
+        """异步列出线程"""
+        result = {}
+        keys = await self._scan_redis_keys()
+
+        for key in keys:
+            thread_id = key.split(":")[-1]
+            state = await self._load_from_redis(thread_id)
+            if state:
+                result[thread_id] = state
+
+        return result
+
+    async def get_stats(self) -> Dict[str, Any]:
+        """异步获取统计信息"""
+        try:
+            redis_count = len(await self._scan_redis_keys())
+        except RedisError:
+            redis_count = -1
+
+        return {
+            "redis_thread_count": redis_count,
+        }
+
+    async def _safe_close(self) -> None:
+        logger.debug("Starting Redis safe cleanup")
+        if self.redis:
+            try:
+                await self.redis.close()
+            except Exception as e:
+                logger.warning(f"Error closing Redis connection: {str(e)}", exc_info=True)
+
+        self.redis = None
+        self._is_initialized = False
+        logger.info("Redis resources cleaned up")
+
+    async def close(self) -> None:
+        async with self._lock:
+            if not self._is_initialized:
+                logger.debug("Redis already closed, skipping")
+                return
+
+            logger.info("Starting Redis graceful shutdown")
+
+            try:
+                await self._safe_close()
+            except Exception as e:
+                logger.error(
+                    "Error during Redis shutdown: %s",
+                    str(e),
+                    exc_info=logger.isEnabledFor(logging.DEBUG)
+                )
+            finally:
+                AsyncRedisThreadStorage._instances.discard(self)
+                logger.debug("Redis instance unregistered")
+
+    async def __aenter__(self) -> 'AsyncRedisThreadStorage':
+        logger.debug("Entering Redis context manager")
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        logger.debug("Exiting Redis context manager")
+        await self.close()
+
+    @property
+    def is_initialized(self):
+        return self._is_initialized
+
+    async def clear_stop_flag(self, thread_id: str) -> bool:
+        """Clear the stop flag for a thread.
+
+        Args:
+            thread_id: The ID of the thread to clear the stop flag for
+
+        Returns:
+            bool: True if the flag was cleared successfully, False otherwise
+        """
+        key = self._redis_key(thread_id)
+        try:
+            async with self.redis.pipeline() as pipe:
+                while True:
+                    try:
+                        await pipe.watch(key)
+                        existing = await pipe.exists(key)
+                        if not existing:
+                            await pipe.unwatch()
+                            return False
+
+                        pipe.multi()
+                        await pipe.hdel(key, "stop_requested")
+                        await pipe.hdel(key, "stop_requested_at")
+                        await pipe.hset(key, "status", "complete")
+                        await pipe.execute()
+                        return True
+                    except WatchError:
+                        continue
+                    finally:
+                        await pipe.unwatch()
+        except RedisError as e:
+            logger.error(f"Failed to clear stop flag for thread {thread_id}: {str(e)}")
+            return False
+
+    async def set_stop_flag(self, thread_id: str) -> bool:
+        key = self._redis_key(thread_id)
+        try:
+            async with self.redis.pipeline() as pipe:
+                while True:
+                    try:
+                        await pipe.watch(key)
+                        existing = await pipe.hget(key, "status")
+                        if existing is None:
+                            await pipe.unwatch()
+                            return False
+
+                        pipe.multi()
+                        await pipe.hset(key, "status", "stopping")
+                        await pipe.hset(key, "stop_requested", "1")
+                        await pipe.hset(key, "stop_requested_at", str(time.time()))
+                        await pipe.execute()
+                        return True
+                    except WatchError:
+                        continue
+                    finally:
+                        await pipe.unwatch()
+        except RedisError as e:
+            logger.error(f"Failed to set stop flag: {str(e)}")
+            return False
+
+    async def should_stop(self, thread_id: str) -> bool:
+        """检查是否需要停止"""
+        key = self._redis_key(thread_id)
+        try:
+            async with self.redis.pipeline() as pipe:
+                await pipe.hexists(key, "stop_requested")
+                await pipe.hget(key, "status")
+                exists, status = await pipe.execute()
+                return bool(exists) and status != b"stopped"
+        except RedisError as e:
+            logger.error(f"Failed to check stop flag: {str(e)}")
+            return False

agent_api_server/callback_handler.py

@@ -0,0 +1,18 @@
+
+_callbacks = []
+
+def register_callback(callback):
+    """
+    注册一个回调函数
+    """
+    _callbacks.append(callback)
+
+def notify_callbacks(**kwargs):
+    """
+    当模型配置更新时，通知所有注册的回调
+    """
+    for callback in _callbacks:
+        try:
+            callback(**kwargs)
+        except Exception as e:
+            print(f"[ERROR] Callback execution failed: {e}")