agentgraph/shared/models/reference_based/optimization_recommendation.py

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

# Predefined recommendation types for common optimization scenarios
RecommendationType = Literal[
    "PROMPT_REFINEMENT",    # Suggests improving an entity's prompt for clarity or efficiency
    "AGENT_MERGING",        # Recommends merging two or more agents with overlapping roles
    "TASK_CONSOLIDATION",   # Suggests combining multiple simple tasks into one
    "TOOL_ENHANCEMENT",     # Recommends improving a tool (e.g., adding parameters, better error handling)
    "WORKFLOW_SIMPLIFICATION", # Suggests changes to the relationship flow (e.g., parallelizing tasks)
]

class OptimizationRecommendation(BaseModel):
    """Represents a suggestion for improving the structure, efficiency, or clarity of the agent system."""

    id: str = Field(default_factory=lambda: f"opt_{uuid.uuid4()}", description="Unique identifier for the optimization recommendation")
    recommendation_type: RecommendationType = Field(..., description="The category of optimization being suggested")
    description: str = Field(..., description="A detailed, human-readable explanation of the recommendation, including the observed pattern and the justification for the change.")
    affected_ids: List[str] = Field(default_factory=list, description="A list of Entity or Relation IDs that are the focus of this recommendation")
    raw_text_ref: List[ContentReference] = Field(default_factory=list, description="List of references to the exact trace locations related to this recommendation")