| import json |
| from enum import Enum |
| from json import JSONDecodeError |
| from typing import Optional |
|
|
| from extensions.ext_redis import redis_client |
|
|
|
|
| class ToolParameterCacheType(Enum): |
| PARAMETER = "tool_parameter" |
|
|
|
|
| class ToolParameterCache: |
| def __init__( |
| self, tenant_id: str, provider: str, tool_name: str, cache_type: ToolParameterCacheType, identity_id: str |
| ): |
| self.cache_key = ( |
| f"{cache_type.value}_secret:tenant_id:{tenant_id}:provider:{provider}:tool_name:{tool_name}" |
| f":identity_id:{identity_id}" |
| ) |
|
|
| def get(self) -> Optional[dict]: |
| """ |
| Get cached model provider credentials. |
| |
| :return: |
| """ |
| cached_tool_parameter = redis_client.get(self.cache_key) |
| if cached_tool_parameter: |
| try: |
| cached_tool_parameter = cached_tool_parameter.decode("utf-8") |
| cached_tool_parameter = json.loads(cached_tool_parameter) |
| except JSONDecodeError: |
| return None |
|
|
| return cached_tool_parameter |
| else: |
| return None |
|
|
| def set(self, parameters: dict) -> None: |
| """ |
| Cache model provider credentials. |
| |
| :param credentials: provider credentials |
| :return: |
| """ |
| redis_client.setex(self.cache_key, 86400, json.dumps(parameters)) |
|
|
| def delete(self) -> None: |
| """ |
| Delete cached model provider credentials. |
| |
| :return: |
| """ |
| redis_client.delete(self.cache_key) |
|
|
