File size: 4,458 Bytes
39c9555 19cb1fb 5634a6b 19cb1fb 39c9555 5634a6b 19cb1fb 1c5042f 19cb1fb 1c5042f 19cb1fb 1c5042f 19cb1fb 5634a6b 19cb1fb 2a2b2fc 19cb1fb 39c9555 19cb1fb e5c6edf 5634a6b 19cb1fb e5c6edf 19cb1fb e5c6edf 19cb1fb 5634a6b 19cb1fb e5c6edf 1c5042f e5c6edf 1c5042f 19cb1fb e5c6edf 19cb1fb e5c6edf 19cb1fb 5634a6b 19cb1fb 1c5042f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
from typing import Dict, Any
from src.core.llm_engine import LLMEngine
from src.ml.edge_case_detector import EdgeCaseDetector
from src.core.code_generator import SeleniumGenerator
from src.api.models import GenerateRequest, TestSuiteResponse, CodeGenRequest, CodeResponse
import uvicorn
import uuid
import os
from loguru import logger
# --- LIFESPAN MANAGER ---
ml_resources: Dict[str, Any] = {}
@asynccontextmanager
async def lifespan(app: FastAPI):
logger.info("⚡ System Startup: Loading AI Engines...")
try:
ml_resources["llm"] = LLMEngine()
ml_resources["ml"] = EdgeCaseDetector()
ml_resources["codegen"] = SeleniumGenerator()
logger.info("✅ Engines Online.")
except Exception as e:
logger.critical(f"❌ Engine Startup Failed: {e}")
raise e
yield
ml_resources.clear()
logger.info("🛑 System Shutdown.")
# --- APP CONFIG ---
app = FastAPI(
title="Zeta AI API",
version="1.0.0",
lifespan=lifespan
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/health")
async def health_check():
return {
"status": "active",
"engines_loaded": list(ml_resources.keys())
}
@app.post("/generate", response_model=TestSuiteResponse)
async def generate_tests(request: GenerateRequest):
logger.info(f"Processing generation request ({len(request.requirements_text)} chars)")
if "llm" not in ml_resources:
raise HTTPException(status_code=503, detail="AI Engines not ready")
try:
# 1. Generate Raw Tests (Contains Title, Steps, Actions)
llm_engine: LLMEngine = ml_resources["llm"]
raw_tests = await llm_engine.generate_test_cases(request.requirements_text)
# 2. Prepare Data for ML Engine
# We temporarily add a 'text' field for the ML model to analyze
for test in raw_tests:
if "text" not in test:
# Combine fields to give the ML more context
test["text"] = f"{test.get('title', '')} {' '.join(test.get('steps', []))}"
ml_engine: EdgeCaseDetector = ml_resources["ml"]
# 3. Get ML Analysis
analysis_objects = await ml_engine.analyze_complexity(raw_tests)
# 4. MERGE LOGIC (The Fix)
# Instead of returning analysis_objects, we inject the risk data back into raw_tests
final_test_cases = []
for original, analysis in zip(raw_tests, analysis_objects):
# Keep original fields (Title, Steps, Actions)
merged = original.copy()
# Inject ML metrics
analysis_dict = analysis.model_dump()
merged["risk_analysis"] = {
"is_edge_case": analysis_dict["is_edge_case"],
"risk_level": analysis_dict["risk_level"],
"complexity_score": analysis_dict["complexity_score"],
"risk_sources": analysis_dict["risk_sources"]
}
final_test_cases.append(merged)
return TestSuiteResponse(
suite_id=str(uuid.uuid4()),
test_cases=final_test_cases,
meta={"source": "Gemini 2.5", "ml_validation": True}
)
except Exception as e:
logger.error(f"Generation logic failed: {e}")
# Return error details to help debugging
raise HTTPException(status_code=500, detail=f"Internal Error: {str(e)}")
@app.post("/codegen", response_model=CodeResponse)
async def generate_code(request: CodeGenRequest):
if "codegen" not in ml_resources:
raise HTTPException(status_code=503, detail="Codegen Engine not ready")
try:
codegen_engine: SeleniumGenerator = ml_resources["codegen"]
code = codegen_engine.generate_test_script(request.test_plan)
return CodeResponse(
filename=f"test_{uuid.uuid4().hex[:8]}.py",
python_code=code
)
except Exception as e:
logger.error(f"Codegen failed: {e}")
raise HTTPException(status_code=500, detail="Code generation failed")
if __name__ == "__main__":
port = int(os.getenv("PORT", 8000))
uvicorn.run("src.api.main:app", host="0.0.0.0", port=port, reload=True) |