agentgraph/shared/models/reference_based/entity.py

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

class Entity(BaseModel):
    id: str = Field(..., description="Unique identifier for the entity")
    type: Literal["Agent", "Task", "Tool", "Input", "Output", "Human"] = Field(
        ..., 
        description="Type of entity defined by prompt type: Agent (system prompt), Task (instruction prompt), Tool (description prompt), Input (input format prompt), Output (output format prompt), Human (optional prompt). The raw_prompt field is the primary distinguishing factor for entity uniqueness and classification."
    )
    name: str = Field(
        ..., 
        description="Name of the entity derived from the prompt content. Names should reflect the specific prompt or specification that defines this entity. For composite entities, use descriptive names that capture the prompt's scope (e.g., 'SQL Query Generation System Prompt', 'Data Analysis Instruction Set')."
    )
    importance: Literal["HIGH", "MEDIUM", "LOW"] = Field(
        ..., 
        description="Importance level of this entity in the system. HIGH: Core agents, critical tasks, essential tools that are central to system function. MEDIUM: Supporting agents, standard tasks, commonly used tools. LOW: Auxiliary entities, simple tasks, rarely used components."
    )
    raw_prompt: str = Field(
        default="", 
        description="PRIMARY DISTINGUISHING CONTENT: The actual prompt, specification, or instruction that defines this entity. This is the core content that makes each entity unique and should contain: For Agents (system prompts defining role/capabilities), For Tasks (instruction prompts defining objectives), For Tools (description prompts defining functionality), For Inputs (format specifications), For Outputs (format specifications), For Humans (interaction patterns). This field is more important than the name for entity distinction and relationship mapping."
    )
    raw_prompt_ref: List[ContentReference] = Field(
        default_factory=list,
        description="A list of references to the locations of the raw prompt content in the original trace. When provided, this allows mapping back to all exact positions in the trace where this prompt was found."
    )