analysis/nl2sql.py

"""
Natural Language Query Parser for Data Explorer
Uses Claude API to parse engineer queries into structured sensor + time range selections.
Adapted from clearskies_agent_v1.1/llm_parser.py
"""
import json
from typing import Dict, Optional, List
from datetime import datetime
from core import config


class NLQueryParser:
    """Parses natural language queries into structured sensor + time range selections"""

    def __init__(self):
        self.api_key = config.ANTHROPIC_API_KEY
        self.api_available = bool(self.api_key and self.api_key != 'your_api_key_here')
        self.model = "claude-sonnet-4-20250514"
        self._client = None

    @property
    def client(self):
        if self._client is None and self.api_available:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self.api_key)
        return self._client

    def parse(
        self,
        user_query: str,
        available_sensors: List[str],
    ) -> Optional[Dict]:
        """
        Parse a natural language query into structured parameters.

        Args:
            user_query: Natural language query from the engineer
            available_sensors: List of available sensor tag names in the database

        Returns:
            Dict with keys:
                sensors: List[str] - matched sensor tag names
                start_time: datetime - query start time
                end_time: datetime - query end time
                explanation: str - human-readable interpretation
            Or None if parsing fails.
        """
        if not self.api_available:
            return None

        static_instructions = self._build_static_instructions(available_sensors)
        dynamic_query = f'User query: "{user_query}"\n\nReturn ONLY the JSON object, no other text.'

        try:
            message = self.client.messages.create(
                model=self.model,
                max_tokens=1000,
                system=[
                    {
                        "type": "text",
                        "text": static_instructions,
                        "cache_control": {"type": "ephemeral"},
                    }
                ],
                messages=[{"role": "user", "content": dynamic_query}],
            )
            response_text = message.content[0].text
            parsed = self._extract_json(response_text)

            if parsed:
                return self._process_result(parsed, available_sensors)
            return None

        except Exception as e:
            return {'error': str(e)}

    def _build_static_instructions(self, available_sensors: List[str]) -> str:
        """Build the cacheable static instruction block with sensor context.

        This is separated from the dynamic user query so the Anthropic API
        can cache this large, stable prefix across repeated calls within
        the same session (5-minute TTL).
        """
        # Build sensor reference with descriptions
        sensor_lines = []
        for tag_name, info in config.SENSOR_TAGS.items():
            desc = info.get('description', '')
            units = info.get('units', '')
            sensor_lines.append(f"  {tag_name}: {desc} ({units})")
        sensor_reference = "\n".join(sensor_lines)

        # Build sensor group reference
        group_lines = []
        for group_name, tags in config.SENSOR_GROUPS.items():
            group_lines.append(f"  {group_name}: {', '.join(tags)}")
        group_reference = "\n".join(group_lines)

        current_date = datetime.now()

        return f"""You are a query parser for a ClearSkies Hydrogen (CSH2) cryopump testing data system. Parse the engineer's natural language query into structured JSON for data retrieval.

Current date/time: {current_date.strftime('%Y-%m-%d %H:%M:%S')}
Data is available from March 2025 to September 2025.

AVAILABLE SENSORS (tag name: description (units)):
{sensor_reference}

SENSOR GROUPS (for category-based queries):
{group_reference}

COMMON ALIASES:
- "pressure" or "pressures" -> PT130 (Discharge), PT01T (Cryotank), PT110 (Inlet)
- "temperature" or "temperatures" -> TT110 (Pump Feed Line), TT130 (Discharge)
- "flow" -> FT140 (Flow Transmitter)
- "motor" or "speed" or "rpm" -> M130_Speed (Motor Speed)
- "power" -> AvgPower, PeakPower, PowerConsumption
- "discharge pressure" -> PT130
- "inlet pressure" or "supply pressure" -> PT110
- "tank pressure" or "cryotank pressure" -> PT01T
- "discharge temperature" -> TT130
- "feed temperature" or "pump temperature" -> TT110
- "valve" -> AOV140
- "density" -> N2Density, H2Density
- "level" -> LT200
- "all sensors" or "everything" -> PT130, PT01T, PT110, TT110, TT130, FT140, M130_Speed

Parse queries into a JSON object with these fields:
{{
  "sensors": ["PT130", "TT110"],
  "start_time": "2025-09-25T10:00:00",
  "end_time": "2025-09-25T14:00:00",
  "explanation": "Discharge pressure and pump feed temperature from Sept 25, 10 AM to 2 PM"
}}

RULES:
- Match sensor names to the available sensors list (case-insensitive)
- Use the aliases above to resolve common terms like "pressure", "temperature" etc.
- For relative dates like "yesterday", "last week", calculate based on current date
- For date-only queries without specific times, use 00:00:00 for start and 23:59:59 for end
- For "between X and Y" time references, parse both times
- Use ISO 8601 format for timestamps (YYYY-MM-DDTHH:MM:SS)
- If a sensor group is mentioned (e.g., "all pressures"), include all sensors from that group
- The explanation should be a concise human-readable summary of what was parsed
- If no sensors are explicitly mentioned, try to infer from context, or default to PT130, TT110, FT140
- If no time range is mentioned, default to the most recent day (today)"""

    def _build_prompt(self, user_query: str, available_sensors: List[str]) -> str:
        """Build the full Claude prompt (used by fallback path).

        The main parse() method uses _build_static_instructions() + dynamic query
        separately for prompt caching. This method is kept for backward compatibility.
        """
        static = self._build_static_instructions(available_sensors)
        return f"""{static}

User query: "{user_query}"

Return ONLY the JSON object, no other text."""

    def _extract_json(self, text: str) -> Optional[Dict]:
        """Extract JSON object from LLM response"""
        try:
            start = text.find('{')
            end = text.rfind('}') + 1

            if start != -1 and end > start:
                json_str = text[start:end]
                parsed = json.loads(json_str)

                if 'sensors' in parsed:
                    return parsed

            return None
        except json.JSONDecodeError:
            return None

    def _process_result(self, parsed: Dict, available_sensors: List[str]) -> Optional[Dict]:
        """Validate and process the parsed result"""
        try:
            # Convert time strings to datetime
            start_time = self._convert_to_datetime(parsed.get('start_time'))
            end_time = self._convert_to_datetime(parsed.get('end_time'))

            if not start_time or not end_time:
                return {'error': 'Could not parse time range from query.'}

            # Validate sensors exist in database
            raw_sensors = parsed.get('sensors', [])
            valid_sensors = []
            # Case-insensitive match against available sensors
            sensor_map = {s.upper(): s for s in available_sensors}
            for s in raw_sensors:
                matched = sensor_map.get(s.upper())
                if matched:
                    valid_sensors.append(matched)

            if not valid_sensors:
                return {'error': f'No matching sensors found. Requested: {", ".join(raw_sensors)}'}

            return {
                'sensors': valid_sensors,
                'start_time': start_time,
                'end_time': end_time,
                'explanation': parsed.get('explanation', ''),
            }

        except Exception as e:
            return {'error': f'Processing failed: {e}'}

    def _convert_to_datetime(self, iso_string: Optional[str]) -> Optional[datetime]:
        """Convert ISO string to datetime object"""
        if not iso_string:
            return None
        try:
            return datetime.fromisoformat(iso_string)
        except (ValueError, TypeError):
            return None