FlexibleTestPlatform/backend/app/api/routes/ai_assistant.py

"""
AI 助手 API 路由
"""

from typing import Optional, List
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from pydantic import BaseModel
from datetime import datetime

from app.database import get_db
from app.models.intent import Intent, AICheckRecord
from app.services.ai_service import get_ai_service

router = APIRouter(prefix="/ai", tags=["AI助手"])


# ============ Pydantic Schemas ============

class GenerateRequest(BaseModel):
    """AI生成请求"""
    prompt: str
    context: Optional[str] = None


class GenerateResponse(BaseModel):
    """AI生成响应"""
    content: str
    success: bool
    error: Optional[str] = None


class CheckRequest(BaseModel):
    """AI检查请求"""
    content: str
    intent_id: Optional[int] = None  # 如果提供，会保存检查记录
    requirements: Optional[List[str]] = None


class CheckResult(BaseModel):
    """AI检查结果"""
    passed: bool
    score: int
    issues: List[str]
    suggestions: List[str]
    raw_response: Optional[str] = None


class CheckResponse(BaseModel):
    """AI检查响应"""
    result: CheckResult
    record_id: Optional[int] = None  # 如果保存了记录
    success: bool
    error: Optional[str] = None


class IntentGenerateRequest(BaseModel):
    """意图编制生成请求"""
    test_type: str  # 测试类型：功能测试、性能测试、安全测试等
    test_target: str  # 测试目标描述
    additional_requirements: Optional[str] = None


class PlanGenerateRequest(BaseModel):
    """测试规划生成请求"""
    requirement_text: str  # 测试需求文本（前端文本框内容）
    source_name: Optional[str] = "用户输入"  # 来源名称


class ExtractStepsRequest(BaseModel):
    """提取测试步骤请求"""
    content: str  # 意图编制内容
    title: Optional[str] = "意图编制"


class ExtractedStepData(BaseModel):
    """提取的步骤数据"""
    id: int
    name: str
    purpose: str
    deviceParams: dict  # {设备类别: {参数名: 参数值}}


class ExtractStepsResponse(BaseModel):
    """提取测试步骤响应"""
    success: bool
    steps: List[ExtractedStepData]
    error: Optional[str] = None


# ============ API Endpoints ============

@router.post("/generate", response_model=GenerateResponse)
async def generate_content(request: GenerateRequest):
    """AI 生成内容"""
    try:
        ai_service = get_ai_service()
        content = await ai_service.generate(
            prompt=request.prompt,
            context=request.context
        )
        return GenerateResponse(content=content, success=True)
    except Exception as e:
        return GenerateResponse(
            content="",
            success=False,
            error=str(e)
        )


@router.post("/check", response_model=CheckResponse)
async def check_content(
    request: CheckRequest,
    db: AsyncSession = Depends(get_db)
):
    """AI 检查内容"""
    try:
        ai_service = get_ai_service()
        result = await ai_service.check(
            content=request.content,
            requirements=request.requirements
        )

        check_result = CheckResult(
            passed=result.get("passed", False),
            score=result.get("score", 0),
            issues=result.get("issues", []),
            suggestions=result.get("suggestions", []),
            raw_response=result.get("raw_response")
        )

        record_id = None

        # 如果提供了 intent_id，保存检查记录
        if request.intent_id:
            # 验证意图存在
            query = select(Intent).where(Intent.id == request.intent_id)
            intent_result = await db.execute(query)
            intent = intent_result.scalar_one_or_none()

            if intent:
                record = AICheckRecord(
                    intent_id=request.intent_id,
                    check_result=result,
                    suggestions="\n".join(result.get("suggestions", []))
                )
                db.add(record)
                await db.commit()
                await db.refresh(record)
                record_id = record.id

        return CheckResponse(
            result=check_result,
            record_id=record_id,
            success=True
        )
    except Exception as e:
        return CheckResponse(
            result=CheckResult(passed=False, score=0, issues=[str(e)], suggestions=[]),
            success=False,
            error=str(e)
        )


@router.post("/generate-intent", response_model=GenerateResponse)
async def generate_intent_content(request: IntentGenerateRequest):
    """AI 生成意图编制初始内容"""
    prompt = f"""请帮我生成一份测试意图编制文档的初始内容。

测试类型：{request.test_type}
测试目标：{request.test_target}
{"额外要求：" + request.additional_requirements if request.additional_requirements else ""}

请按照以下格式生成内容：

## 1. 测试目标
[详细描述测试的目标和预期达成的效果]

## 2. 测试范围
[明确测试的边界和覆盖范围]

## 3. 测试条件
### 3.1 前置条件
[列出测试开始前需要满足的条件]

### 3.2 测试环境
[描述测试所需的硬件、软件环境]

### 3.3 测试数据
[描述测试所需的数据准备]

## 4. 测试用例概述
[列出主要的测试场景和用例]

## 5. 预期结果
[描述测试成功的判定标准]

## 6. 风险与注意事项
[列出可能的风险和需要注意的事项]
"""

    context = """你是一位专业的软件测试工程师，擅长编写测试意图编制文档。
请生成规范、完整、专业的测试意图编制内容。"""

    try:
        ai_service = get_ai_service()
        content = await ai_service.generate(prompt=prompt, context=context)
        return GenerateResponse(content=content, success=True)
    except Exception as e:
        return GenerateResponse(
            content="",
            success=False,
            error=str(e)
        )


