feat: add cache layer for llm (#981)

Author: earayu

aperag/app.py

@@ -15,7 +15,7 @@
 from fastapi import FastAPI
 
 from aperag.exception_handlers import register_exception_handlers
-from aperag.llm.litellm_track import register_opik_llm_track
+from aperag.llm.litellm_track import register_custom_llm_track
 from aperag.views.api_key import router as api_key_router
 from aperag.views.audit import router as audit_router
 from aperag.views.auth import router as auth_router
@@ -34,7 +34,7 @@
 # Register global exception handlers
 register_exception_handlers(app)
 
-register_opik_llm_track()
+register_custom_llm_track()
 
 app.include_router(auth_router, prefix="/api/v1")
 app.include_router(main_router, prefix="/api/v1")

aperag/db/redis_manager.py

@@ -16,68 +16,114 @@
 Redis connection manager for the application.
 
 This module provides a simple and efficient Redis connection management system
-using redis-py's built-in connection pooling capabilities.
+using redis-py's built-in connection pooling capabilities for both sync and async operations.
 """
 
 import logging
 from typing import Optional
 
-import redis.asyncio as redis
+import redis
+import redis.asyncio as async_redis
 
 logger = logging.getLogger(__name__)
 
 
 class RedisConnectionManager:
     """
-    Simple Redis connection manager using redis-py's built-in connection pooling.
+    Redis connection manager supporting both sync and async operations.
 
-    This provides a shared connection pool for the entire application, avoiding
+    This provides shared connection pools for the entire application, avoiding
     the overhead of creating multiple connections for different Redis operations.
 
     Features:
-    - Automatic connection pooling with redis-py
+    - Automatic connection pooling with redis-py (both sync and async)
     - Configurable pool size and timeouts
     - Global shared instance for efficiency
     - Proper cleanup handling
     """
 
     _instance: Optional["RedisConnectionManager"] = None
-    _client: Optional[redis.Redis] = None
-    _pool: Optional[redis.ConnectionPool] = None
+    _async_client: Optional[async_redis.Redis] = None
+    _sync_client: Optional[redis.Redis] = None
+    _async_pool: Optional[async_redis.ConnectionPool] = None
+    _sync_pool: Optional[redis.ConnectionPool] = None
 
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super().__new__(cls)
         return cls._instance
 
     @classmethod
-    async def get_client(cls, redis_url: str = None) -> redis.Redis:
+    async def get_async_client(cls, redis_url: str = None) -> async_redis.Redis:
         """
-        Get Redis client with shared connection pool.
+        Get async Redis client with shared connection pool.
 
         Args:
             redis_url: Redis connection URL. If None, will use from config.
 
         Returns:
-            Redis client instance with shared connection pool
+            Async Redis client instance with shared connection pool
         """
-        if cls._client is None:
-            await cls._initialize_client(redis_url)
-        return cls._client
+        if cls._async_client is None:
+            await cls._initialize_async_client(redis_url)
+        return cls._async_client
 
     @classmethod
-    async def _initialize_client(cls, redis_url: str = None):
-        """Initialize Redis client with connection pool."""
+    def get_sync_client(cls, redis_url: str = None) -> redis.Redis:
+        """
+        Get sync Redis client with shared connection pool.
+
+        Args:
+            redis_url: Redis connection URL. If None, will use from config.
+
+        Returns:
+            Sync Redis client instance with shared connection pool
+        """
+        if cls._sync_client is None:
+            cls._initialize_sync_client(redis_url)
+        return cls._sync_client
+
+    @classmethod
+    async def _initialize_async_client(cls, redis_url: str = None):
+        """Initialize async Redis client with connection pool."""
         if redis_url is None:
-            # Import here to avoid circular imports
-            from aperag.config import settings
+            redis_url = cls._get_redis_url()
+
+        logger.debug(f"Initializing async Redis connection pool: {redis_url}")
+
+        # Create async connection pool
+        cls._async_pool = async_redis.ConnectionPool.from_url(
+            redis_url,
+            max_connections=20,  # Pool size
+            encoding="utf-8",
+            decode_responses=True,
+            socket_connect_timeout=5,
+            socket_timeout=5,
+            retry_on_timeout=True,
+            health_check_interval=30,
+        )
+
+        # Create async client using the pool
+        cls._async_client = async_redis.Redis(connection_pool=cls._async_pool)
+
+        # Test connection
+        try:
+            await cls._async_client.ping()
+            logger.debug("Async Redis connection pool initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to connect to async Redis: {e}")
+            raise ConnectionError(f"Cannot connect to async Redis: {e}")
 
-            redis_url = settings.memory_redis_url
+    @classmethod
+    def _initialize_sync_client(cls, redis_url: str = None):
+        """Initialize sync Redis client with connection pool."""
+        if redis_url is None:
+            redis_url = cls._get_redis_url()
 
-        logger.debug(f"Initializing Redis connection pool: {redis_url}")
+        logger.debug(f"Initializing sync Redis connection pool: {redis_url}")
 
-        # Create connection pool with reasonable defaults
-        cls._pool = redis.ConnectionPool.from_url(
+        # Create sync connection pool
+        cls._sync_pool = redis.ConnectionPool.from_url(
             redis_url,
             max_connections=20,  # Pool size
             encoding="utf-8",
@@ -88,46 +134,88 @@ async def _initialize_client(cls, redis_url: str = None):
             health_check_interval=30,
         )
 
-        # Create client using the pool
-        cls._client = redis.Redis(connection_pool=cls._pool)
+        # Create sync client using the pool
+        cls._sync_client = redis.Redis(connection_pool=cls._sync_pool)
 
         # Test connection
         try:
-            await cls._client.ping()
-            logger.debug("Redis connection pool initialized successfully")
+            cls._sync_client.ping()
+            logger.debug("Sync Redis connection pool initialized successfully")
         except Exception as e:
-            logger.error(f"Failed to connect to Redis: {e}")
-            raise ConnectionError(f"Cannot connect to Redis: {e}")
+            logger.error(f"Failed to connect to sync Redis: {e}")
+            raise ConnectionError(f"Cannot connect to sync Redis: {e}")
+
+    @classmethod
+    def _get_redis_url(cls) -> str:
+        """Get Redis URL from configuration."""
+        # Import here to avoid circular imports
+        from aperag.config import settings
+
+        return settings.MEMORY_REDIS_URL
 
     @classmethod
     async def close(cls):
-        """Close Redis connection pool and clean up resources."""
-        if cls._client:
-            logger.debug("Closing Redis connection pool")
-            await cls._client.close()
-            cls._client = None
-            cls._pool = None
-            logger.debug("Redis connection pool closed")
+        """Close Redis connection pools and clean up resources."""
+        # Close async client
+        if cls._async_client:
+            logger.debug("Closing async Redis connection pool")
+            await cls._async_client.close()
+            cls._async_client = None
+            cls._async_pool = None
+
+        # Close sync client
+        if cls._sync_client:
+            logger.debug("Closing sync Redis connection pool")
+            cls._sync_client.close()
+            cls._sync_client = None
+            cls._sync_pool = None
+
+        logger.debug("All Redis connection pools closed")
 
     @classmethod
     def get_pool_info(cls) -> dict:
         """Get connection pool information for monitoring."""
-        if cls._pool:
-            return {
-                "max_connections": cls._pool.max_connections,
-                "created_connections": cls._pool.created_connections,
-                "available_connections": len(cls._pool._available_connections),
-                "in_use_connections": len(cls._pool._in_use_connections),
+        info = {}
+
+        if cls._async_pool:
+            info["async_pool"] = {
+                "max_connections": cls._async_pool.max_connections,
+                "created_connections": cls._async_pool.created_connections,
+                "available_connections": len(cls._async_pool._available_connections),
+                "in_use_connections": len(cls._async_pool._in_use_connections),
+            }
+
+        if cls._sync_pool:
+            info["sync_pool"] = {
+                "max_connections": cls._sync_pool.max_connections,
+                "created_connections": cls._sync_pool.created_connections,
+                "available_connections": len(cls._sync_pool._available_connections),
+                "in_use_connections": len(cls._sync_pool._in_use_connections),
             }
-        return {"status": "not_initialized"}
+
+        if not info:
+            info["status"] = "not_initialized"
+
+        return info
 
 
-# Convenience functions for backward compatibility with history.py
-async def get_async_redis_client() -> redis.Redis:
+# Convenience functions for backward compatibility
+async def get_async_redis_client() -> async_redis.Redis:
     """Get async Redis client - backward compatible with history.py"""
-    return await RedisConnectionManager.get_client()
+    return await RedisConnectionManager.get_async_client()
+
+
+def get_sync_redis_client() -> redis.Redis:
+    """Get sync Redis client for cache and other sync operations."""
+    return RedisConnectionManager.get_sync_client()
 
 
 def get_redis_connection_manager() -> RedisConnectionManager:
     """Get the Redis connection manager instance."""
     return RedisConnectionManager()
+
+
+# Legacy compatibility - keep the old function name
+async def get_client(redis_url: str = None) -> async_redis.Redis:
+    """Legacy function name - use get_async_client instead."""
+    return await RedisConnectionManager.get_async_client(redis_url)

aperag/flow/runners/rerank.py

@@ -114,7 +114,7 @@ async def run(self, ui: RerankInput, si: SystemInput) -> Tuple[RerankOutput, dic
             )
 
             # Perform reranking
-            result = await rerank_service.rank(query, docs)
+            result = await rerank_service.async_rerank(query, docs)
             logger.info(f"Successfully reranked {len(result)} documents")
 
         except (InvalidConfigurationError, ProviderNotFoundError) as e: