PMFlow/backend/app/api/experiments.py

60 lines
1.7 KiB
Python

"""实验相关 API"""
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from ..core.database import get_db
from ..models import Experiment
from ..schemas import ExperimentCreate, ExperimentUpdate
router = APIRouter(prefix="/experiments", tags=["experiments"])
@router.get("/")
async def list_experiments(db: AsyncSession = Depends(get_db)):
result = await db.execute(select(Experiment))
return result.scalars().all()
@router.post("/")
async def create_experiment(
experiment: ExperimentCreate,
db: AsyncSession = Depends(get_db),
):
new_experiment = Experiment(
title=experiment.title,
hypothesis=experiment.hypothesis,
target_metric=experiment.target_metric,
linked_context_nodes=experiment.linked_context_nodes,
status="running",
)
db.add(new_experiment)
await db.commit()
await db.refresh(new_experiment)
return new_experiment
@router.patch("/{experiment_id}")
async def update_experiment_score(
experiment_id: int,
update: ExperimentUpdate,
db: AsyncSession = Depends(get_db),
):
result = await db.execute(select(Experiment).where(Experiment.id == experiment_id))
experiment = result.scalar_one_or_none()
if not experiment:
raise HTTPException(status_code=404, detail="Experiment not found")
experiment.current_score = update.current_score
if update.current_score >= 0.9:
experiment.status = "success"
elif update.current_score < 0.5:
experiment.status = "failed"
else:
experiment.status = "running"
await db.commit()
await db.refresh(experiment)
return experiment