mirror of
https://github.com/OpenHands/OpenHands.git
synced 2025-12-26 05:48:36 +08:00
53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
"""
|
|
Store for managing experiment assignments.
|
|
|
|
This store handles creating and updating experiment assignments for conversations.
|
|
"""
|
|
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
from storage.database import session_maker
|
|
from storage.experiment_assignment import ExperimentAssignment
|
|
|
|
from openhands.core.logger import openhands_logger as logger
|
|
|
|
|
|
class ExperimentAssignmentStore:
|
|
"""Store for managing experiment assignments."""
|
|
|
|
def update_experiment_variant(
|
|
self,
|
|
conversation_id: str,
|
|
experiment_name: str,
|
|
variant: str,
|
|
) -> None:
|
|
"""
|
|
Update the variant for a specific experiment.
|
|
|
|
Args:
|
|
conversation_id: The conversation ID
|
|
experiment_name: The name of the experiment
|
|
variant: The variant assigned
|
|
"""
|
|
with session_maker() as session:
|
|
# Use PostgreSQL's INSERT ... ON CONFLICT DO NOTHING to handle unique constraint
|
|
stmt = insert(ExperimentAssignment).values(
|
|
conversation_id=conversation_id,
|
|
experiment_name=experiment_name,
|
|
variant=variant,
|
|
)
|
|
stmt = stmt.on_conflict_do_nothing(
|
|
constraint='uq_experiment_assignments_conversation_experiment'
|
|
)
|
|
|
|
session.execute(stmt)
|
|
session.commit()
|
|
|
|
logger.info(
|
|
'experiment_assignment_store:upserted_variant',
|
|
extra={
|
|
'conversation_id': conversation_id,
|
|
'experiment_name': experiment_name,
|
|
'variant': variant,
|
|
},
|
|
)
|