AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-15 01:15:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: optimize memory & invoke error output (#1931)

Browse Source
This commit is contained in:
takatost 2024-01-05 03:47:46 +08:00 committed by GitHub
parent 97c972f14d
commit faf40a42bc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
api/core/model_runtime/errors/invoke.py
View File

@ -8,6 +8,9 @@ class InvokeError(Exception):
def __init__(self, description: Optional[str] = None) -> None: def __init__(self, description: Optional[str] = None) -> None:
self.description = description self.description = description
def __str__(self):
return self.description or self.__class__.__name__
class InvokeConnectionError(InvokeError): class InvokeConnectionError(InvokeError):
"""Raised when the Invoke returns connection error.""" """Raised when the Invoke returns connection error."""

11
api/libs/rsa.py
View File

@ -5,7 +5,6 @@ from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.PublicKey import RSA from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes from Crypto.Random import get_random_bytes
from core.helper.lru_cache import LRUCache
from extensions.ext_redis import redis_client from extensions.ext_redis import redis_client
from extensions.ext_storage import storage from extensions.ext_storage import storage
@ -46,15 +45,7 @@ def encrypt(text, public_key):
return prefix_hybrid + encrypted_data return prefix_hybrid + encrypted_data
tenant_rsa_keys = LRUCache(capacity=1000)
def get_decrypt_decoding(tenant_id): def get_decrypt_decoding(tenant_id):
rsa_key = tenant_rsa_keys.get(tenant_id)
if rsa_key:
cipher_rsa = PKCS1_OAEP.new(rsa_key)
return rsa_key, cipher_rsa
filepath = "privkeys/{tenant_id}".format(tenant_id=tenant_id) + "/private.pem" filepath = "privkeys/{tenant_id}".format(tenant_id=tenant_id) + "/private.pem"
cache_key = 'tenant_privkey:{hash}'.format(hash=hashlib.sha3_256(filepath.encode()).hexdigest()) cache_key = 'tenant_privkey:{hash}'.format(hash=hashlib.sha3_256(filepath.encode()).hexdigest())
@ -70,8 +61,6 @@ def get_decrypt_decoding(tenant_id):
rsa_key = RSA.import_key(private_key) rsa_key = RSA.import_key(private_key)
cipher_rsa = PKCS1_OAEP.new(rsa_key) cipher_rsa = PKCS1_OAEP.new(rsa_key)
tenant_rsa_keys.put(tenant_id, rsa_key)
return rsa_key, cipher_rsa return rsa_key, cipher_rsa