agent/langgraph_agent.py

from typing import TypedDict, Dict, Any
import json
from urllib.error import URLError
from urllib.request import Request, urlopen

from core.config import settings
from core.logger import logger

try:
    from langgraph.graph import StateGraph, END
    from langchain_ollama import OllamaLLM
except Exception:  # pragma: no cover - optional dependency safety
    StateGraph = None
    END = None
    OllamaLLM = None


# -----------------------------
# STATE DEFINITION
# -----------------------------

class InspectionState(TypedDict):
    input_data: Dict[str, Any]
    summary: str
    decision: str
    recommendation: str


# -----------------------------
# LLM FACTORY (IMPORTANT)
# -----------------------------

def get_llm():
    """Create LLM instance (better than global init)."""
    if OllamaLLM is None:
        raise RuntimeError("LLM dependencies are not installed")
    return OllamaLLM(
        model=settings.OLLAMA_MODEL,
        base_url=settings.OLLAMA_BASE_URL,
        timeout=settings.OLLAMA_TIMEOUT_SECONDS,
    )


def is_ollama_available() -> bool:
    try:
        with urlopen(
            f"{settings.OLLAMA_BASE_URL}/api/tags",
            timeout=settings.OLLAMA_TIMEOUT_SECONDS,
        ) as response:
            return response.status == 200
    except (URLError, ValueError, TimeoutError):
        return False


def is_huggingface_available() -> bool:
    return bool(settings.HF_TOKEN and settings.HF_CHAT_MODEL and settings.HF_ROUTER_BASE_URL)


def is_openrouter_available() -> bool:
    return bool(settings.OPENROUTER_API_KEY and settings.OPENROUTER_MODEL and settings.OPENROUTER_BASE_URL)


def resolve_llm_provider() -> str:
    configured = settings.LLM_PROVIDER.strip().lower()
    if configured and configured != "auto":
        return configured

    if is_openrouter_available():
        return "openrouter"
    if is_huggingface_available():
        return "huggingface"
    if OllamaLLM is not None and is_ollama_available():
        return "ollama"
    return "none"


