Updated type hints

Author: IgorChvyrov-sm

src/conductor/client/event/event_client.py

@@ -1,3 +1,5 @@
+from typing import Dict
+
 from conductor.client.event.queue.queue_configuration import QueueConfiguration
 from conductor.client.http.api import EventResourceApi
 from conductor.client.http.api_client import ApiClient
@@ -13,13 +15,13 @@ def delete_queue_configuration(self, queue_configuration: QueueConfiguration) ->
             queue_type=queue_configuration.queue_type,
         )
 
-    def get_kafka_queue_configuration(self, queue_topic: str) -> QueueConfiguration:
+    def get_kafka_queue_configuration(self, queue_topic: str) -> Dict[str, object]:
         return self.get_queue_configuration(
             queue_type="kafka",
             queue_name=queue_topic,
         )
 
-    def get_queue_configuration(self, queue_type: str, queue_name: str) -> QueueConfiguration:
+    def get_queue_configuration(self, queue_type: str, queue_name: str) -> Dict[str, object]:
         return self.client.get_queue_config(queue_type, queue_name)
 
     def put_queue_configuration(self, queue_configuration: QueueConfiguration):

src/conductor/client/orkes/orkes_authorization_client.py

@@ -4,6 +4,7 @@
 
 from conductor.client.authorization_client import AuthorizationClient
 from conductor.client.configuration.configuration import Configuration
+from conductor.client.http.models import ExtendedConductorApplication
 from conductor.client.http.models.authorization_request import AuthorizationRequest
 from conductor.client.http.models.conductor_application import ConductorApplication
 from conductor.client.http.models.conductor_user import ConductorUser
@@ -44,7 +45,7 @@ def get_app_by_access_key_id(self, access_key_id: str) -> ConductorApplication:
         app_obj = self.applicationResourceApi.get_app_by_access_key_id(access_key_id)
         return self.api_client.deserialize_class(app_obj, "ConductorApplication")
 
-    def list_applications(self) -> List[ConductorApplication]:
+    def list_applications(self) -> List[ExtendedConductorApplication]:
         return self.applicationResourceApi.list_applications()
 
     def update_application(

src/conductor/client/orkes/orkes_integration_client.py

@@ -1,15 +1,15 @@
 from __future__ import absolute_import, annotations
 
-from typing import Dict, List, Optional
+from typing import List, Optional
 
 from conductor.client.codegen.rest import ApiException
 from conductor.client.configuration.configuration import Configuration
+from conductor.client.http.models import MessageTemplate
 from conductor.client.http.models.integration import Integration
 from conductor.client.http.models.integration_api import IntegrationApi
 from conductor.client.http.models.integration_api_update import IntegrationApiUpdate
 from conductor.client.http.models.integration_def import IntegrationDef
 from conductor.client.http.models.integration_update import IntegrationUpdate
-from conductor.client.http.models.prompt_template import PromptTemplate
 from conductor.client.integration_client import IntegrationClient
 from conductor.client.orkes.orkes_base_client import OrkesBaseClient
 
@@ -53,7 +53,7 @@ def get_integration(self, integration_name: str) -> Optional[Integration]:
     def get_integrations(self) -> List[Integration]:
         return self.integrationApi.get_integration_providers()
 
-    def get_integration_provider(self, name: str) -> Optional[IntegrationDef]:
+    def get_integration_provider(self, name: str) -> Optional[Integration]:
         """Get integration provider by name"""
         try:
             return self.integrationApi.get_integration_provider(name)
@@ -64,7 +64,7 @@ def get_integration_provider(self, name: str) -> Optional[IntegrationDef]:
 
     def get_integration_providers(
         self, category: Optional[str] = None, active_only: Optional[bool] = None
-    ) -> List[IntegrationDef]:
+    ) -> List[Integration]:
         """Get all integration providers with optional filtering"""
         kwargs = {}
         if category is not None:
@@ -79,7 +79,7 @@ def get_integration_provider_defs(self) -> List[IntegrationDef]:
 
     def get_prompts_with_integration(
         self, ai_integration: str, model_name: str
-    ) -> List[PromptTemplate]:
+    ) -> List[MessageTemplate]:
         return self.integrationApi.get_prompts_with_integration(ai_integration, model_name)
 
     def save_integration_api(self, integration_name, api_name, api_details: IntegrationApiUpdate):
@@ -127,15 +127,15 @@ def get_tags_for_integration_provider(self, name):
     # Utility Methods for Integration Provider Management
     def get_integration_provider_by_category(
         self, category: str, active_only: bool = True
-    ) -> List[IntegrationDef]:
+    ) -> List[Integration]:
         """Get integration providers filtered by category"""
         return self.get_integration_providers(category=category, active_only=active_only)
 
