Added type hints for experiment manager (#10851)

Co-authored-by: Ray Myers <ray.myers@gmail.com>

enterprise/experiments/experiment_manager.py

@@ -10,11 +10,14 @@ from experiments.experiment_versions import (
 from openhands.core.config.openhands_config import OpenHandsConfig
 from openhands.core.logger import openhands_logger as logger
 from openhands.experiments.experiment_manager import ExperimentManager
+from openhands.server.session.conversation_init_data import ConversationInitData
 
 
 class SaaSExperimentManager(ExperimentManager):
     @staticmethod
-    def run_conversation_variant_test(user_id, conversation_id, conversation_settings):
+    def run_conversation_variant_test(
+        user_id, conversation_id, conversation_settings
+    ) -> ConversationInitData:
         """
         Run conversation variant test and potentially modify the conversation settings
         based on the PostHog feature flags.
@@ -53,7 +56,7 @@ class SaaSExperimentManager(ExperimentManager):
     @staticmethod
     def run_config_variant_test(
         user_id: str, conversation_id: str, config: OpenHandsConfig
-    ):
+    ) -> OpenHandsConfig:
         """
         Run agent config variant test and potentially modify the OpenHands config
         based on the current experiment type and PostHog feature flags.

enterprise/experiments/experiment_versions/_003_llm_claude4_vs_gpt5_experiment.py

@@ -14,9 +14,10 @@ from server.constants import (
 from storage.experiment_assignment_store import ExperimentAssignmentStore
 
 from openhands.core.logger import openhands_logger as logger
+from openhands.server.session.conversation_init_data import ConversationInitData
 
 
-def _get_model_variant(user_id, conversation_id) -> str | None:
+def _get_model_variant(user_id: str | None, conversation_id: str) -> str | None:
     if not EXPERIMENT_CLAUDE4_VS_GPT5:
         logger.info(
             'experiment_manager:ab_testing:skipped',
@@ -104,7 +105,11 @@ def _get_model_variant(user_id, conversation_id) -> str | None:
     return enabled_variant
 
 
-def handle_claude4_vs_gpt5_experiment(user_id, conversation_id, conversation_settings):
+def handle_claude4_vs_gpt5_experiment(
+    user_id: str | None,
+    conversation_id: str,
+    conversation_settings: ConversationInitData,
+) -> ConversationInitData:
     """
     Handle the LiteLLM model experiment.
 
@@ -120,7 +125,7 @@ def handle_claude4_vs_gpt5_experiment(user_id, conversation_id, conversation_set
     enabled_variant = _get_model_variant(user_id, conversation_id)
 
     if not enabled_variant:
-        return None
+        return conversation_settings
 
     # Set the model based on the feature flag variant
     if enabled_variant == 'gpt5':

enterprise/experiments/experiment_versions/_004_condenser_max_step_experiment.py

@@ -11,6 +11,7 @@ from server.constants import IS_FEATURE_ENV
 from storage.experiment_assignment_store import ExperimentAssignmentStore
 
 from openhands.core.logger import openhands_logger as logger
+from openhands.server.session.conversation_init_data import ConversationInitData
 
 
 def _get_condenser_max_step_variant(user_id, conversation_id):
@@ -114,8 +115,10 @@ def _get_condenser_max_step_variant(user_id, conversation_id):
 
 
 def handle_condenser_max_step_experiment(
-    user_id: str, conversation_id: str, conversation_settings
-):
+    user_id: str | None,
+    conversation_id: str,
+    conversation_settings: ConversationInitData,
+) -> ConversationInitData:
     """
     Handle the condenser max step experiment for conversation settings.