diff --git a/openhands/experiments/experiment_manager.py b/openhands/experiments/experiment_manager.py new file mode 100644 index 0000000000..39c6af6284 --- /dev/null +++ b/openhands/experiments/experiment_manager.py @@ -0,0 +1,19 @@ +import os + +from openhands.server.session.conversation_init_data import ConversationInitData +from openhands.utils.import_utils import get_impl + + +class ExperimentManager: + @staticmethod + def run_conversation_variant_test( + user_id: str, conversation_id: str, conversation_settings: ConversationInitData + ) -> ConversationInitData: + return conversation_settings + + +experiment_manager_cls = os.environ.get( + 'OPENHANDS_EXPERIMENT_MANAGER_CLS', + 'openhands.experiments.experiment_manager.ExperimentManager', +) +ExperimentManagerImpl = get_impl(ExperimentManager, experiment_manager_cls) diff --git a/openhands/server/listen_socket.py b/openhands/server/listen_socket.py index 365c63d1d0..dd49f30bce 100644 --- a/openhands/server/listen_socket.py +++ b/openhands/server/listen_socket.py @@ -19,6 +19,7 @@ from openhands.events.observation.agent import ( AgentStateChangedObservation, ) from openhands.events.serialization import event_to_dict +from openhands.experiments.experiment_manager import ExperimentManagerImpl from openhands.integrations.provider import PROVIDER_TOKEN_TYPE, ProviderToken from openhands.integrations.service_types import ProviderType from openhands.server.session.conversation_init_data import ConversationInitData @@ -49,7 +50,7 @@ def create_provider_tokens_object( async def setup_init_convo_settings( - user_id: str | None, providers_set: list[ProviderType] + user_id: str | None, conversation_id: str, providers_set: list[ProviderType] ) -> ConversationInitData: settings_store = await SettingsStoreImpl.get_instance(config, user_id) settings = await settings_store.load() @@ -73,7 +74,11 @@ async def setup_init_convo_settings( if user_secrets: session_init_args['custom_secrets'] = user_secrets.custom_secrets - return ConversationInitData(**session_init_args) + convo_init_data = ConversationInitData(**session_init_args) + # We should recreate the same experiment conditions when restarting a conversation + return ExperimentManagerImpl.run_conversation_variant_test( + user_id, conversation_id, convo_init_data + ) @sio.event @@ -119,7 +124,9 @@ async def connect(connection_id: str, environ: dict) -> None: f'User {user_id} is allowed to connect to conversation {conversation_id}' ) - conversation_init_data = await setup_init_convo_settings(user_id, providers_set) + conversation_init_data = await setup_init_convo_settings( + user_id, conversation_id, providers_set + ) agent_loop_info = await conversation_manager.join_conversation( conversation_id, connection_id, diff --git a/openhands/server/services/conversation_service.py b/openhands/server/services/conversation_service.py index e0e9867d4c..ccb3c75e90 100644 --- a/openhands/server/services/conversation_service.py +++ b/openhands/server/services/conversation_service.py @@ -4,6 +4,7 @@ from typing import Any from openhands.core.logger import openhands_logger as logger from openhands.events.action.message import MessageAction +from openhands.experiments.experiment_manager import ExperimentManagerImpl from openhands.integrations.provider import ( CUSTOM_SECRETS_TYPE_WITH_JSON_SCHEMA, PROVIDER_TOKEN_TYPE, @@ -78,6 +79,8 @@ async def create_new_conversation( session_init_args['git_provider'] = git_provider session_init_args['conversation_instructions'] = conversation_instructions conversation_init_data = ConversationInitData(**session_init_args) + + logger.info('Loading conversation store') conversation_store = await ConversationStoreImpl.get_instance(config, user_id) logger.info('ServerConversation store loaded') @@ -93,6 +96,7 @@ async def create_new_conversation( extra={'user_id': user_id, 'session_id': conversation_id}, ) + conversation_init_data = ExperimentManagerImpl.run_conversation_variant_test(user_id, conversation_id, conversation_init_data) conversation_title = get_default_conversation_title(conversation_id) logger.info(f'Saving metadata for conversation {conversation_id}') @@ -105,6 +109,7 @@ async def create_new_conversation( selected_repository=selected_repository, selected_branch=selected_branch, git_provider=git_provider, + llm_model=settings.llm_model, ) ) diff --git a/openhands/storage/data_models/conversation_metadata.py b/openhands/storage/data_models/conversation_metadata.py index d089811afa..8adea2835a 100644 --- a/openhands/storage/data_models/conversation_metadata.py +++ b/openhands/storage/data_models/conversation_metadata.py @@ -25,6 +25,7 @@ class ConversationMetadata: trigger: ConversationTrigger | None = None pr_number: list[int] = field(default_factory=list) created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + llm_model: str | None = None # Cost and token metrics accumulated_cost: float = 0.0 prompt_tokens: int = 0