-    def get_active_integration_providers(self) -> List[IntegrationDef]:
+    def get_active_integration_providers(self) -> List[Integration]:
         """Get only active integration providers"""
         return self.get_integration_providers(active_only=True)
 
-    def get_integration_available_apis(self, name: str) -> List[IntegrationApi]:
+    def get_integration_available_apis(self, name: str) -> List[str]:
         """Get available APIs for an integration"""
         return self.integrationApi.get_integration_available_apis(name)
 
@@ -156,7 +156,7 @@ def get_all_integrations(
 
     def get_providers_and_integrations(
         self, integration_type: Optional[str] = None, active_only: Optional[bool] = None
-    ) -> Dict[str, object]:
+    ) -> List[str]:
         """Get providers and integrations together"""
         kwargs = {}
         if integration_type is not None:

src/conductor/client/orkes/orkes_metadata_client.py

@@ -10,19 +10,25 @@
 from conductor.client.orkes.models.metadata_tag import MetadataTag
 from conductor.client.orkes.models.ratelimit_tag import RateLimitTag
 from conductor.client.orkes.orkes_base_client import OrkesBaseClient
+from conductor.client.http.models.extended_task_def import ExtendedTaskDef
+from conductor.client.http.models.extended_workflow_def import ExtendedWorkflowDef
 
 
 class OrkesMetadataClient(OrkesBaseClient, MetadataClient):
     def __init__(self, configuration: Configuration):
         super(OrkesMetadataClient, self).__init__(configuration)
 
-    def register_workflow_def(self, workflow_def: WorkflowDef, overwrite: Optional[bool] = True):
-        self.metadataResourceApi.create(workflow_def, overwrite=overwrite)
+    def register_workflow_def(
+        self, workflow_def: ExtendedWorkflowDef, overwrite: Optional[bool] = True
+    ) -> object:
+        return self.metadataResourceApi.create(workflow_def, overwrite=overwrite)
 
-    def update_workflow_def(self, workflow_def: WorkflowDef, overwrite: Optional[bool] = True):
-        self.metadataResourceApi.update([workflow_def], overwrite=overwrite)
+    def update_workflow_def(
+        self, workflow_def: ExtendedWorkflowDef, overwrite: Optional[bool] = True
+    ) -> object:
+        return self.metadataResourceApi.update([workflow_def], overwrite=overwrite)
 
-    def unregister_workflow_def(self, name: str, version: int):
+    def unregister_workflow_def(self, name: str, version: int) -> None:
         self.metadataResourceApi.unregister_workflow_def(name, version)
 
     def get_workflow_def(self, name: str, version: Optional[int] = None) -> WorkflowDef:
@@ -37,16 +43,16 @@ def get_workflow_def(self, name: str, version: Optional[int] = None) -> Workflow
     def get_all_workflow_defs(self) -> List[WorkflowDef]:
         return self.metadataResourceApi.get_workflow_defs()
 
-    def register_task_def(self, task_def: TaskDef):
-        self.metadataResourceApi.register_task_def([task_def])
+    def register_task_def(self, task_def: ExtendedTaskDef) -> object:
+        return self.metadataResourceApi.register_task_def([task_def])
 
     def update_task_def(self, task_def: TaskDef):
         self.metadataResourceApi.update_task_def(task_def)
 
     def unregister_task_def(self, task_type: str):
         self.metadataResourceApi.unregister_task_def(task_type)
 
-    def get_task_def(self, task_type: str) -> TaskDef:
+    def get_task_def(self, task_type: str) -> object:
         return self.metadataResourceApi.get_task_def(task_type)
 
     def get_all_task_defs(self) -> List[TaskDef]:
@@ -59,7 +65,7 @@ def delete_workflow_tag(self, tag: MetadataTag, workflow_name: str):
         tagStr = TagString(tag.key, tag.type, tag.value)
         self.tagsApi.delete_workflow_tag(tagStr, workflow_name)
 
-    def get_workflow_tags(self, workflow_name: str) -> List[MetadataTag]:
+    def get_workflow_tags(self, workflow_name: str) -> object:
         return self.tagsApi.get_workflow_tags(workflow_name)
 
     def set_workflow_tags(self, tags: List[MetadataTag], workflow_name: str):
@@ -72,7 +78,7 @@ def deleteTaskTag(self, tag: MetadataTag, taskName: str):
         tagStr = TagString(tag.key, tag.type, tag.value)
         self.tagsApi.delete_task_tag(tagStr, taskName)
 