def request_huggingface_recommendation(prompt: str) -> str:
    if not is_huggingface_available():
        raise RuntimeError("Hugging Face Inference Providers is not configured")

    payload = json.dumps({
        "model": settings.HF_CHAT_MODEL,
        "messages": [
            {
                "role": "user",
                "content": prompt,
            }
        ],
        "max_tokens": 220,
        "response_format": {"type": "text"},
    }).encode("utf-8")

    request = Request(
        f"{settings.HF_ROUTER_BASE_URL.rstrip('/')}/chat/completions",
        data=payload,
        headers={
            "Authorization": f"Bearer {settings.HF_TOKEN}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    with urlopen(request, timeout=settings.OLLAMA_TIMEOUT_SECONDS) as response:
        body = json.loads(response.read().decode("utf-8"))

    choices = body.get("choices") or []
    if not choices:
        raise RuntimeError("No choices were returned by Hugging Face")

    message = choices[0].get("message") or {}
    content = message.get("content", "")

    if isinstance(content, str):
        return content

    if isinstance(content, list):
        text_parts: list[str] = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                text_parts.append(str(item.get("text", "")))
        if text_parts:
            return "\n".join(part for part in text_parts if part)

    raise RuntimeError("Unable to extract recommendation content from Hugging Face response")


def request_openrouter_recommendation(prompt: str) -> str:
    if not is_openrouter_available():
        raise RuntimeError("OpenRouter is not configured")

    payload = json.dumps({
        "model": settings.OPENROUTER_MODEL,
        "messages": [
            {
                "role": "user",
                "content": prompt,
            }
        ],
        "max_tokens": 220,
    }).encode("utf-8")

    request = Request(
        f"{settings.OPENROUTER_BASE_URL.rstrip('/')}/chat/completions",
        data=payload,
        headers={
            "Authorization": f"Bearer {settings.OPENROUTER_API_KEY}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    with urlopen(request, timeout=settings.OLLAMA_TIMEOUT_SECONDS) as response:
        body = json.loads(response.read().decode("utf-8"))

    choices = body.get("choices") or []
    if not choices:
        raise RuntimeError("No choices were returned by OpenRouter")

    message = choices[0].get("message") or {}
    content = message.get("content", "")
    if not content:
        raise RuntimeError("Unable to extract recommendation content from OpenRouter")
    return str(content)


# -----------------------------
# NODE 1: SUMMARY
# -----------------------------

def summarize_node(state: InspectionState) -> Dict[str, str]:

    data = state["input_data"]

    summary = (
        f"Detected defects:\n"
        f"- Minor: {data['summary'].get('Minor', 0)}\n"
        f"- Moderate: {data['summary'].get('Moderate', 0)}\n"
        f"- Critical: {data['summary'].get('Critical', 0)}\n"
    )

    return {"summary": summary}


# -----------------------------
# NODE 2: RULE-BASED DECISION
# -----------------------------

def decision_node(state: InspectionState) -> Dict[str, str]:

    data = state["input_data"]

    critical = data["summary"].get("Critical", 0)
    moderate = data["summary"].get("Moderate", 0)

    if critical > 0:
        decision = "FAIL"
    elif moderate > 0:
        decision = "REVIEW"
    else:
        decision = "PASS"

    return {"decision": decision}


# -----------------------------
# NODE 3: LLM RECOMMENDATION
# -----------------------------

def recommendation_node(state: InspectionState) -> Dict[str, str]:

    defects = state["input_data"].get("defects", [])

    prompt = f"""
You are a senior steel quality control engineer working in a manufacturing plant.

Inspection Summary:
{state['summary']}

Detected Defects (detailed):
{defects}

Final Decision: {state['decision']}

Give a STRICTLY INDUSTRIAL RESPONSE:

1. Defect-wise Analysis
2. Severity Impact
3. Action Plan (Accept / Rework / Downgrade / Reject)
4. Material Disposition Recommendation
5. Process Improvement Suggestions

Rules:
- Be specific to EACH defect type
- Avoid generic advice
- Be practical and engineering-focused
- Include whether the material can still be downgraded for a lower-grade use case
"""

    try:
        provider = resolve_llm_provider()

        if provider == "huggingface":
            response = request_huggingface_recommendation(prompt)
        elif provider == "openrouter":
            response = request_openrouter_recommendation(prompt)
        elif provider == "ollama":
            llm = get_llm()
            response = llm.invoke(prompt)
        else:
            raise RuntimeError("No hosted LLM provider is configured")

    except Exception as e:
        logger.error(f"LLM failed: {e}")
        response = "LLM failed to generate recommendation. Please review manually."

    return {"recommendation": response}


# -----------------------------
# BUILD GRAPH (ONCE)
# -----------------------------

def build_agent():
    if StateGraph is None or END is None:
        raise RuntimeError("LangGraph dependencies are not installed")

    builder = StateGraph(InspectionState)

    builder.add_node("summarize", summarize_node)
    builder.add_node("decision", decision_node)
    builder.add_node("recommendation", recommendation_node)

    builder.set_entry_point("summarize")

    builder.add_edge("summarize", "decision")
    builder.add_edge("decision", "recommendation")
    builder.add_edge("recommendation", END)

    return builder.compile()


AGENT = None


def _build_fallback_response(input_data: Dict[str, Any]) -> Dict[str, Any]:
    summary = (
        f"Detected defects:\n"
        f"- Minor: {input_data.get('summary', {}).get('Minor', 0)}\n"
        f"- Moderate: {input_data.get('summary', {}).get('Moderate', 0)}\n"
        f"- Critical: {input_data.get('summary', {}).get('Critical', 0)}\n"
    )

    critical = input_data.get("summary", {}).get("Critical", 0)
    moderate = input_data.get("summary", {}).get("Moderate", 0)
    minor = input_data.get("summary", {}).get("Minor", 0)

    if critical > 0:
        decision = "FAIL"
        recommendation = (
            "Critical surface damage detected. Stop automatic acceptance, isolate the coil or sheet, "
            "and route the material for immediate engineering review or rejection."
        )
    elif moderate > 0:
        decision = "REVIEW"
        recommendation = (
            "Moderate defects detected. Hold the batch for operator review, confirm defect spread, "
            "and schedule corrective process checks before release."
        )
    elif minor > 0:
        decision = "PASS"
        recommendation = (
            "Only minor defects were detected. Material can proceed with monitoring and a short-term "
            "process capability check to prevent escalation."
        )
    else:
        decision = "PASS"
        recommendation = (
            "No actionable defects were detected. Continue production and keep the inspection line active "
            "for routine monitoring."
        )

    return {
        "summary": summary,
        "decision": decision,
        "recommendation": recommendation,
        "agent_mode": "heuristic",
        "agent_provider": "Rule-Based Safety Engine",
        "agent_model": "fallback",
    }


def _normalize_llm_mode(llm_mode: str | None) -> str:
    normalized = (llm_mode or "auto").strip().lower()
    if normalized in {"off", "auto", "always"}:
        return normalized
    return "auto"


# -----------------------------
# RUN AGENT
# -----------------------------

def run_agent(input_data: Dict[str, Any], *, llm_mode: str = "auto") -> Dict[str, Any]:
    fallback = _build_fallback_response(input_data)
    normalized_mode = _normalize_llm_mode(llm_mode)

    if normalized_mode == "off":
        return fallback

    if normalized_mode == "auto":
        source = str(input_data.get("source", "")).strip().lower()
        metadata = input_data.get("metadata") or {}
        persist_requested = bool(metadata.get("persist", True))

        if source in {"camera", "live", "stream"} or not persist_requested:
            return fallback

    if not settings.ENABLE_LLM_REPORTS:
        return fallback

    provider = resolve_llm_provider()

    if StateGraph is None or END is None:
        logger.warning("LLM report generation is enabled, but LangGraph dependencies are unavailable.")
        return fallback

    if provider == "none":
        logger.warning("No hosted LLM provider is configured. Using fallback inspection recommendation.")
        return fallback
    if provider == "huggingface":
        if not is_huggingface_available():
            logger.warning("Hugging Face Inference Providers is unavailable. Using fallback inspection recommendation.")
            return fallback
    elif provider == "openrouter":
        if not is_openrouter_available():
            logger.warning("OpenRouter is unavailable. Using fallback inspection recommendation.")
            return fallback
    else:
        if OllamaLLM is None:
            logger.warning("Ollama dependencies are unavailable. Using fallback inspection recommendation.")
            return fallback
        if not is_ollama_available():
            logger.warning("Ollama server is unavailable. Using fallback inspection recommendation.")
            return fallback

    try:
        global AGENT

        if AGENT is None:
            AGENT = build_agent()

        result = AGENT.invoke({
            "input_data": input_data
        })

        return {
            **result,
            "agent_mode": "llm",
            "agent_provider": (
                "Hugging Face Inference Providers"
                if provider == "huggingface"
                else "OpenRouter"
                if provider == "openrouter"
                else "Ollama"
            ),
            "agent_model": (
                settings.HF_CHAT_MODEL
                if provider == "huggingface"
                else settings.OPENROUTER_MODEL
                if provider == "openrouter"
                else settings.OLLAMA_MODEL
            ),
        }

    except Exception as e:
        logger.error(f"Agent execution failed: {e}")
        return fallback