agentgraph/shared/models/direct_based/models.py

from typing import List, Literal, Optional

from pydantic import BaseModel, Field

# -----------------------------
# Schema v3 core data models
# -----------------------------

class Entity(BaseModel):
    """Core entity in the agent knowledge graph (schema v3)."""

    id: str = Field(..., description="Unique identifier for the entity")
    type: Literal["Agent", "Task", "Tool", "Input", "Output", "Human"] = Field(
        ..., description="Entity category"
    )
    name: str = Field(..., description="Concise, reusable name for the entity")
    # Critical prompt fragment that defines the entity.  Empty string if not available.
    raw_prompt: Optional[str] = Field("", description="Prompt fragment that uniquely defines the entity")


class Relation(BaseModel):
    """Directed relationship between two entities (schema v3)."""

    id: str = Field(..., 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="Relation category")
    # Prompt fragment representing the concrete interaction (optional except for certain relation types)
    raw_prompt: Optional[str] = Field("", description="Prompt fragment illustrating the interaction")


class KnowledgeGraph(BaseModel):
    """Complete knowledge graph for an agent system (schema v3)."""

    entities: List[Entity]
    relations: List[Relation]
    system_name: str
    system_summary: str