pcb-afm-simulation-system/web/backend/app/routers/ai_plan.py

"""AI Plan Generation router (P3-M3).

Endpoints for generating simulation plans from natural language
using Kimi k3 model, with L0 pre-screening validation.
"""
from typing import Dict, Any, Optional, List
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from ..services.plan_generator import get_plan_generator

router = APIRouter(prefix="/api/ai-plan", tags=["AI Plan Generation"])


class GeneratePlanRequest(BaseModel):
    """Request to generate a plan from natural language."""
    requirement: str = Field(..., description="Natural language simulation requirement")
    project_context: Optional[Dict[str, Any]] = Field(default=None, description="Optional project context")
    existing_experience: Optional[List[Dict[str, Any]]] = Field(default=None, description="Optional similar experience cases")


class RefinePlanRequest(BaseModel):
    """Request to refine an existing plan."""
    original_plan: Dict[str, Any] = Field(..., description="The previously generated plan")
    feedback: str = Field(..., description="Natural language feedback for refinement")


@router.post("/generate")
def generate_plan(request: GeneratePlanRequest):
    """Generate a simulation plan from natural language.

    Uses Kimi k3 model to convert natural language requirements into
    a structured simulation plan, with L0 pre-screening validation.
    """
    try:
        generator = get_plan_generator()
        result = generator.generate(
            user_requirement=request.requirement,
            project_context=request.project_context,
            existing_experience=request.existing_experience,
        )
        return result
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Plan generation failed: {str(e)}")


@router.post("/refine")
def refine_plan(request: RefinePlanRequest):
    """Refine an existing plan based on user feedback."""
    try:
        generator = get_plan_generator()
        result = generator.refine_plan(
            original_plan=request.original_plan,
            user_feedback=request.feedback,
        )
        return result
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Plan refinement failed: {str(e)}")


@router.get("/templates")
def list_templates():
    """List available plan generation prompt templates."""
    return {
        "templates": [
            {
                "name": "generate",
                "path": "prompts/plan_generation/generate.txt",
                "description": "Natural language to structured simulation plan",
            }
        ],
        "model": "kimi-k3",
        "endpoint": "https://api.kimi.com/coding/v1",
    }