schemas.py

from typing import Dict, Any, Optional, List
from pydantic import BaseModel, Field

class InterventionOutput(BaseModel):
    """Schema for intervention/correction outputs from causal attribution."""

    corrected_reasoning: Optional[str] = Field(
        None,
        description="The corrected reasoning text for the step"
    )
    corrected_tool_name: Optional[str] = Field(
        None,
        description="The corrected tool name (if applicable)"
    )
    corrected_tool_args_json: Optional[str] = Field(
        None,
        description="The corrected tool arguments as a JSON string (e.g. {\"key\": \"value\"})"
    )
    corrected_text: Optional[str] = Field(
        None,
        description="The corrected text for the step that should have been generated by the LLM"
    )
    explanation: str = Field(
        ...,
        description="Brief explanation of what was corrected and why"
    )


class OutcomePrediction(BaseModel):
    """Schema for predicting if an intervention would fix the failure."""

    would_succeed: bool = Field(
        ...,
        description="Whether the intervention would likely fix the failure"
    )
    confidence: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="Confidence level between 0.0 and 1.0"
    )
    reasoning: str = Field(
        ...,
        description="Brief explanation of the prediction"
    )


class RepairOutput(BaseModel):
    """Schema for counterfactual repair outputs."""

    repaired_text: Optional[str] = Field(
        None,
        description="The minimally repaired text content"
    )
    repaired_tool_name: Optional[str] = Field(
        None,
        description="The repaired tool name (if step is a tool call)"
    )
    repaired_tool_args_json: Optional[str] = Field(
        None,
        description="The repaired tool arguments as a JSON string (if step is a tool call). Must be valid JSON object like {\"key\": \"value\"}"
    )
    changes_made: List[str] = Field(
        ...,
        description="List of specific changes made to achieve minimal repair"
    )
    minimality_justification: str = Field(
        ...,
        description="Explanation of why this is the minimal necessary change"
    )


class CritiqueOutput(BaseModel):
    """Schema for multi-agent critique outputs."""

    agreement: str = Field(
        ...,
        pattern="^(AGREE|DISAGREE|PARTIAL)$",
        description="Level of agreement with the causal claim"
    )
    confidence: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="Confidence in this critique between 0.0 and 1.0"
    )
    reasoning: str = Field(
        ...,
        description="Detailed explanation of the critique"
    )
    alternative_explanation: Optional[str] = Field(
        None,
        description="Alternative causal explanation if disagreeing"
    )
    evidence_strength: str = Field(
        ...,
        pattern="^(STRONG|MODERATE|WEAK)$",
        description="Strength of evidence for the causal claim"
    )


class SkillAttributionOutput(BaseModel):
    """Schema for mapping a causal step to an underlying skill."""

    skill_label: str = Field(
        ...,
        description="Short, general label for the underlying skill (2-4 words)"
    )
    skill_description: str = Field(
        ...,
        description="1-2 sentence description of the skill and typical failure pattern"
    )
    confidence: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="Confidence in the skill attribution between 0.0 and 1.0"
    )
    rationale: str = Field(
        ...,
        description="Brief rationale grounded in the step content"
    )


class SkillGroupOutput(BaseModel):
    """Schema for grouping skill labels into higher-level skill clusters."""

    group_name: str = Field(
        ...,
        description="Short, human-readable name for the skill group"
    )
    group_description: str = Field(
        ...,
        description="1-2 sentence description of the grouped failure pattern"
    )
    member_labels: List[str] = Field(
        ...,
        description="List of skill labels included in this group"
    )
    rationale: str = Field(
        ...,
        description="Brief explanation of why these labels belong together"
    )


class SkillGroupingOutput(BaseModel):
    """Schema for LLM-based skill label grouping."""

    groups: List[SkillGroupOutput] = Field(
        ...,
        description="List of semantic skill groups"
    )


class ToolArgsOutput(BaseModel):
    """Schema for parsing tool arguments from text."""

    parsed_args_json: str = Field(
        ...,
        description="Extracted tool arguments as a JSON string (e.g. {\"key\": \"value\"})"
    )
    confidence: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="Confidence in the parsing between 0.0 and 1.0"
    )