-    def getTaskTags(self, taskName: str) -> List[MetadataTag]:
+    def getTaskTags(self, taskName: str) -> object:
         return self.tagsApi.get_task_tags(taskName)
 
     def setTaskTags(self, tags: List[MetadataTag], taskName: str):

src/conductor/client/orkes/orkes_prompt_client.py

@@ -6,7 +6,7 @@
 from conductor.client.configuration.configuration import Configuration
 from conductor.client.http.models.prompt_template import PromptTemplate
 from conductor.client.http.models.prompt_template_test_request import PromptTemplateTestRequest
-from conductor.client.orkes.models.metadata_tag import MetadataTag
+from conductor.client.http.models.tag import Tag
 from conductor.client.orkes.orkes_base_client import OrkesBaseClient
 from conductor.client.prompt_client import PromptClient
 
@@ -32,13 +32,13 @@ def get_prompts(self):
     def delete_prompt(self, prompt_name: str):
         self.promptApi.delete_message_template(prompt_name)
 
-    def get_tags_for_prompt_template(self, prompt_name: str) -> List[MetadataTag]:
+    def get_tags_for_prompt_template(self, prompt_name: str) -> List[Tag]:
         return self.promptApi.get_tags_for_prompt_template(prompt_name)
 
-    def update_tag_for_prompt_template(self, prompt_name: str, tags: List[MetadataTag]):
+    def update_tag_for_prompt_template(self, prompt_name: str, tags: List[Tag]):
         self.promptApi.put_tag_for_prompt_template(tags, prompt_name)
 
