llm_api.py

"""
REST API layer for LLM inference — farmer chat assistant.

Exposes HuggingFace Inference API through standardized endpoints.
Conversation state management is handled by the client/external service.

Why this exists:
    The FarmerAssistant class only wraps HTTP calls to HuggingFace.
    This module translates that into REST endpoints so external clients
    (web UIs, mobile apps, third-party services) can interact with the LLM
    without importing Python modules.

The endpoints here are stateless — conversation history, user sessions,
and context persistence are the caller's responsibility.
"""
from __future__ import annotations

import time
from typing import Any, Dict, List, Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from llm_chat import FarmerAssistant


# ──────────────────────────────────────────────────────────────────────────────
# Request/Response Models
# ──────────────────────────────────────────────────────────────────────────────


class ChatRequest(BaseModel):
    """Single LLM inference request.
    
    The caller provides the user's message. Conversation history and context
    are optional — the client maintains state.
    """
    content: str = Field(
        ...,
        min_length=1,
        max_length=2000,
        description="User's question or statement for the farmer assistant"
    )
    
    conversation_history: Optional[List[Dict[str, str]]] = Field(
        default=None,
        description=(
            "Previous messages in [{'role': 'user', 'content': '...'}, "
            "{'role': 'assistant', 'content': '...'}] format. "
            "Used to ground the response in prior context."
        )
    )
    
    farm_context: Optional[Dict[str, Any]] = Field(
        default=None,
        description=(
            "Optional farm metadata to include in the system context. "
            "Example: {'farm_name': 'Johnson Farm', 'crop': 'tomato', 'area_ha': 2.5}"
        )
    )
    
    max_tokens: int = Field(
        default=256,
        ge=10,
        le=1024,
        description="Maximum tokens in the response"
    )
    
    temperature: float = Field(
        default=0.7,
        ge=0.0,
        le=2.0,
        description="Creativity level (0=deterministic, 1=creative, 2=very creative)"
    )


class ChatResponse(BaseModel):
    """LLM inference response."""
    content: str = Field(..., description="Assistant's response")
    timestamp: str = Field(..., description="ISO8601 timestamp when response was generated")
    model: str = Field(..., description="Model ID used (e.g., 'mistral-7b-instruct')")
    tokens_used: Optional[int] = Field(
        default=None,
        description="Approximate tokens consumed by request (input + output)"
    )
    latency_ms: float = Field(..., description="Response time in milliseconds")
    metadata: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Additional metadata (e.g., warning flags, model-specific info)"
    )


class HealthResponse(BaseModel):
    """LLM backend health status."""
    status: str = Field(..., description="'ok' or 'unavailable'")
    model: str = Field(..., description="Active model ID")
    latency_ms: Optional[float] = Field(
        default=None,
        description="Milliseconds to reach the LLM backend"
    )
    message: Optional[str] = Field(
        default=None,
        description="Human-readable status message"
    )


class ContextValidationRequest(BaseModel):
    """Validate farm context before using in chat."""
    farm_context: Dict[str, Any] = Field(
        ...,
        description="Farm metadata to validate (e.g., design_summary, crop, area)"
    )


class ContextValidationResponse(BaseModel):
    """Result of context validation."""
    valid: bool = Field(..., description="Whether the context is valid")
    warnings: List[str] = Field(
        default_factory=list,
        description="Non-fatal issues (e.g., missing recommended fields)"
    )
    errors: List[str] = Field(
        default_factory=list,
        description="Fatal issues that should be resolved before use"
    )


# ──────────────────────────────────────────────────────────────────────────────
# Helper: Lazy-load FarmerAssistant singleton
# ──────────────────────────────────────────────────────────────────────────────

_ASSISTANT = None


def get_assistant() -> FarmerAssistant:
    """Lazy-load the FarmerAssistant on first use.
    
    Raises HTTPException(503) if the token is not configured.
    """
    global _ASSISTANT
    if _ASSISTANT is None:
        try:
            _ASSISTANT = FarmerAssistant()
        except ValueError as e:
            raise HTTPException(
                status_code=503,
                detail=f"LLM backend unavailable: {str(e)}",
            )
    return _ASSISTANT


# ──────────────────────────────────────────────────────────────────────────────
# Context Validation
# ──────────────────────────────────────────────────────────────────────────────


def _validate_farm_context(context: Dict[str, Any]) -> tuple[bool, List[str], List[str]]:
    """
    Validate farm context for use in chat.
    
    Returns: (is_valid, warnings, errors)
    """
    warnings = []
    errors = []
    
    # Optional but recommended fields
    recommended = ["farm_name", "crop", "area_ha"]
    present = set(context.keys())
    missing_recommended = [f for f in recommended if f not in present]
    if missing_recommended:
        warnings.append(f"Missing recommended fields: {', '.join(missing_recommended)}")
    
    # Validate field types if present
    if "area_ha" in context:
        try:
            float(context["area_ha"])
        except (TypeError, ValueError):
            errors.append("'area_ha' must be a number")
    
    if "crop" in context:
        valid_crops = ["tomato", "pepper", "lettuce", "cucumber", "orchard", "generic"]
        if str(context["crop"]).lower() not in valid_crops:
            warnings.append(
                f"Unknown crop '{context['crop']}'. "
                f"Expected one of: {', '.join(valid_crops)}"
            )
    
    # design_summary can be complex; just check it exists if referenced
    if "design_summary" in context and not isinstance(context["design_summary"], dict):
        warnings.append("'design_summary' should be a dict")
    
    is_valid = len(errors) == 0
    return is_valid, warnings, errors


