agentgraph/shared/models/reference_based/relation.py

from pydantic import BaseModel, Field
from typing import Dict, List, Optional, Any, Literal
import uuid
from .content_reference import ContentReference

class RelationType(BaseModel):
    name: str = Field(..., description="Name of the relation type")
    description: str = Field(..., description="Description of what this relation type means")
    source_type: str = Field(..., description="Type of entity that can be the source")
    target_type: str = Field(..., description="Type of entity that can be the target")

class Relation(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Unique identifier for the relation")
    source: str = Field(..., description="ID of the source entity")
    target: str = Field(..., description="ID of the target entity")
    type: Literal["CONSUMED_BY", "PERFORMS", "ASSIGNED_TO", "USES", "REQUIRED_BY", "SUBTASK_OF", "NEXT", "PRODUCES", "DELIVERS_TO", "INTERVENES"] = Field(
        ..., 
        description="Type of relation (only predefined types are allowed)"
    )
    importance: Literal["HIGH", "MEDIUM", "LOW"] = Field(
        ..., 
        description="Importance level of this relationship in the system. HIGH: Critical data flows, core agent-task assignments, essential tool usage. MEDIUM: Standard workflows, common interactions, regular data processing. LOW: Auxiliary connections, optional steps, rarely activated relationships."
    )
    interaction_prompt: str = Field(
        default="", 
        description="Actual runtime interaction message/log that shows this relationship occurring during execution. Contains the exact text from the trace where this interaction happened (e.g., 'Agent started task X', 'Calling tool Y with parameters Z', 'User provided feedback: ABC'). This is NOT the static prompt definition but the dynamic interaction evidence."
    )
    interaction_prompt_ref: List[ContentReference] = Field(
        default_factory=list,
        description="List of references to the locations of interaction prompt content in the original trace. Enables mapping back to all occurrences of the interaction prompt."
    )

# Predefined relation types
DEFAULT_RELATION_TYPES = [
    RelationType(
        name="CONSUMED_BY",
        description="Indicates that an input is consumed by an agent. Requires input format specification in interaction_prompt.",
        source_type="Input",
        target_type="Agent"
    ),
    RelationType(
        name="PERFORMS",
        description="Indicates that an agent actively executes a task. No prompt required.",
        source_type="Agent",
        target_type="Task"
    ),
    RelationType(
        name="ASSIGNED_TO",
        description="Indicates that a task is delegated or assigned to an agent. No prompt required.",
        source_type="Task",
        target_type="Agent"
    ),
    RelationType(
        name="USES",
        description="Indicates that an agent uses a tool. Requires system prompt components in interaction_prompt.",
        source_type="Agent",
        target_type="Tool"
    ),
    RelationType(
        name="REQUIRED_BY",
        description="Indicates that a tool is required by a task. Requires instruction prompt components in interaction_prompt.",
        source_type="Tool",
        target_type="Task"
    ),
    RelationType(
        name="SUBTASK_OF",
        description="Indicates that a task is a subtask of another task. No prompt required.",
        source_type="Task",
        target_type="Task"
    ),
    RelationType(
        name="NEXT",
        description="Indicates that a task follows another task in sequence. No prompt required.",
        source_type="Task",
        target_type="Task"
    ),
    RelationType(
        name="PRODUCES",
        description="Indicates that a task produces an output. No prompt required.",
        source_type="Task",
        target_type="Output"
    ),
    RelationType(
        name="DELIVERS_TO",
        description="Indicates that an output is delivered to a human. Requires output format specification in interaction_prompt.",
        source_type="Output",
        target_type="Human"
    ),
    RelationType(
        name="INTERVENES",
        description="Indicates that an agent or human intervenes in a task. Requires feedback format specification in interaction_prompt.",
        source_type="Agent,Human",
        target_type="Task"
    )
]