@router.post("/generate-plan", response_model=GenerateResponse)
async def generate_test_plan(request: PlanGenerateRequest):
    """
    根据前端输入的测试需求文本，调用 planner 生成测试规划。
    直接返回生成的 Markdown 内容（不保存文件）。
    """
    from planner.planning_agent.planner import build_plan_from_text

    try:
        # 直接调用规划生成函数，返回 Markdown 内容
        md_content = build_plan_from_text(
            requirement_text=request.requirement_text,
            source_doc=request.source_name or "用户输入"
        )
        
        if not md_content:
            return GenerateResponse(
                content="",
                success=False,
                error="生成内容为空，请检查输入内容或后端服务"
            )
        
        return GenerateResponse(content=md_content, success=True)
    except Exception as e:
        import traceback
        traceback.print_exc()
        return GenerateResponse(
            content="",
            success=False,
            error=f"规划生成失败: {str(e)}"
        )


@router.post("/extract-steps", response_model=ExtractStepsResponse)
async def extract_steps_from_intent(request: ExtractStepsRequest):
    """
    从意图编制内容中提取测试步骤和参数。
    使用 LLM 解析文本并返回结构化的步骤数据。
    """
    try:
        ai_service = get_ai_service()
        
        # 构建提取 prompt
        extract_prompt = f"""请从以下测试规划内容中提取测试步骤和仪器参数。

请严格按照以下 JSON 格式返回，不要添加任何其他内容：
```json
{{
  "steps": [
    {{
      "id": 1,
      "name": "步骤名称",
      "purpose": "步骤目的",
      "deviceParams": {{
        "程控电源参数": {{
          "输出电压": "28V",
          "输出电流": "6A"
        }},
        "频谱分析仪参数": {{
          "中心频率": "2.7GHz",
          "扫宽Span": "500MHz"
        }}
      }}
    }}
  ]
}}
```

设备类别必须使用以下名称之一：
- 程控电源参数
- 功率探头参数
- 频谱分析仪参数
- 矢量网络分析仪参数
- 矢量信号分析仪参数
- 矢量信号源参数
- 示波器参数
- 信号源基础参数

测试规划内容：
{request.content}
"""

        context = """你是一个专业的测试工程师，擅长从测试规划文档中提取结构化的测试步骤和仪器参数。
请仔细分析文档，提取每个测试步骤及其涉及的仪器配置参数。
只返回 JSON 格式的结果，不要添加任何解释或其他文字。"""

        response_text = await ai_service.generate(prompt=extract_prompt, context=context)
        
        # 解析 JSON 响应
        import json
        import re
        
        # 尝试从响应中提取 JSON
        json_match = re.search(r'```json\s*([\s\S]*?)\s*```', response_text)
        if json_match:
            json_str = json_match.group(1)
        else:
            # 尝试直接解析整个响应
            json_str = response_text.strip()
            # 清理可能的 markdown 代码块标记
            json_str = re.sub(r'^```\w*\n?', '', json_str)
            json_str = re.sub(r'\n?```$', '', json_str)
        
        try:
            parsed = json.loads(json_str)
            steps_data = parsed.get("steps", [])
        except json.JSONDecodeError:
            # 如果解析失败，返回空列表
            return ExtractStepsResponse(
                success=False,
                steps=[],
                error="无法解析 LLM 返回的 JSON 格式"
            )
        
        # 构建响应，过滤掉空值
        def filter_device_params(params: dict) -> dict:
            """过滤掉空的设备参数"""
            if not params or not isinstance(params, dict):
                return {}
            filtered = {}
            for device_category, device_params in params.items():
                if not device_params or not isinstance(device_params, dict):
                    continue
                # 过滤掉空值的参数
                non_empty_params = {
                    k: v for k, v in device_params.items() 
                    if v and isinstance(v, str) and v.strip()
                }
                if non_empty_params:
                    filtered[device_category] = non_empty_params
            return filtered
        
        steps = []
        for step in steps_data:
            filtered_params = filter_device_params(step.get("deviceParams", {}))
            steps.append(ExtractedStepData(
                id=step.get("id", len(steps) + 1),
                name=step.get("name", f"步骤 {len(steps) + 1}"),
                purpose=step.get("purpose", ""),
                deviceParams=filtered_params
            ))
        
        return ExtractStepsResponse(
            success=True,
            steps=steps
        )
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return ExtractStepsResponse(
            success=False,
            steps=[],
            error=f"步骤提取失败: {str(e)}"
        )


@router.get("/check-records/{intent_id}")
async def get_check_records(
    intent_id: int,
    db: AsyncSession = Depends(get_db)
):
    """获取意图的所有检查记录"""
    query = select(AICheckRecord).where(
        AICheckRecord.intent_id == intent_id
    ).order_by(AICheckRecord.checked_at.desc())

    result = await db.execute(query)
    records = result.scalars().all()

    return {
        "intent_id": intent_id,
        "total": len(records),
        "records": [
            {
                "id": r.id,
                "check_result": r.check_result,
                "suggestions": r.suggestions,
                "checked_at": r.checked_at.isoformat()
            }
            for r in records
        ]
    }