class GSM8KCalculationStep(BaseModel):
    description: str = Field(
        ...,
        description="Clear description of what is being calculated in this step"
    )
    operation: str = Field(
        ...,
        pattern="^(addition|subtraction|multiplication|division|other)$",
        description="Type of mathematical operation: addition, subtraction, multiplication, division, or other"
    )
    expression: str = Field(
        ...,
        description="The mathematical expression to evaluate (e.g., '16 - 3 - 4' or '9 * 2')"
    )


class GSM8KSolution(BaseModel):
    """Schema for complete GSM8K problem solution with structured steps."""

    reasoning: str = Field(
        ...,
        description="Brief overview of the approach to solve the problem"
    )
    steps: List[GSM8KCalculationStep] = Field(
        ...,
        description="List of calculation steps needed to solve the problem"
    )
    final_answer: str = Field(
        ...,
        description="The final numerical answer to the problem"
    )


class BrowseCompAgentStep(BaseModel):
    
    action_type: str = Field(
        ...,
        pattern="^(search|open_url|extract|answer)$",
        description="Type of action: 'search' to query the web, 'open_url' to fetch a page, 'extract' to note facts from current context, 'answer' to provide the final answer"
    )
    query: Optional[str] = Field(
        None,
        description="Search query string (required when action_type is 'search')"
    )
    url: Optional[str] = Field(
        None,
        description="URL to fetch (required when action_type is 'open_url')"
    )
    note: str = Field(
        ...,
        description="Brief rationale explaining why this action is being taken"
    )
    extracted_facts: List[str] = Field(
        default_factory=list,
        description="List of facts extracted/noted from current context (especially useful for 'extract' action)"
    )
    exact_answer: Optional[str] = Field(
        None,
        description="The final exact answer (required when action_type is 'answer')"
    )
    confidence: Optional[float] = Field(
        None,
        ge=0.0,
        le=100.0,
        description="Confidence percentage (0-100) in the answer (used when action_type is 'answer')"
    )
    explanation: Optional[str] = Field(
        None,
        description="Explanation for the final answer (used when action_type is 'answer')"
    )


