File size: 5,508 Bytes
a2498f7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
"""
Agent Output Validation
========================

JSON schemas for validating LLM agent outputs.
Ensures data integrity between pipeline stages.
"""

from typing import Any, Optional

try:
    from jsonschema import validate, ValidationError

    HAS_JSONSCHEMA = True
except ImportError:
    HAS_JSONSCHEMA = False

from core.logging import get_logger

logger = get_logger("validation")


# =============================================================================
# SCHEMAS
# =============================================================================

BRAND_IDENTIFICATION_SCHEMA = {
    "type": "object",
    "properties": {
        "brand_primary": {"type": ["string", "null"]},
        "brand_secondary": {"type": ["string", "null"]},
        "brand_accent": {"type": ["string", "null"]},
        "palette_strategy": {"type": "string"},
        "cohesion_score": {"type": ["number", "integer"]},
        "cohesion_notes": {"type": "string"},
        "semantic_names": {"type": "object"},
        "self_evaluation": {"type": "object"},
    },
    "required": ["brand_primary", "palette_strategy"],
}

BENCHMARK_ADVICE_SCHEMA = {
    "type": "object",
    "properties": {
        "recommended_benchmark": {"type": "string"},
        "recommended_benchmark_name": {"type": "string"},
        "reasoning": {"type": "string"},
        "alignment_changes": {"type": "array"},
        "pros_of_alignment": {"type": "array"},
        "cons_of_alignment": {"type": "array"},
        "alternative_benchmarks": {"type": "array"},
        "self_evaluation": {"type": "object"},
    },
    "required": ["recommended_benchmark", "reasoning"],
}

BEST_PRACTICES_SCHEMA = {
    "type": "object",
    "properties": {
        "overall_score": {"type": ["number", "integer"]},
        "checks": {"type": "array"},
        "priority_fixes": {"type": "array"},
        "passing_practices": {"type": "array"},
        "failing_practices": {"type": "array"},
        "self_evaluation": {"type": "object"},
    },
    "required": ["overall_score", "priority_fixes"],
}

HEAD_SYNTHESIS_SCHEMA = {
    "type": "object",
    "properties": {
        "executive_summary": {"type": "string"},
        "scores": {"type": "object"},
        "benchmark_fit": {"type": "object"},
        "brand_analysis": {"type": "object"},
        "top_3_actions": {"type": "array"},
        "color_recommendations": {"type": "array"},
        "type_scale_recommendation": {"type": "object"},
        "spacing_recommendation": {"type": "object"},
        "self_evaluation": {"type": "object"},
    },
    "required": ["executive_summary", "top_3_actions"],
}

# Map agent names to schemas
AGENT_SCHEMAS = {
    "aurora": BRAND_IDENTIFICATION_SCHEMA,
    "brand_identifier": BRAND_IDENTIFICATION_SCHEMA,
    "atlas": BENCHMARK_ADVICE_SCHEMA,
    "benchmark_advisor": BENCHMARK_ADVICE_SCHEMA,
    "sentinel": BEST_PRACTICES_SCHEMA,
    "best_practices": BEST_PRACTICES_SCHEMA,
    "nexus": HEAD_SYNTHESIS_SCHEMA,
    "head_synthesizer": HEAD_SYNTHESIS_SCHEMA,
}


# =============================================================================
# VALIDATION FUNCTIONS
# =============================================================================

def validate_agent_output(data: Any, agent_name: str) -> tuple[bool, Optional[str]]:
    """
    Validate an agent's output against its expected schema.

    Args:
        data: The output data (dict or dataclass with to_dict())
        agent_name: Name of the agent (e.g., 'aurora', 'nexus')

    Returns:
        (is_valid, error_message) tuple
    """
    agent_key = agent_name.lower().strip()
    schema = AGENT_SCHEMAS.get(agent_key)

    if not schema:
        logger.warning(f"No schema found for agent: {agent_name}")
        return True, None  # No schema = pass (don't block)

    # Convert dataclass to dict if needed
    if hasattr(data, "to_dict"):
        data_dict = data.to_dict()
    elif hasattr(data, "__dataclass_fields__"):
        from dataclasses import asdict
        data_dict = asdict(data)
    elif isinstance(data, dict):
        data_dict = data
    else:
        return False, f"Cannot validate: unexpected type {type(data)}"

    if not HAS_JSONSCHEMA:
        # Fallback: manual required-field check
        return _manual_validate(data_dict, schema, agent_name)

    try:
        validate(instance=data_dict, schema=schema)
        logger.debug(f"Validation passed for {agent_name}")
        return True, None
    except ValidationError as e:
        error_msg = f"Validation failed for {agent_name}: {e.message}"
        logger.warning(error_msg)
        return False, error_msg


def _manual_validate(data: dict, schema: dict, agent_name: str) -> tuple[bool, Optional[str]]:
    """Fallback validation without jsonschema library."""
    required = schema.get("required", [])
    missing = [field for field in required if field not in data]

    if missing:
        error_msg = f"{agent_name} output missing required fields: {missing}"
        logger.warning(error_msg)
        return False, error_msg

    return True, None


def validate_all_agents(outputs: dict) -> dict[str, tuple[bool, Optional[str]]]:
    """
    Validate all agent outputs at once.

    Args:
        outputs: Dict mapping agent_name → output data

    Returns:
        Dict mapping agent_name → (is_valid, error_message)
    """
    results = {}
    for agent_name, data in outputs.items():
        results[agent_name] = validate_agent_output(data, agent_name)
    return results