src/language/reasoner.py

"""
Language Reasoning Module — Natural-language command interpretation via LLM.

Takes a user command and the list of detected objects, then asks an LLM
to produce a structured action interpretation:

    {"action": "move", "source": "bottle", "destination": "box"}

Supported providers (set via LLM_PROVIDER env var):
    - "openai"  → OpenAI API  (GPT-4o-mini by default)
    - "groq"    → Groq API    (Llama 3.3 70B by default, ~5-10x faster)
"""

import json
import logging
from typing import Any

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# System prompt sent to the LLM
# ---------------------------------------------------------------------------
_SYSTEM_PROMPT = """\
You are an AI assistant for a Vision Language Action robotics system.

Given:
  1. A list of objects detected in a scene (with labels and bounding boxes).
  2. A natural-language command from a human operator.

Your task:
  Interpret the command and return a JSON object with exactly these keys:
    - "action"      : The verb describing what to do (e.g. "move", "pick_up",
                       "place", "inspect", "rotate", "remove").
    - "source"      : The object to act on.
    - "destination"  : The target location or object (if applicable, else null).

CRITICAL RULES:
  • For "source" and "destination", you MUST use the EXACT label string from
    the detected objects list. The user may refer to objects using different
    words (e.g. "wrench" might actually be detected as "knife", "toolbox"
    might be "suitcase", "mug" might be "cup"). Your job is to figure out
    which detected object the user is referring to and use that detected label.
  • If the user says "wrench" but only "knife" is detected, return "knife".
  • If the user says "toolbox" but only "suitcase" is detected, return "suitcase".
  • If the user says "bin" but only "bowl" is detected, return "bowl".
  • NEVER return a label that is not in the detected objects list.
  • If truly no detected object can reasonably match, use the closest detected
    object label and note it.
  • Return ONLY the JSON object, no extra text.
  • If the command is ambiguous, make your best interpretation.
  • If no destination is implied, set "destination" to null.
"""


# Provider → base URL mapping (Groq uses the OpenAI-compatible SDK)
_PROVIDER_BASE_URLS: dict[str, str | None] = {
    "openai": None,                           # uses default OpenAI endpoint
    "groq":   "https://api.groq.com/openai/v1",
}


