pepperplus-cb/src/control_backend/schemas/program.py

from enum import Enum
from typing import Literal

from pydantic import UUID4, BaseModel


class ProgramElement(BaseModel):
    """
    Represents a basic element of our behavior program.

    :ivar name: The researcher-assigned name of the element.
    :ivar id: Unique identifier.
    """

    name: str
    id: UUID4

    # To make program elements hashable
    model_config = {"frozen": True}


class LogicalOperator(Enum):
    AND = "AND"
    OR = "OR"


type Belief = KeywordBelief | SemanticBelief | InferredBelief
type BasicBelief = KeywordBelief | SemanticBelief


class KeywordBelief(ProgramElement):
    """
    Represents a belief that is set when the user spoken text contains a certain keyword.

    :ivar keyword: The keyword on which this belief gets set.
    """

    name: str = ""
    keyword: str


class SemanticBelief(ProgramElement):
    """
    Represents a belief that is set by semantic LLM validation.

    :ivar description: Description of how to form the belief, used by the LLM.
    """

    description: str


class InferredBelief(ProgramElement):
    """
    Represents a belief that gets formed by combining two beliefs with a logical AND or OR.

    These beliefs can also be :class:`InferredBelief`, leading to arbitrarily deep nesting.

    :ivar operator: The logical operator to apply.
    :ivar left: The left part of the logical expression.
    :ivar right: The right part of the logical expression.
    """

    name: str = ""
    operator: LogicalOperator
    left: Belief
    right: Belief


class Norm(ProgramElement):
    name: str = ""
    norm: str
    critical: bool = False


class BasicNorm(Norm):
    """
    Represents a behavioral norm.

    :ivar norm: The actual norm text describing the behavior.
    :ivar critical: When true, this norm should absolutely not be violated (checked separately).
    """

    pass


class ConditionalNorm(Norm):
    """
    Represents a norm that is only active when a condition is met (i.e., a certain belief holds).

    :ivar condition: When to activate this norm.
    """

    condition: Belief


type PlanElement = Goal | Action


class Plan(ProgramElement):
    """
    Represents a list of steps to execute. Each of these steps can be a goal (with its own plan)
    or a simple action.

    :ivar steps: The actions or subgoals to execute, in order.
    """

    name: str = ""
    steps: list[PlanElement]


class BaseGoal(ProgramElement):
    """
    Represents an objective to be achieved. This base version does not include a plan to achieve
    this goal, and is used in semantic belief extraction.

    :ivar description: A description of the goal, used to determine if it has been achieved.
    :ivar can_fail: Whether we can fail to achieve the goal after executing the plan.
    """

    description: str = ""
    can_fail: bool = True


class Goal(BaseGoal):
    """
    Represents an objective to be achieved. To reach the goal, we should execute the corresponding
    plan. It inherits from the BaseGoal a variable `can_fail`, which if true will cause the
    completion to be determined based on the conversation.

    Instances of this goal are not hashable because a plan is not hashable.

    :ivar plan: The plan to execute.
    """

    plan: Plan


type Action = SpeechAction | GestureAction | LLMAction


class SpeechAction(ProgramElement):
    """
    Represents the action of the robot speaking a literal text.

    :ivar text: The text to speak.
    """

    name: str = ""
    text: str


class Gesture(BaseModel):
    """
    Represents a gesture to be performed. Can be either a single gesture,
    or a random gesture from a category (tag).

    :ivar type: The type of the gesture, "tag" or "single".
    :ivar name: The name of the single gesture or tag.
    """

    type: Literal["tag", "single"]
    name: str


class GestureAction(ProgramElement):
    """
    Represents the action of the robot performing a gesture.

    :ivar gesture: The gesture to perform.
    """

    name: str = ""
    gesture: Gesture


class LLMAction(ProgramElement):
    """
    Represents the action of letting an LLM generate a reply based on its chat history
    and an additional goal added in the prompt.

    :ivar goal: The extra (temporary) goal to add to the LLM.
    """

    name: str = ""
    goal: str


class Trigger(ProgramElement):
    """
    Represents a belief-based trigger. When a belief is set, the corresponding plan is executed.

    :ivar condition: When to activate the trigger.
    :ivar plan: The plan to execute.
    """

    condition: Belief
    plan: Plan


class Phase(ProgramElement):
    """
    A distinct phase within a program, containing norms, goals, and triggers.

    :ivar norms: List of norms active in this phase.
    :ivar goals: List of goals to pursue in this phase.
    :ivar triggers: List of triggers that define transitions out of this phase.
    """

    name: str = ""
    norms: list[BasicNorm | ConditionalNorm]
    goals: list[Goal]
    triggers: list[Trigger]


class Program(BaseModel):
    """
    Represents a complete interaction program, consisting of a sequence or set of phases.

    :ivar phases: The list of phases that make up the program.
    """

    phases: list[Phase]