# ──────────────────────────────────────────────────────────────────────────────
# Router
# ──────────────────────────────────────────────────────────────────────────────


def build_router() -> APIRouter:
    """Build the LLM inference API router."""
    router = APIRouter(prefix="/api/v1/llm", tags=["chat"])
    
    @router.get("/health")
    def health() -> HealthResponse:
        """
        Check if the LLM backend is accessible.
        
        Returns latency to help clients decide whether to retry or fallback.
        """
        try:
            assistant = get_assistant()
            start = time.perf_counter()
            # Simple test: check that we can reach the API
            # (Without actually sending tokens to the model)
            _ = assistant.api_token is not None
            elapsed_ms = (time.perf_counter() - start) * 1000
            
            return HealthResponse(
                status="ok",
                model=assistant.model_id,
                latency_ms=elapsed_ms,
                message="LLM backend is reachable",
            )
        except HTTPException as e:
            # Token not configured
            raise e
        except Exception as e:
            raise HTTPException(
                status_code=503,
                detail=f"LLM health check failed: {str(e)}",
            )
    
    @router.post("/chat")
    def chat(req: ChatRequest) -> ChatResponse:
        """
        Send a message to the farmer assistant and get a response.
        
        The caller is responsible for maintaining conversation history.
        Pass prior messages in `conversation_history` if you want context.
        
        Optional `farm_context` grounds the response in your farm data
        (e.g., crop, area, design summary).
        """
        try:
            assistant = get_assistant()
        except HTTPException:
            raise
        
        # Enrich the message with farm context if provided
        farm_context_str = ""
        if req.farm_context:
            # Validate context first
            is_valid, warnings, errors = _validate_farm_context(req.farm_context)
            if errors:
                raise HTTPException(
                    status_code=422,
                    detail={
                        "code": "invalid_farm_context",
                        "message": "Farm context validation failed",
                        "errors": errors,
                    },
                )
            
            # Build a brief context string to prepend to the user message.
            farm_name = req.farm_context.get("farm_name", "Your farm")
            crop = req.farm_context.get("crop", "unknown crop")
            area = req.farm_context.get("area_ha", "unknown area")
            farm_context_str = f"Farm context: {farm_name}, growing {crop}, {area} ha. "
        
        # Send to LLM
        start = time.perf_counter()
        try:
            response_text = assistant.chat(
                user_message=f"{farm_context_str}{req.content}",
                conversation_history=req.conversation_history,
                max_tokens=req.max_tokens,
                temperature=req.temperature,
            )
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"LLM inference failed: {str(e)}",
            )
        
        elapsed_ms = (time.perf_counter() - start) * 1000
        
        # Rough token estimation (1 token ≈ 4 chars)
        estimated_tokens = (
            len(req.content) + len(response_text) + len(farm_context_str)
        ) // 4
        
        from datetime import datetime, timezone
        timestamp = datetime.now(timezone.utc).isoformat()
        
        return ChatResponse(
            content=response_text,
            timestamp=timestamp,
            model=assistant.model_id,
            tokens_used=estimated_tokens,
            latency_ms=elapsed_ms,
            metadata={
                "farm_context_provided": req.farm_context is not None,
                "conversation_history_length": len(req.conversation_history or []),
            },
        )
    
    @router.post("/validate-context")
    def validate_context(req: ContextValidationRequest) -> ContextValidationResponse:
        """
        Validate farm context before using it in chat requests.
        
        Helps clients catch issues early without spending LLM tokens.
        """
        is_valid, warnings, errors = _validate_farm_context(req.farm_context)
        
        return ContextValidationResponse(
            valid=is_valid,
            warnings=warnings,
            errors=errors,
        )
    
    @router.post("/chat/batch")
    def chat_batch(requests: List[ChatRequest]) -> List[ChatResponse]:
        """
        Send multiple messages in a batch (useful for bulk analysis).
        
        Note: Each request is independent — no conversation history carried
        between items. For multi-turn conversations, use POST /api/v1/llm/chat
        with explicit history in each call.
        
        Rate-limited: max 10 requests per batch.
        """
        if len(requests) > 10:
            raise HTTPException(
                status_code=422,
                detail="Batch size limited to 10 requests",
            )
        
        try:
            assistant = get_assistant()
        except HTTPException:
            raise
        
        responses = []
        for req in requests:
            try:
                start = time.perf_counter()
                response_text = assistant.chat(
                    user_message=req.content,
                    conversation_history=req.conversation_history,
                    max_tokens=req.max_tokens,
                    temperature=req.temperature,
                )
                elapsed_ms = (time.perf_counter() - start) * 1000
                estimated_tokens = (len(req.content) + len(response_text)) // 4
                
                from datetime import datetime, timezone
                timestamp = datetime.now(timezone.utc).isoformat()
                
                responses.append(ChatResponse(
                    content=response_text,
                    timestamp=timestamp,
                    model=assistant.model_id,
                    tokens_used=estimated_tokens,
                    latency_ms=elapsed_ms,
                    metadata={
                        "farm_context_provided": req.farm_context is not None,
                    },
                ))
            except Exception as e:
                # Return error in same list position
                responses.append(ChatResponse(
                    content=f"Error: {str(e)}",
                    timestamp="",
                    model="",
                    latency_ms=0,
                    metadata={"error": True},
                ))
        
        return responses
    
    return router