src/chemgraph/schemas/multi_agent_response.py

from typing import Any, Literal, Optional, Union

from pydantic import BaseModel, Field, model_validator
from chemgraph.schemas.atomsdata import AtomsData


class WorkerTask(BaseModel):
    """
    Represents a task assigned to an executor agent for performing tool-based computations.

    Attributes:
        task_index (int): The index or ID of the task, typically used to track execution order.
        prompt (str): A natural language prompt that describes the task or request for which
                       the executor is expected to generate tool calls.
        retry_count (int): How many times this task has been previously attempted.
                           Defaults to 0 for new tasks.  When the planner re-dispatches
                           a failed task, the router increments this value automatically.
    """

    task_index: int = Field(..., description="Task index")
    prompt: str = Field(..., description="Prompt to send to executor for tool calls")
    retry_count: int = Field(
        default=0,
        description="Number of previous attempts for this task (0 = first attempt)",
    )


class PlannerResponse(BaseModel):
    """
    Response model from the Planner agent.

    The planner acts as a router: it decides whether to dispatch tasks
    to executor subgraphs (``executor_subgraph``), ask the human user
    for clarification (``ask_human``), or finish (``FINISH``) when all
    work is done.

    Attributes:
        thought_process (str): The planner's reasoning for the current decision.
        next_step (str): The next node to activate — ``"executor_subgraph"``
            to fan-out tasks, ``"ask_human"`` to request human input, or
            ``"FINISH"`` to end the workflow.
        tasks (list[WorkerTask] | None): Tasks to assign when routing to executors.
        clarification (str | None): Question to ask the human when
            ``next_step`` is ``"ask_human"``.
    """

    thought_process: str = Field(
        description="Your reasoning for the current decision."
    )
    next_step: Literal["executor_subgraph", "ask_human", "FINISH"] = Field(
        description="The next node to activate in the workflow."
    )
    tasks: list[WorkerTask] = Field(
        default=None,
        description="List of tasks to assign to executor subgraphs.",
    )
    clarification: Optional[str] = Field(
        default=None,
        description=(
            "Question to present to the human user when next_step is 'ask_human'. "
            "Must be provided when next_step is 'ask_human'."
        ),
    )

    @model_validator(mode="before")
    @classmethod
    def normalize_planner_payload(cls, data: Any) -> Any:
        """Accept common planner variants and coerce into PlannerResponse shape.

        Parameters
        ----------
        data : Any
            Raw planner payload before Pydantic validation.

        Returns
        -------
        Any
            Normalized payload compatible with ``PlannerResponse``.
        """
        if isinstance(data, list):
            return {
                "thought_process": "Delegating parsed tasks to executors.",
                "next_step": "executor_subgraph",
                "tasks": data,
            }

        if isinstance(data, dict):
            normalized = dict(data)
            # Accept legacy "worker_tasks" key
            if "tasks" not in normalized and "worker_tasks" in normalized:
                normalized["tasks"] = normalized.pop("worker_tasks")
            if "tasks" in normalized and "next_step" not in normalized:
                normalized["next_step"] = "executor_subgraph"
            if "tasks" in normalized and "thought_process" not in normalized:
                normalized["thought_process"] = (
                    "Delegating parsed tasks to executors."
                )
            # Accept legacy "question" key as clarification
            if (
                normalized.get("next_step") == "ask_human"
                and "clarification" not in normalized
                and "question" in normalized
            ):
                normalized["clarification"] = normalized.pop("question")
            return normalized

        return data


class VibrationalFrequency(BaseModel):
    """
    Schema for storing vibrational frequency results from a simulation.

    Attributes
    ----------
    frequency_cm1 : list[str]
        List of vibrational frequencies in inverse centimeters (cm⁻¹).
        Each entry is a string representation of the frequency value.
    """

    frequency_cm1: list[str] = Field(
        ...,
        description="List of vibrational frequencies in cm-1.",
    )

class IRSpectrum(BaseModel):
    """
    Schema for storing vibrational frequency  and intensities from a simulation.

    Attributes
    ----------
    frequency_cm1 : list[str]
        List of vibrational frequencies in inverse centimeters (cm⁻¹).
        Each entry is a string representation of the frequency value.
    intensity : list[str]
        List of vibrational intensities.
        Each entry is a string representation of the intensity value.
    plot : Optional[str]
        Base64-encoded PNG image of the IR spectrum plot.
    """

    frequency_cm1: list[str] = Field(
        ...,
        description="List of vibrational frequencies in cm-1.",
    )

    intensity: list[str] = Field(
        ...,
        description="List of intensities in D/Å^2 amu^-1.",
    )

    plot: Optional[str] = None   # base64 PNG image


class InfraredSpectrum(BaseModel):
    """
    Schema for calculating infrared spectrum from a simulation.

    Attributes
    ----------
    frequency_spec_cm1 : list[str]
        List of range of frequencies in inverse centimeters (cm⁻¹)
        Each entry is a string representation of the frequency value.
    intensity_spec_D2A2amu1 : list[str]
        List of range of intensities in (D/Å)^2 amu⁻¹
        Each entry is a string representation of the intensity value.
    """
    frequency_spec_cm1: list[str] = Field(
        ...,
        description="Range of frequencies for plotting spectrum in cm-1.",
    )
    
    intensity_spec_D2A2amu1: list[str] = Field(
        ...,
        description="Values of intensities for plotting spectrum in (D/Å)^2 amu^-1.",
    )


class ScalarResult(BaseModel):
    """
    Schema for storing a scalar numerical result from a simulation or calculation.

    Attributes
    ----------
    value : float
        The numerical value of the scalar result (e.g., 1.23).
    property : str
        The name of the physical or chemical property represented (e.g., 'enthalpy', 'Gibbs free energy').
    unit : str
        The unit associated with the result (e.g., 'eV', 'kJ/mol').
    """

    value: float = Field(..., description="Scalar numerical result like enthalpy")
    property: str = Field(
        ...,
        description="Name of the property, e.g. 'enthalpy', 'Gibbs free energy'",
    )
    unit: str = Field(..., description="Unit of the result, e.g. 'eV'")


class ResponseFormatter(BaseModel):
    """Defined structured response to the user."""

    answer: Union[
        str,
        ScalarResult,
        VibrationalFrequency,
        IRSpectrum,
        AtomsData,
    ] = Field(
        description=(
            "Structured answer to the user's query. Use:\n"
            "1. `str` for general or explanatory responses or SMILES string.\n"
            "2. `VibrationalFrequency` for vibrational frequencies.\n"
            "3. `ScalarResult` for single numerical properties (e.g. enthalpy).\n"
            "4. `AtomsData` for atomic geometries (XYZ coordinate, etc.) and optimized structures."
            "5. `InfraredSpectrum` for calculating infrared spectra."
        )
    )