from pydantic import BaseModel, Field
from typing import Dict, List, Optional, Any
# Risk assessment functionality removed
# noqa
from .entity import Entity
from .relation import Relation
from .failure import Failure
from .optimization_recommendation import OptimizationRecommendation

class KnowledgeGraph(BaseModel):
    system_name: str = Field("", description="A concise, descriptive name for the agent system")
    system_summary: str = Field("", description="A short 2-3 sentence summary of the agent system's purpose and structure")
    entities: List[Entity] = Field(default_factory=list, description="List of entities in the knowledge graph")
    relations: List[Relation] = Field(default_factory=list, description="List of relations in the knowledge graph")
    failures: Optional[List['Failure']] = Field(default=None, description="Optional list of detected risk or failure events across the trace")
    optimizations: Optional[List[OptimizationRecommendation]] = Field(default=None, description="Optional list of recommendations for optimizing the agent system")