diff --git a/backend/prompt_studio/prompt_profile_manager/profile_manager_helper.py b/backend/prompt_studio/prompt_profile_manager/profile_manager_helper.py index 19592b812..04725a16a 100644 --- a/backend/prompt_studio/prompt_profile_manager/profile_manager_helper.py +++ b/backend/prompt_studio/prompt_profile_manager/profile_manager_helper.py @@ -4,7 +4,7 @@ class ProfileManagerHelper: @classmethod - def get_profile_manager(cls, profile_manager_id): + def get_profile_manager(cls, profile_manager_id: str) -> ProfileManager: try: return ProfileManager.objects.get(profile_id=profile_manager_id) except ProfileManager.DoesNotExist: diff --git a/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py index 2468c3135..712d850a4 100644 --- a/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py @@ -835,7 +835,7 @@ def dynamic_indexer( ) from e @staticmethod - def wait_for_document_indexing(doc_id_key): + def wait_for_document_indexing(doc_id_key: str) -> str: max_wait_time = settings.MAX_WAIT_TIME # 30 minutes wait_time = 0 polling_interval = settings.POLLING_INTERVAL # Poll every 5 seconds diff --git a/backend/prompt_studio/prompt_studio_core/redis_utils.py b/backend/prompt_studio/prompt_studio_core/redis_utils.py index e00879502..862c3443d 100644 --- a/backend/prompt_studio/prompt_studio_core/redis_utils.py +++ b/backend/prompt_studio/prompt_studio_core/redis_utils.py @@ -2,26 +2,26 @@ from utils.cache_service import CacheService -def set_document_indexing(doc_id_key, ttl=1800): +def set_document_indexing(doc_id_key: str, ttl: int = 1800) -> None: CacheService.set_key(f"document_indexing:{doc_id_key}", "started", expire=ttl) -def is_document_indexing(doc_id_key): +def is_document_indexing(doc_id_key: str) -> bool: return CacheService.get_key(f"document_indexing:{doc_id_key}") == b"started" -def mark_document_indexed(doc_id_key, doc_id): +def mark_document_indexed(doc_id_key: str, doc_id: str) -> None: CacheService.set_key( f"document_indexing:{doc_id_key}", doc_id, expire=settings.INDEXING_FLAG_TTL ) -def get_indexed_document_id(doc_id_key): +def get_indexed_document_id(doc_id_key: str) -> str | None: result = CacheService.get_key(f"document_indexing:{doc_id_key}") if result and result != b"started": return result return None -def remove_document_indexing(doc_id_key): +def remove_document_indexing(doc_id_key: str) -> None: CacheService.delete_a_key(f"document_indexing:{doc_id_key}")