Spaces:

EATosin
/

Zeta

Sleeping

Zeta / src /api /main.py

Update src/api/main.py

2a2b2fc verified 10 months ago

4.46 kB

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)

	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)