src/test_api.py

#!/usr/bin/env python3
"""
Test FastAPI for crop yield prediction
"""

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Optional
import json
import subprocess
import tempfile
import os
from datetime import datetime

app = FastAPI(
    title="Crop Yield Prediction API",
    description="API for predicting crop yields using Random Forest model",
    version="1.0.0"
)

class CropPredictionRequest(BaseModel):
    year: int = Field(..., description="Crop year (e.g., 2024)", example=2024)
    state: str = Field(..., description="State name", example="Punjab")
    crop: str = Field(..., description="Crop name", example="Rice")
    season: str = Field(..., description="Season", example="Kharif")
    area: float = Field(..., description="Area in hectares", example=10.0)
    production: float = Field(..., description="Production in tons", example=25.0)
    rainfall: Optional[float] = Field(1000.0, description="Annual rainfall in mm", example=1200)
    fertilizer: Optional[float] = Field(50.0, description="Fertilizer usage in kg", example=75)
    pesticide: Optional[float] = Field(5.0, description="Pesticide usage in kg", example=8)

class CropPredictionResponse(BaseModel):
    model: str = Field(..., description="Model used for prediction", example="Random Forest")
    predicted_yield: str = Field(..., description="Predicted yield with units", example="2017.7 kg/hectare")
    total_expected_production: str = Field(..., description="Total expected production with units", example="20.18 tons")
    assessment: str = Field(..., description="Yield assessment", example="Good yield expected")

@app.get("/")
async def root():
    return {
        "message": "Crop Yield Prediction API. Use /docs for interactive API documentation.",
        "version": "1.0.0",
        "endpoints": {
            "predict": "/predict",
            "health": "/health",
            "docs": "/docs"
        }
    }

@app.get("/health")
async def health_check():
    # Check if model files exist
    model_files_exist = (
        os.path.exists("trained_models/preprocessor.pkl") and 
        os.path.exists("trained_models/random_forest_model.pkl")
    )
    
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "model_loaded": model_files_exist
    }

@app.post("/predict", response_model=CropPredictionResponse)
async def predict_yield(request: CropPredictionRequest):
    """
    Predict crop yield using the existing CLI tool
    """
    try:
        # Create input JSON
        input_data = {
            "year": request.year,
            "state": request.state,
            "crop": request.crop,
            "season": request.season,
            "area": request.area,
            "production": request.production,
            "rainfall": request.rainfall,
            "fertilizer": request.fertilizer,
            "pesticide": request.pesticide
        }
        
        # Use the existing JSON CLI tool
        process = subprocess.Popen(
            ["python3", "crop_yield_predictor.py", "json"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        stdout, stderr = process.communicate(json.dumps(input_data))
        
        if process.returncode != 0:
            raise HTTPException(status_code=500, detail=f"Prediction failed: {stderr}")
        
        try:
            result = json.loads(stdout)
            if "error" in result:
                raise HTTPException(status_code=400, detail=result["error"])
            
            return CropPredictionResponse(**result)
            
        except json.JSONDecodeError:
            raise HTTPException(status_code=500, detail=f"Invalid response from predictor: {stdout}")
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    import os
    port = int(os.environ.get("PORT", 8000))
    print(f"Starting server on port: {port}")
    uvicorn.run(app, host="0.0.0.0", port=port)