FusionAGI/fusionagi/maa/layers/intent_engine.py

"""Layer 1 — Intent Formalization Engine.

Responsible for:
1. Intent decomposition - breaking natural language into structured requirements
2. Requirement typing - classifying requirements (dimensional, load, environmental, process)
3. Load case enumeration - identifying operational scenarios
"""

import re
import uuid
from typing import Any

from fusionagi._logger import logger
from fusionagi.maa.schemas.intent import (
    EngineeringIntentGraph,
    IntentNode,
    LoadCase,
    RequirementType,
)


class IntentIncompleteError(Exception):
    """Raised when intent formalization cannot be completed due to missing information."""

    def __init__(self, message: str, missing_fields: list[str] | None = None):
        self.missing_fields = missing_fields or []
        super().__init__(message)


class IntentEngine:
    """
    Intent decomposition, requirement typing, and load case enumeration.

    Features:
    - Pattern-based requirement extraction from natural language
    - Automatic requirement type classification
    - Load case identification
    - Environmental bounds extraction
    - LLM-assisted formalization (optional)
    """

    # Patterns for dimensional requirements (measurements, tolerances)
    DIMENSIONAL_PATTERNS = [
        r"(\d+(?:\.\d+)?)\s*(mm|cm|m|in|inch|inches|ft|feet)\b",
        r"tolerance[s]?\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"±\s*(\d+(?:\.\d+)?)",
        r"(\d+(?:\.\d+)?)\s*×\s*(\d+(?:\.\d+)?)",
        r"diameter\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"radius\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"thickness\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"length\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"width\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"height\s*(?:of\s*)?(\d+(?:\.\d+)?)",
    ]

    # Patterns for load requirements (forces, pressures, stresses)
    LOAD_PATTERNS = [
        r"(\d+(?:\.\d+)?)\s*(N|kN|MN|lb|lbf|kg|kgf)\b",
        r"(\d+(?:\.\d+)?)\s*(MPa|GPa|Pa|psi|ksi)\b",
        r"load\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"force\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"stress\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"pressure\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"factor\s*of\s*safety\s*(?:of\s*)?(\d+(?:\.\d+)?)",
        r"yield\s*strength",
        r"tensile\s*strength",
        r"fatigue\s*(?:life|limit|strength)",
    ]

    # Patterns for environmental requirements
    ENVIRONMENTAL_PATTERNS = [
        r"(\d+(?:\.\d+)?)\s*(?:°|deg|degrees?)?\s*(C|F|K|Celsius|Fahrenheit|Kelvin)\b",
        r"temperature\s*(?:range|of)?\s*(\d+)",
        r"humidity\s*(?:of\s*)?(\d+)",
        r"corrosion\s*resist",
        r"UV\s*resist",
        r"water\s*(?:proof|resist)",
        r"chemical\s*resist",
        r"outdoor",
        r"marine",
        r"aerospace",
    ]

    # Patterns for process requirements
    PROCESS_PATTERNS = [
        r"CNC|machining|milling|turning|drilling",
        r"3D\s*print|additive|FDM|SLA|SLS|DMLS",
        r"cast|injection\s*mold|die\s*cast",
        r"weld|braze|solder",
        r"heat\s*treat|anneal|harden|temper",
        r"surface\s*finish|polish|anodize|plate",
        r"assembly|sub-assembly",
        r"material:\s*(\w+)",
        r"aluminum|steel|titanium|plastic|composite",
    ]

    # Load case indicator patterns
    LOAD_CASE_PATTERNS = [
        r"(?:during|under|in)\s+(\w+(?:\s+\w+)?)\s+(?:conditions?|operation|mode)",
        r"(\w+)\s+load\s+case",
        r"(?:static|dynamic|cyclic|impact|thermal)\s+load",
        r"(?:normal|extreme|emergency|failure)\s+(?:operation|conditions?|mode)",
        r"operating\s+(?:at|under|in)",
    ]

    def __init__(self, llm_adapter: Any | None = None):
        """
        Initialize the IntentEngine.

        Args:
            llm_adapter: Optional LLM adapter for enhanced natural language processing.
        """
        self._llm = llm_adapter

    def formalize(
        self,
        intent_id: str,
        natural_language: str | None = None,
        file_refs: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        use_llm: bool = True,
    ) -> EngineeringIntentGraph:
        """
        Formalize engineering intent from natural language and file references.

        Args:
            intent_id: Unique identifier for this intent.
            natural_language: Natural language description of requirements.
            file_refs: References to CAD files, specifications, etc.
            metadata: Additional metadata.
            use_llm: Whether to use LLM for enhanced processing (if available).

        Returns:
            EngineeringIntentGraph with extracted requirements.

        Raises:
            IntentIncompleteError: If required information is missing.
        """
        if not intent_id:
            raise IntentIncompleteError("intent_id required", ["intent_id"])

        if not natural_language and not file_refs:
            raise IntentIncompleteError(
                "At least one of natural_language or file_refs required",
                ["natural_language", "file_refs"],
            )

        nodes: list[IntentNode] = []
        load_cases: list[LoadCase] = []
        environmental_bounds: dict[str, Any] = {}

        # Process natural language if provided
        if natural_language:
            # Use LLM if available and requested
            if use_llm and self._llm:
                llm_result = self._formalize_with_llm(intent_id, natural_language)
                if llm_result:
                    return llm_result

            # Fall back to pattern-based extraction
            extracted = self._extract_requirements(intent_id, natural_language)
            nodes.extend(extracted["nodes"])
            load_cases.extend(extracted["load_cases"])
            environmental_bounds.update(extracted["environmental_bounds"])

        # Process file references
        if file_refs:
            for ref in file_refs:
                nodes.append(
                    IntentNode(
                        node_id=f"{intent_id}_file_{uuid.uuid4().hex[:8]}",
                        requirement_type=RequirementType.OTHER,
                        description=f"Reference: {ref}",
                        metadata={"file_ref": ref},
                    )
                )

        # If no nodes were extracted, create a general requirement
        if not nodes and natural_language:
            nodes.append(
                IntentNode(
                    node_id=f"{intent_id}_general_0",
                    requirement_type=RequirementType.OTHER,
                    description=natural_language[:500],
                )
            )

        logger.info(
            "Intent formalized",
            extra={
                "intent_id": intent_id,
                "num_nodes": len(nodes),
                "num_load_cases": len(load_cases),
            },
        )

        return EngineeringIntentGraph(
            intent_id=intent_id,
            nodes=nodes,
            load_cases=load_cases,
            environmental_bounds=environmental_bounds,
            metadata=metadata or {},
        )

    def _extract_requirements(
        self,
        intent_id: str,
        text: str,
    ) -> dict[str, Any]:
        """
        Extract requirements from text using pattern matching.

        Returns dict with nodes, load_cases, and environmental_bounds.
        """
        nodes: list[IntentNode] = []
        load_cases: list[LoadCase] = []
        environmental_bounds: dict[str, Any] = {}

        # Split into sentences for processing
        sentences = re.split(r'[.!?]+', text)

        node_counter = 0
        load_case_counter = 0

        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue

            # Check for dimensional requirements
            for pattern in self.DIMENSIONAL_PATTERNS:
                if re.search(pattern, sentence, re.IGNORECASE):
                    nodes.append(
                        IntentNode(
                            node_id=f"{intent_id}_dim_{node_counter}",
                            requirement_type=RequirementType.DIMENSIONAL,
                            description=sentence,
                            metadata={"pattern": "dimensional"},
                        )
                    )
                    node_counter += 1
                    break

            # Check for load requirements
            for pattern in self.LOAD_PATTERNS:
                if re.search(pattern, sentence, re.IGNORECASE):
                    nodes.append(
                        IntentNode(
                            node_id=f"{intent_id}_load_{node_counter}",
                            requirement_type=RequirementType.LOAD,
                            description=sentence,
                            metadata={"pattern": "load"},
                        )
                    )
                    node_counter += 1
                    break

            # Check for environmental requirements
            for pattern in self.ENVIRONMENTAL_PATTERNS:
                match = re.search(pattern, sentence, re.IGNORECASE)
                if match:
                    nodes.append(
                        IntentNode(
                            node_id=f"{intent_id}_env_{node_counter}",
                            requirement_type=RequirementType.ENVIRONMENTAL,
                            description=sentence,
                            metadata={"pattern": "environmental"},
                        )
                    )
                    node_counter += 1

                    # Extract specific bounds if possible
                    if "temperature" in sentence.lower():
                        temp_match = re.search(r"(-?\d+(?:\.\d+)?)", sentence)
                        if temp_match:
                            environmental_bounds["temperature"] = float(temp_match.group(1))
                    break

            # Check for process requirements
            for pattern in self.PROCESS_PATTERNS:
                if re.search(pattern, sentence, re.IGNORECASE):
                    nodes.append(
                        IntentNode(
                            node_id=f"{intent_id}_proc_{node_counter}",
                            requirement_type=RequirementType.PROCESS,
                            description=sentence,
                            metadata={"pattern": "process"},
                        )
                    )
                    node_counter += 1
                    break

            # Check for load cases
            for pattern in self.LOAD_CASE_PATTERNS:
                match = re.search(pattern, sentence, re.IGNORECASE)
                if match:
                    load_case_desc = match.group(0) if match.group(0) else sentence
                    load_cases.append(
                        LoadCase(
                            load_case_id=f"{intent_id}_lc_{load_case_counter}",
                            description=load_case_desc,
                            metadata={"source_sentence": sentence},
                        )
                    )
                    load_case_counter += 1
                    break

        return {
            "nodes": nodes,
            "load_cases": load_cases,
            "environmental_bounds": environmental_bounds,
        }

    def _formalize_with_llm(
        self,
        intent_id: str,
        natural_language: str,
    ) -> EngineeringIntentGraph | None:
        """
        Use LLM to extract structured requirements from natural language.

        Returns None if LLM processing fails (falls back to pattern matching).
        """
        if not self._llm:
            return None

        import json

        prompt = f"""Extract engineering requirements from the following text.
Return a JSON object with:
- "nodes": list of requirements, each with:
  - "requirement_type": one of "dimensional", "load", "environmental", "process", "other"
  - "description": the requirement text
- "load_cases": list of operational scenarios, each with:
  - "description": the scenario description
- "environmental_bounds": dict of environmental limits (e.g., {{"temperature_max": 85, "humidity_max": 95}})

Text: {natural_language[:2000]}

Return only valid JSON, no markdown."""

        try:
            messages = [
                {"role": "system", "content": "You are an engineering requirements extraction system."},
                {"role": "user", "content": prompt},
            ]

            # Try structured output if available
            if hasattr(self._llm, "complete_structured"):
                result = self._llm.complete_structured(messages)
                if result:
                    return self._parse_llm_result(intent_id, result)

            # Fall back to text completion
            raw = self._llm.complete(messages)
            if raw:
                # Clean up response
                if raw.startswith("```"):
                    raw = raw.split("```")[1]
                    if raw.startswith("json"):
                        raw = raw[4:]
                result = json.loads(raw)
                return self._parse_llm_result(intent_id, result)

        except Exception as e:
            logger.warning(f"LLM formalization failed: {e}")

        return None

    def _parse_llm_result(
        self,
        intent_id: str,
        result: dict[str, Any],
    ) -> EngineeringIntentGraph:
        """Parse LLM result into EngineeringIntentGraph."""
        nodes = []
        for i, node_data in enumerate(result.get("nodes", [])):
            req_type_str = node_data.get("requirement_type", "other")
            try:
                req_type = RequirementType(req_type_str)
            except ValueError:
                req_type = RequirementType.OTHER

            nodes.append(
                IntentNode(
                    node_id=f"{intent_id}_llm_{i}",
                    requirement_type=req_type,
                    description=node_data.get("description", ""),
                    metadata={"source": "llm"},
                )
            )

        load_cases = []
        for i, lc_data in enumerate(result.get("load_cases", [])):
            load_cases.append(
                LoadCase(
                    load_case_id=f"{intent_id}_lc_llm_{i}",
                    description=lc_data.get("description", ""),
                    metadata={"source": "llm"},
                )
            )

        environmental_bounds = result.get("environmental_bounds", {})

        return EngineeringIntentGraph(
            intent_id=intent_id,
            nodes=nodes,
            load_cases=load_cases,
            environmental_bounds=environmental_bounds,
            metadata={"formalization_source": "llm"},
        )

    def validate_completeness(self, graph: EngineeringIntentGraph) -> tuple[bool, list[str]]:
        """
        Validate that an intent graph has sufficient information.

        Returns:
            Tuple of (is_complete, list_of_missing_items)
        """
        missing = []

        if not graph.nodes:
            missing.append("No requirements extracted")

        # Check for at least one dimensional or load requirement for manufacturing
        has_dimensional = any(n.requirement_type == RequirementType.DIMENSIONAL for n in graph.nodes)
        any(n.requirement_type == RequirementType.LOAD for n in graph.nodes)

        if not has_dimensional:
            missing.append("No dimensional requirements specified")

        # Load cases are recommended but not required
        if not graph.load_cases:
            logger.info("No load cases specified for intent", extra={"intent_id": graph.intent_id})

        return len(missing) == 0, missing