-    def delete_tag_for_prompt_template(self, prompt_name: str, tags: List[MetadataTag]):
+    def delete_tag_for_prompt_template(self, prompt_name: str, tags: List[Tag]):
         self.promptApi.delete_tag_for_prompt_template(tags, prompt_name)
 
     def test_prompt(

src/conductor/client/orkes/orkes_scheduler_client.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
-from typing import List, Optional, Tuple
+from typing import List, Optional
 
 from conductor.client.configuration.configuration import Configuration
+from conductor.client.http.models import WorkflowScheduleModel
 from conductor.client.http.models.save_schedule_request import SaveScheduleRequest
 from conductor.client.http.models.search_result_workflow_schedule_execution_model import (
     SearchResultWorkflowScheduleExecutionModel,
 )
 from conductor.client.http.models.workflow_schedule import WorkflowSchedule
 from conductor.client.orkes.models.metadata_tag import MetadataTag
+from conductor.client.http.models.tag import Tag
 from conductor.client.orkes.orkes_base_client import OrkesBaseClient
 from conductor.client.scheduler_client import SchedulerClient
 
@@ -20,10 +22,10 @@ def __init__(self, configuration: Configuration):
     def save_schedule(self, save_schedule_request: SaveScheduleRequest):
         self.schedulerResourceApi.save_schedule(save_schedule_request)
 
-    def get_schedule(self, name: str) -> Tuple[Optional[WorkflowSchedule], str]:
+    def get_schedule(self, name: str) -> WorkflowSchedule:
         return self.schedulerResourceApi.get_schedule(name)
 
-    def get_all_schedules(self, workflow_name: Optional[str] = None) -> List[WorkflowSchedule]:
+    def get_all_schedules(self, workflow_name: Optional[str] = None) -> List[WorkflowScheduleModel]:
         kwargs = {}
         if workflow_name:
             kwargs.update({"workflow_name": workflow_name})
@@ -88,8 +90,8 @@ def requeue_all_execution_records(self):
     def set_scheduler_tags(self, tags: List[MetadataTag], name: str):
         self.schedulerResourceApi.put_tag_for_schedule(tags, name)
 
-    def get_scheduler_tags(self, name: str) -> List[MetadataTag]:
+    def get_scheduler_tags(self, name: str) -> List[Tag]:
         return self.schedulerResourceApi.get_tags_for_schedule(name)
 
-    def delete_scheduler_tags(self, tags: List[MetadataTag], name: str) -> List[MetadataTag]:
+    def delete_scheduler_tags(self, tags: List[MetadataTag], name: str) -> None:
         return self.schedulerResourceApi.delete_tag_for_schedule(tags, name)

src/conductor/client/orkes/orkes_secret_client.py

@@ -1,7 +1,7 @@
 from typing import List, Set
 
 from conductor.client.configuration.configuration import Configuration
-from conductor.client.orkes.models.metadata_tag import MetadataTag
+from conductor.client.http.models.tag import Tag
 from conductor.client.orkes.orkes_base_client import OrkesBaseClient
 from conductor.client.secret_client import SecretClient
 
@@ -25,14 +25,14 @@ def list_secrets_that_user_can_grant_access_to(self) -> List[str]:
     def delete_secret(self, key: str):
         self.secretResourceApi.delete_secret(key)
 
-    def secret_exists(self, key: str) -> bool:
+    def secret_exists(self, key: str) -> object:
         return self.secretResourceApi.secret_exists(key)
 
-    def set_secret_tags(self, tags: List[MetadataTag], key: str):
+    def set_secret_tags(self, tags: List[Tag], key: str):
         self.secretResourceApi.put_tag_for_secret(tags, key)
 
-    def get_secret_tags(self, key: str) -> List[MetadataTag]:
+    def get_secret_tags(self, key: str) -> List[Tag]:
         return self.secretResourceApi.get_tags(key)
 
-    def delete_secret_tags(self, tags: List[MetadataTag], key: str) -> List[MetadataTag]:
+    def delete_secret_tags(self, tags: List[Tag], key: str) -> None:
         return self.secretResourceApi.delete_tag_for_secret(tags, key)

src/conductor/client/orkes/orkes_workflow_client.py

@@ -9,7 +9,6 @@
 from conductor.client.http.models.scrollable_search_result_workflow_summary import (
     ScrollableSearchResultWorkflowSummary,
 )
-from conductor.client.http.models.signal_response import SignalResponse
 from conductor.client.http.models.skip_task_request import SkipTaskRequest
 from conductor.client.http.models.start_workflow_request import StartWorkflowRequest
 from conductor.client.http.models.workflow import Workflow
@@ -70,7 +69,7 @@ def execute_workflow_with_return_strategy(
         wait_for_seconds: int = 30,
         consistency: Optional[str] = None,
         return_strategy: Optional[str] = None,
-    ) -> SignalResponse:
+    ) -> WorkflowRun:
         """Execute a workflow synchronously with optional reactive features
         Args:
             start_workflow_request: StartWorkflowRequest containing workflow details

src/conductor/client/workflow/executor/workflow_executor.py

@@ -7,6 +7,7 @@
 from conductor.client.http.api.metadata_resource_api import MetadataResourceApi
 from conductor.client.http.api.task_resource_api import TaskResourceApi
 from conductor.client.http.api_client import ApiClient
+from conductor.client.http.models import Task
 from conductor.client.http.models.correlation_ids_search_request import CorrelationIdsSearchRequest
 from conductor.client.http.models.rerun_workflow_request import RerunWorkflowRequest
 from conductor.client.http.models.scrollable_search_result_workflow_summary import (
@@ -17,7 +18,7 @@
 from conductor.client.http.models.start_workflow_request import StartWorkflowRequest
 from conductor.client.http.models.task_result import TaskResult
 from conductor.client.http.models.workflow import Workflow
-from conductor.client.http.models.workflow_def import WorkflowDef
+from conductor.client.http.models.extended_workflow_def import ExtendedWorkflowDef
 from conductor.client.http.models.workflow_run import WorkflowRun
 from conductor.client.http.models.workflow_status import WorkflowStatus
 from conductor.client.orkes.orkes_workflow_client import OrkesWorkflowClient
@@ -30,7 +31,9 @@ def __init__(self, configuration: Configuration) -> None:
         self.task_client = TaskResourceApi(api_client)
         self.workflow_client = OrkesWorkflowClient(configuration)
 
-    def register_workflow(self, workflow: WorkflowDef, overwrite: Optional[bool] = None) -> object:
+    def register_workflow(
+        self, workflow: ExtendedWorkflowDef, overwrite: Optional[bool] = None
+    ) -> object:
         """Create a new workflow definition"""
         kwargs = {}
         if overwrite is not None:
@@ -81,7 +84,7 @@ def execute_workflow_with_return_strategy(
         request_id: Optional[str] = None,
         consistency: Optional[str] = None,
         return_strategy: Optional[str] = None,
-    ) -> SignalResponse:
+    ) -> WorkflowRun:
         """Execute a workflow synchronously with optional reactive features"""
         if request_id is None:
             request_id = str(uuid.uuid4())
@@ -291,7 +294,7 @@ def update_task_by_ref_name_sync(
             status=status,
         )
 
-    def get_task(self, task_id: str) -> str:
+    def get_task(self, task_id: str) -> Task:
         """Get task by Id"""
         return self.task_client.get_task(task_id=task_id)