class CommandReasoner:
    """
    Interprets natural-language commands using an OpenAI-compatible LLM.

    Supports OpenAI and Groq as backends — both expose the same
    Chat Completions API so the same prompt and parsing logic work for both.

    Attributes:
        provider: Active provider name ("openai" or "groq").
        model:    Model identifier (e.g. "gpt-4o-mini" or "llama-3.3-70b-versatile").
        client:   OpenAI SDK client instance (or None in fallback mode).
    """

    def __init__(
        self,
        provider: str,
        api_key: str,
        model: str,
    ):
        """
        Initialise the reasoner for a specific provider.

        Args:
            provider: "openai" or "groq".
            api_key:  API key for the chosen provider.
            model:    Model identifier string.
        """
        self.provider = provider
        self.model = model
        self.client = None

        if api_key:
            self._init_client(provider, api_key)
        else:
            logger.warning(
                "No API key provided for provider='%s' — "
                "CommandReasoner will use fallback mode.",
                provider,
            )

    def _init_client(self, provider: str, api_key: str) -> None:
        """Create an OpenAI-SDK client pointed at the chosen provider endpoint."""
        try:
            from openai import OpenAI

            base_url = _PROVIDER_BASE_URLS.get(provider)  # None → OpenAI default
            kwargs: dict = {"api_key": api_key}
            if base_url:
                kwargs["base_url"] = base_url

            self.client = OpenAI(**kwargs)
            logger.info(
                "LLM client initialised (provider=%s, model=%s).",
                provider,
                self.model,
            )
        except Exception as exc:
            logger.error("Failed to initialise LLM client (%s): %s", provider, exc)
            self.client = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def interpret(
        self,
        command: str,
        detected_objects: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """
        Interpret a user command in the context of detected objects.

        Args:
            command: Free-form user instruction.
            detected_objects: List of detection dicts from the vision module.

        Returns:
            Dict with keys ``action``, ``source``, ``destination``.
        """
        if not command or not command.strip():
            logger.warning("Empty command received.")
            return self._error_result("No command provided.")

        if self.client is None:
            logger.info("Using fallback reasoning (no LLM client).")
            return self._fallback_interpret(command, detected_objects)

        return self._llm_interpret(command, detected_objects)

    # ------------------------------------------------------------------
    # LLM-based interpretation
    # ------------------------------------------------------------------

    def _llm_interpret(
        self,
        command: str,
        detected_objects: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """Call the LLM to interpret the command."""
        objects_description = self._format_objects(detected_objects)

        user_message = (
            f"Detected objects in the scene:\n{objects_description}\n\n"
            f"User command: \"{command}\"\n\n"
            f"Return the structured JSON interpretation."
        )

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": _SYSTEM_PROMPT},
                    {"role": "user", "content": user_message},
                ],
                temperature=0.1,
                max_tokens=256,
            )

            raw = response.choices[0].message.content.strip()
            logger.info("LLM raw response: %s", raw)

            return self._parse_response(raw)

        except Exception as exc:
            logger.error("LLM call failed: %s", exc)
            return self._fallback_interpret(command, detected_objects)

    # ------------------------------------------------------------------
    # Fallback (no LLM)
    # ------------------------------------------------------------------

    def _fallback_interpret(
        self,
        command: str,
        detected_objects: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """
        Smart keyword-based fallback when no LLM is available.

        Uses the ObjectMapper for synonym-aware matching between user
        terms and detected YOLO labels.
        """
        from src.utils.object_mapper import find_best_match

        logger.info("Fallback interpretation for: '%s'", command)

        # Try to determine action verb
        action = "move"
        for verb in ("pick", "place", "move", "inspect", "rotate", "remove", "find"):
            if verb in command.lower():
                action = verb
                break

        labels = [obj["label"] for obj in detected_objects]
        cmd_lower = command.lower()

        # Extract noun phrases from command by removing common verbs/prepositions
        stop_words = {
            "the", "a", "an", "and", "or", "to", "from", "in", "on", "at",
            "into", "onto", "of", "with", "for", "it", "up", "back", "next",
            "pick", "place", "move", "inspect", "rotate", "remove", "find",
            "put", "take", "grab", "get", "bring", "carry", "push", "pull",
            "near", "close", "closest", "nearest", "under", "over", "above",
            "below", "beside", "behind",
        }
        words = cmd_lower.replace(",", " ").replace(".", " ").split()

        # Build candidate noun phrases (single words + bigrams)
        candidates: list[str] = []
        for w in words:
            if w not in stop_words and len(w) > 1:
                candidates.append(w)
        for i in range(len(words) - 1):
            bigram = f"{words[i]} {words[i+1]}"
            if words[i] not in stop_words or words[i+1] not in stop_words:
                candidates.append(bigram)

        # Match candidates against detected labels using ObjectMapper
        source = None
        destination = None
        used_labels: set[str] = set()

        for candidate in candidates:
            match = find_best_match(candidate, labels)
            if match and match not in used_labels:
                if source is None:
                    source = match
                    used_labels.add(match)
                elif destination is None:
                    destination = match
                    used_labels.add(match)
                    break

        # Fallback: if no matches, use first two detected objects
        if source is None and labels:
            source = labels[0]
        if destination is None and len(labels) > 1:
            remaining = [l for l in labels if l != source]
            if remaining:
                destination = remaining[0]

        return {
            "action": action,
            "source": source or "unknown",
            "destination": destination,
        }

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _format_objects(detected_objects: list[dict[str, Any]]) -> str:
        """Format detected objects as a readable string for the LLM prompt."""
        if not detected_objects:
            return "No objects detected."

        lines: list[str] = []
        for i, obj in enumerate(detected_objects, 1):
            bbox = obj.get("bbox", [])
            conf = obj.get("confidence", 0)
            lines.append(
                f"  {i}. {obj['label']} (confidence: {conf:.2f}, bbox: {bbox})"
            )
        return "\n".join(lines)

    @staticmethod
    def _parse_response(raw: str) -> dict[str, Any]:
        """
        Parse the LLM response string into a dict.

        Handles markdown code fences and extra whitespace.
        """
        cleaned = raw.strip()
        # Strip optional markdown code fences
        if cleaned.startswith("```"):
            lines = cleaned.split("\n")
            lines = [ln for ln in lines if not ln.strip().startswith("```")]
            cleaned = "\n".join(lines).strip()

        try:
            result = json.loads(cleaned)
        except json.JSONDecodeError:
            logger.warning("Failed to parse LLM JSON: %s", cleaned)
            return {
                "action": "unknown",
                "source": "unknown",
                "destination": None,
                "error": "Failed to parse LLM response.",
                "raw_response": raw,
            }

        # Ensure required keys
        return {
            "action": result.get("action", "unknown"),
            "source": result.get("source", "unknown"),
            "destination": result.get("destination"),
        }

    @staticmethod
    def _error_result(message: str) -> dict[str, Any]:
        """Return a structured error result."""
        return {
            "action": "error",
            "source": "unknown",
            "destination": None,
            "error": message,
        }