generation.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. """Schemas for plan generation and rule engine."""
  2. from typing import Optional, Any
  3. from pydantic import BaseModel, Field
  4. class ParameterInfo(BaseModel):
  5. """Information about a scannable parameter."""
  6. name: str
  7. display_name: str
  8. unit: str
  9. category: str
  10. default_start: float
  11. default_stop: float
  12. default_step: float
  13. min_allowed: float
  14. max_allowed: float
  15. description: str = ""
  16. topology_supported: list[str] = Field(default_factory=list)
  17. class ParameterRegistryResponse(BaseModel):
  18. """Response containing all scannable parameters."""
  19. total: int
  20. parameters: list[ParameterInfo]
  21. class RangeRecommendation(BaseModel):
  22. """Recommended scan range for a parameter."""
  23. start: float
  24. stop: float
  25. step: float
  26. notes: str = ""
  27. class RangeRecommendRequest(BaseModel):
  28. """Request for range recommendation."""
  29. parameter_name: str
  30. boundary_conditions: dict[str, Any] = Field(default_factory=dict)
  31. class PlanGenerateRequest(BaseModel):
  32. """Request to generate a plan from boundary conditions."""
  33. boundary_conditions: dict[str, Any] = Field(default_factory=dict)
  34. parameter_names: Optional[list[str]] = None
  35. model_path: str = ""
  36. class PlanGenerateResponse(BaseModel):
  37. """Response containing a generated plan draft."""
  38. name: str
  39. model_path: str
  40. topology: str
  41. variables: list[dict[str, Any]]
  42. cases: list[dict[str, Any]] = Field(default_factory=list)
  43. estimated_points: int
  44. estimated_time_min: float
  45. generation_summary: dict[str, Any] = Field(default_factory=dict)