src/chemgraph/schemas/agent_response.py

from typing import List, Optional

from pydantic import BaseModel, Field

from chemgraph.schemas.atomsdata import AtomsData


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^-1).
        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^-1).
        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.
    """

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

    intensity: list[str] = Field(
        ...,
        description="List of intensities in D/A^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^-1)
        Each entry is a string representation of the frequency value.
    intensity_spec_D2A2amu1 : list[str]
        List of range of intensities in (D/A)^2 amu^-1
        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/A)^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 DipoleResult(BaseModel):
    """
    Schema for storing a dipole moment vector from a simulation.

    Attributes
    ----------
    value : List[float]
        The dipole moment vector [dx, dy, dz].
    unit : str
        The unit of the dipole moment (e.g., 'e * Angstrom').
    """

    value: List[float] = Field(..., description="Dipole moment vector [dx, dy, dz].")
    unit: str = Field(..., description="Unit of the dipole moment, e.g. 'e * Angstrom'")


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

    Supports simultaneous multi-modal answers.  For example, the user
    can ask for a structure and a spectrum at the same time.

    The ``smiles`` field holds one or more SMILES strings returned by
    cheminformatics tools.  Each SMILES is a separate list element.
    """

    smiles: Optional[List[str]] = Field(
        default=None,
        description="SMILES strings for one or more molecules.",
    )
    scalar_answer: Optional[ScalarResult] = Field(
        default=None,
        description="Single numerical properties (e.g. enthalpy).",
    )
    scalar_answers: Optional[List[ScalarResult]] = Field(
        default=None,
        description=(
            "Multiple related scalar results from one validated workflow, "
            "for example reaction enthalpy and reaction Gibbs free energy."
        ),
    )
    dipole: Optional[DipoleResult] = Field(
        default=None,
        description="Dipole moment vector.",
    )
    vibrational_answer: Optional[VibrationalFrequency] = Field(
        default=None,
        description="Vibrational frequencies.",
    )
    ir_spectrum: Optional[IRSpectrum] = Field(
        default=None,
        description="Infrared spectra.",
    )
    atoms_data: Optional[AtomsData] = Field(
        default=None,
        description="Atomic geometries (XYZ coordinate, etc.) and optimized structures.",
    )