app/core/state.py

from typing import Annotated, Any, Literal, TypedDict

from .messaging import AgentMessage, merge_messages
from .schemas import AgentResponse, TeamRole


def merge_dicts(a: dict, b: dict) -> dict:
    return {**a, **b}

def merge_lists(a: list, b: list) -> list:
    return a + b

def replace_reducer(a: Any, b: Any) -> Any:
    return b

class AgentState(TypedDict):
    """
    Agent state for LangGraph pipeline.
    """
    context: str
    retrieval_context: str
    messages: Annotated[list[AgentMessage], merge_messages]
    history: Annotated[list[AgentResponse], merge_lists]
    outputs: Annotated[dict[str, str], merge_dicts]
    current_role: Annotated[str, replace_reducer]
    feedback: Annotated[str, replace_reducer]
    retry_count: Annotated[int, replace_reducer]
    judge_results: Annotated[dict[str, dict[str, Any]], merge_dicts]
    prd_context: Annotated[dict[str, Any], replace_reducer]
    human_approval_required: Annotated[bool, replace_reducer]
    approved_by: Annotated[str | None, replace_reducer]
    # Knowledge Graph fields (added in Phase 02)
    knowledge_graph: Annotated[dict[str, Any], merge_dicts]  # role-scoped, keyed by role_name
    kg_schema_version: Annotated[int, replace_reducer]

    # Quality Gate fields (Phase 03)
    quality_reports: Annotated[list[dict[str, Any]], merge_lists]   # Append-only list of QualityReport dicts
    quality_metrics: Annotated[dict[str, dict[str, Any]], merge_dicts]  # role -> latest QualityReport metrics
    re_prompt_info: Annotated[dict[str, dict[str, Any]], merge_dicts]   # role -> {attempt, stagnant, strategy}
    agent_token_usage: Annotated[dict[str, dict[str, int]], merge_dicts]  # role -> {input_tokens, output_tokens}
    contradictions: Annotated[list[dict[str, Any]], merge_lists]  # Append-only list of contradiction dicts

    # Adversarial Quality Review fields
    critic_scores: Annotated[dict[str, dict[str, Any]], merge_dicts]  # role -> CriticScores dict
    skeptic_reviews: Annotated[dict[str, dict[str, Any]], merge_dicts]  # role -> SkepticAttackVectors dict

PRD_CONTEXT_FOR_ROLE: dict[TeamRole, list[str]] = {
    TeamRole.BUSINESS_ANALYST: ["user_stories", "features", "assumptions"],
    TeamRole.SOLUTION_ARCHITECT: ["technical_constraints", "assumptions"],
    TeamRole.DATA_ARCHITECT: ["data_requirements"],
    TeamRole.SECURITY_ANALYST: ["security_requirements", "compliance"],
    TeamRole.UX_DESIGNER: ["user_stories", "user_personas"],
    TeamRole.API_DESIGNER: ["technical_constraints"],
    TeamRole.QA_STRATEGIST: ["acceptance_criteria", "user_stories"],
    TeamRole.DEVOPS_ARCHITECT: ["technical_constraints", "deployment_requirements"],
}

ALL_DEPENDENCIES: list[TeamRole] = list(TeamRole)
AGENT_DEPENDENCIES: dict[TeamRole, list[TeamRole | Literal["*"]]] = {
    TeamRole.PRODUCT_OWNER: [],
    TeamRole.BUSINESS_ANALYST: [TeamRole.PRODUCT_OWNER],
    TeamRole.SOLUTION_ARCHITECT: [TeamRole.PRODUCT_OWNER],
    TeamRole.DATA_ARCHITECT: [TeamRole.PRODUCT_OWNER],
    TeamRole.SECURITY_ANALYST: [TeamRole.PRODUCT_OWNER],
    TeamRole.UX_DESIGNER: [TeamRole.PRODUCT_OWNER],
    TeamRole.API_DESIGNER: [TeamRole.SOLUTION_ARCHITECT, TeamRole.DATA_ARCHITECT],
    TeamRole.QA_STRATEGIST: [TeamRole.BUSINESS_ANALYST],
    TeamRole.DEVOPS_ARCHITECT: [TeamRole.SOLUTION_ARCHITECT, TeamRole.SECURITY_ANALYST],
}