class LLMSchemas:
    """Collection of schema utilities for structured LLM outputs."""

    SCHEMA_MAP: Dict[str, type[BaseModel]] = {
        "intervention": InterventionOutput,
        "outcome_prediction": OutcomePrediction,
        "repair": RepairOutput,
        "critique": CritiqueOutput,
        "skill_attribution": SkillAttributionOutput,
        "skill_grouping": SkillGroupingOutput,
        "tool_args": ToolArgsOutput,
        "gsm8k_solution": GSM8KSolution,
        "browsecomp_step": BrowseCompAgentStep,
    }

    @staticmethod
    def get_model(schema_name: str) -> type[BaseModel]:
        if schema_name not in LLMSchemas.SCHEMA_MAP:
            raise ValueError(
                f"Unknown schema: {schema_name}. "
                f"Available: {list(LLMSchemas.SCHEMA_MAP.keys())}"
            )
        return LLMSchemas.SCHEMA_MAP[schema_name]

    @staticmethod
    def get_response_format(schema_name: str) -> Dict[str, Any]:
        model = LLMSchemas.get_model(schema_name)

        # Generate JSON schema from Pydantic model with mode='serialization'
        # This produces cleaner schemas compatible with more providers
        json_schema = model.model_json_schema(mode='serialization')

        # Clean up schema for compatibility with Google Gemini and other providers
        # Remove $defs and flatten the schema
        if "$defs" in json_schema:
            defs = json_schema.pop("$defs")
            # Inline any references
            LLMSchemas._inline_refs(json_schema, defs)

        # Simplify schema for Google Gemini compatibility
        LLMSchemas._simplify_schema(json_schema)
        LLMSchemas._ensure_required_fields(json_schema)

        # Remove metadata that might interfere
        json_schema.pop("title", None)
        json_schema.pop("description", None)

        # Ensure additionalProperties is set
        if "additionalProperties" not in json_schema:
            json_schema["additionalProperties"] = False

        return {
            "type": "json_schema",
            "json_schema": {
                "name": schema_name,
                "strict": True,
                "schema": json_schema
            }
        }

    @staticmethod
    def _simplify_schema(schema: Dict[str, Any]) -> None:

        if isinstance(schema, dict):
            # Remove title from all properties
            schema.pop("title", None)

            # Handle anyOf for optional fields (common pattern: anyOf with null)
            if "anyOf" in schema:
                any_of = schema.pop("anyOf")
                non_null_options = [opt for opt in any_of if opt.get("type") != "null"]
                has_null = any(opt.get("type") == "null" for opt in any_of)
                if non_null_options:
                    option = non_null_options[0]
                    schema.update(option)
                    if has_null and "type" in option:
                        base_type = option["type"]
                        if isinstance(base_type, list):
                            type_list = list(dict.fromkeys(base_type + ["null"]))
                        else:
                            type_list = [base_type, "null"]
                        schema["type"] = type_list

            # Handle oneOf similarly
            if "oneOf" in schema:
                one_of = schema.pop("oneOf")
                non_null_options = [opt for opt in one_of if opt.get("type") != "null"]
                has_null = any(opt.get("type") == "null" for opt in one_of)
                if non_null_options:
                    option = non_null_options[0]
                    schema.update(option)
                    if has_null and "type" in option:
                        base_type = option["type"]
                        if isinstance(base_type, list):
                            type_list = list(dict.fromkeys(base_type + ["null"]))
                        else:
                            type_list = [base_type, "null"]
                        schema["type"] = type_list

            # Check if type is or contains "object"
            schema_type = schema.get("type")
            is_object_type = (
                schema_type == "object" or 
                (isinstance(schema_type, list) and "object" in schema_type)
            )
            
            if is_object_type:
                # OpenAI strict mode requires additionalProperties: false
                if "additionalProperties" not in schema:
                    schema["additionalProperties"] = False

            # Recursively process nested objects
            for _, value in list(schema.items()):
                if isinstance(value, dict):
                    LLMSchemas._simplify_schema(value)
                elif isinstance(value, list):
                    for item in value:
                        if isinstance(item, dict):
                            LLMSchemas._simplify_schema(item)

    @staticmethod
    def _ensure_required_fields(schema: Dict[str, Any]) -> None:

        if not isinstance(schema, dict):
            return

        if schema.get("type") == "object":
            properties = schema.get("properties")
            if isinstance(properties, dict):
                schema["required"] = list(properties.keys())

        for value in list(schema.values()):
            if isinstance(value, dict):
                LLMSchemas._ensure_required_fields(value)
            elif isinstance(value, list):
                for item in value:
                    if isinstance(item, dict):
                        LLMSchemas._ensure_required_fields(item)

    @staticmethod
    def _inline_refs(schema: Dict[str, Any], defs: Dict[str, Any]) -> None:
        """
        Recursively inline $ref references in a schema.
        Modifies schema in-place.

        Args:
            schema: The schema object to process
            defs: The definitions to inline
        """
        if isinstance(schema, dict):
            if "$ref" in schema:
                # Extract the reference name
                ref = schema["$ref"]
                if ref.startswith("#/$defs/"):
                    def_name = ref.replace("#/$defs/", "")
                    if def_name in defs:
                        # Replace the reference with the actual definition
                        definition = defs[def_name].copy()
                        schema.clear()
                        schema.update(definition)
                        # Recursively inline any nested refs
                        LLMSchemas._inline_refs(schema, defs)
            else:
                # Recursively process nested objects
                for key, value in list(schema.items()):
                    if isinstance(value, dict):
                        LLMSchemas._inline_refs(value, defs)
                    elif isinstance(value, list):
                        for item in value:
                            if isinstance(item, dict):
                                LLMSchemas._inline_refs(item, defs)

    @staticmethod
    def parse_response(schema_name: str, response_data: Dict[str, Any]) -> BaseModel:
        model = LLMSchemas.get_model(schema_name)
        return model(**response_data)