Implement Redis caching (#12)

.env

@@ -4,4 +4,5 @@
 # https://devcenter.heroku.com/articles/config-vars
 
 OPENAI_API_KEY=<your-openai-api-key>
-GOOGLE_CREDENTIALS=<your-google-credentials>
+GOOGLE_CREDENTIALS=
+REDIS_URL=redis://localhost:6379/1

.pre-commit-config.yaml

@@ -23,3 +23,4 @@ repos:
         args:
           - --no-warn-no-return
           - --ignore-missing-imports
+        additional_dependencies: [redis==5.0.1]

chatgpt_plugins/market_right_side/openapi.yaml

@@ -6,24 +6,6 @@ info:
 servers:
   - url: https://admin.marketrightside.com/api
 paths:
-  /oauth/token:
-    post:
-      operationId: getAuthToken
-      summary: Obtain an authorization token
-      responses:
-        "200":
-          description: Returns an authorization token
-          content:
-            application/json:
-              schema:
-                type: object
-                properties:
-                  token_type:
-                    type: string
-                  expires_in:
-                    type: integer
-                  access_token:
-                    type: string
   /public/market-list:
     get:
       operationId: getMarketList

nalgonda/agency_manager.py

@@ -5,6 +5,7 @@
 
 from agency_swarm import Agency, Agent
 
+from nalgonda.caching.redis_cache_manager import RedisCacheManager
 from nalgonda.custom_tools import TOOL_MAPPING
 from nalgonda.models.agency_config import AgencyConfig
 
@@ -13,7 +14,7 @@
 
 class AgencyManager:
     def __init__(self) -> None:
-        self.cache: dict[str, Agency] = {}  # Mapping from agency_id+thread_id to Agency class instance
+        self.cache_manager = RedisCacheManager()
         self.lock = asyncio.Lock()
 
     async def create_agency(self, agency_id: str | None = None) -> tuple[Agency, str]:
@@ -23,26 +24,30 @@ async def create_agency(self, agency_id: str | None = None) -> tuple[Agency, str
         async with self.lock:
             # Note: Async-to-Sync Bridge
             agency = await asyncio.to_thread(self.load_agency_from_config, agency_id)
-            self.cache[agency_id] = agency
+            await self.cache_agency(agency, agency_id, None)
             return agency, agency_id
 
     async def get_agency(self, agency_id: str, thread_id: str | None) -> Agency | None:
         """Get the agency from the cache."""
-        async with self.lock:
-            return self.cache.get(self.get_cache_key(agency_id, thread_id))
+        cache_key = self.get_cache_key(agency_id, thread_id)
+        agency = await self.cache_manager.get(cache_key)
+        return agency
 
     async def cache_agency(self, agency: Agency, agency_id: str, thread_id: str | None) -> None:
-        """Cache the agency for the given agency ID and thread ID."""
-        async with self.lock:
-            cache_key = self.get_cache_key(agency_id, thread_id)
-            self.cache[cache_key] = agency
+        """Cache the agency."""
+        cache_key = self.get_cache_key(agency_id, thread_id)
+        await self.cache_manager.set(cache_key, agency)
 
     async def delete_agency_from_cache(self, agency_id: str, thread_id: str | None) -> None:
-        async with self.lock:
-            cache_key = self.get_cache_key(agency_id, thread_id)
-            self.cache.pop(cache_key, None)
+        """Delete the agency from the cache."""
+        cache_key = self.get_cache_key(agency_id, thread_id)
+        await self.cache_manager.delete(cache_key)
 
     async def refresh_thread_id(self, agency: Agency, agency_id: str, thread_id: str | None) -> str | None:
+        """Refresh the thread ID for the given agency.
+        If the thread ID has changed, update the cache and return the new thread ID.
+        Otherwise, return None.
+        """
         new_thread_id = agency.main_thread.id
         if thread_id != new_thread_id:
             await self.cache_agency(agency, agency_id, new_thread_id)

nalgonda/caching/cache_manager.py

@@ -0,0 +1,22 @@
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class CacheManager(Generic[T]):
+    """
+    Abstract base class for a cache manager to handle caching operations.
+    Specific cache backends should extend this class and implement its methods.
+    """
+
+    async def get(self, key: str) -> T | None:
+        raise NotImplementedError()
+
+    async def set(self, key: str, value: T) -> None:
+        raise NotImplementedError()
+
+    async def delete(self, key: str) -> None:
+        raise NotImplementedError()
+
+    async def close(self) -> None:
+        raise NotImplementedError()

nalgonda/caching/redis_cache_manager.py

@@ -0,0 +1,40 @@
+import asyncio
+import pickle
+
+from agency_swarm import Agency
+from redis import asyncio as aioredis
+
+from nalgonda.caching.cache_manager import CacheManager
+from nalgonda.settings import settings
+
+
+class RedisCacheManager(CacheManager):
+    """Redis cache manager
+    This class implements the CacheManager interface using Redis as the cache backend.
+    """
+
+    def __init__(self):
+        """Initializes the Redis cache manager"""
+        self.redis = aioredis.from_url(str(settings.redis_dsn), decode_responses=True)
+
+    def __del__(self):
+        """Wait for the Redis connection to close"""
+        asyncio.run(self.close())
+
+    async def get(self, key: str) -> Agency | None:
+        """Gets the value for the given key from the cache"""
+        serialized_data = await self.redis.get(key)
+        return pickle.loads(serialized_data) if serialized_data else None
+
+    async def set(self, key: str, value: Agency) -> None:
+        """Sets the value for the given key in the cache"""
+        serialized_data = pickle.dumps(value)
+        await self.redis.set(key, serialized_data)
+
+    async def delete(self, key: str) -> None:
+        """Deletes the value for the given key from the cache"""
+        await self.redis.delete(key)
+
+    async def close(self) -> None:
+        """Closes the Redis connection"""
+        await self.redis.close()

nalgonda/settings.py

@@ -1,15 +1,19 @@
-from pydantic import Field
+from pydantic import AliasChoices, Field, RedisDsn
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 LATEST_GPT_MODEL = "gpt-4-1106-preview"
 
 
 class Settings(BaseSettings):
-    openai_api_key: str | None = Field(default=None, validation_alias="OPENAI_API_KEY")
-    gpt_model: str = Field(default=LATEST_GPT_MODEL, validation_alias="GPT_MODEL")
-    google_credentials: str | None = Field(default=None, validation_alias="GOOGLE_CREDENTIALS")
+    openai_api_key: str | None = Field(default=None)
+    gpt_model: str = Field(default=LATEST_GPT_MODEL)
+    google_credentials: str | None = Field(default=None)
+    redis_dsn: RedisDsn = Field(
+        "redis://localhost:6379/1",
+        validation_alias=AliasChoices("service_redis_dsn", "redis_url"),
+    )
 
-    model_config = SettingsConfigDict(env_file=".env", case_sensitive=True)
+    model_config = SettingsConfigDict(env_file=".env")
 
 
 settings = Settings()

pyproject.toml

@@ -19,6 +19,7 @@ instructor = "^0.3.4"
 openai = "^1.3.0"
 pydantic = "^2.5.2"
 pydantic-settings = "^2.1.0"
+redis = "^5.0.1"
 uvicorn = {extras = ["standard"], version = "^0.24.0.post1"}
 # a dependency from a GitHub repository
 agency-swarm = { git = "https://github.com/VRSEN/agency-swarm.git", branch = "main" }