339 lines
12 KiB
Python
339 lines
12 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from datetime import datetime
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, List
|
||
|
|
|
||
|
|
from fastapi import APIRouter, BackgroundTasks, Depends, File, Form, HTTPException, Query, UploadFile
|
||
|
|
from fastapi.responses import Response
|
||
|
|
from sqlalchemy.orm import Session
|
||
|
|
|
||
|
|
from app.core.security import get_current_user
|
||
|
|
from app.db.session import get_db
|
||
|
|
from app.models.tooling import ConsistencyResult
|
||
|
|
from app.models.user import User
|
||
|
|
from app.schemas.consistency import (
|
||
|
|
AutoConsistencyJobCreateResponse,
|
||
|
|
AutoConsistencyJobStatusResponse,
|
||
|
|
CodeKnowledgeBaseCreate,
|
||
|
|
CodeKnowledgeBaseResponse,
|
||
|
|
CodeKnowledgeBaseUploadResponse,
|
||
|
|
CodeQuestionRequest,
|
||
|
|
CodeQuestionResponse,
|
||
|
|
ConsistencyJobCreate,
|
||
|
|
ConsistencyJobCreateResponse,
|
||
|
|
ConsistencyJobResponse,
|
||
|
|
ConsistencyResultResponse,
|
||
|
|
)
|
||
|
|
from app.services.consistency.exporter import export_excel, export_json, export_markdown
|
||
|
|
from app.services.consistency_job_service import (
|
||
|
|
AUTO_UPLOAD_ROOT,
|
||
|
|
CODE_UPLOAD_ROOT,
|
||
|
|
ask_code_kb,
|
||
|
|
create_code_kb,
|
||
|
|
create_auto_consistency_tool_job,
|
||
|
|
create_consistency_job,
|
||
|
|
create_uploaded_code_kb,
|
||
|
|
get_owned_auto_job,
|
||
|
|
get_owned_code_kb,
|
||
|
|
get_owned_consistency_job,
|
||
|
|
list_code_kbs,
|
||
|
|
list_consistency_jobs,
|
||
|
|
result_model_to_export_dict,
|
||
|
|
run_auto_consistency_job,
|
||
|
|
run_code_kb_build,
|
||
|
|
run_consistency_job,
|
||
|
|
safe_upload_name,
|
||
|
|
save_uploaded_bytes,
|
||
|
|
)
|
||
|
|
from app.services.model_config import ModelConfigService
|
||
|
|
|
||
|
|
router = APIRouter()
|
||
|
|
|
||
|
|
|
||
|
|
def datetime_path() -> str:
|
||
|
|
return datetime.utcnow().strftime("%Y%m%d%H%M%S%f")
|
||
|
|
|
||
|
|
|
||
|
|
async def _save_code_uploads(files: List[UploadFile], target_dir: Path) -> str:
|
||
|
|
if not files:
|
||
|
|
raise HTTPException(status_code=400, detail="No code files uploaded.")
|
||
|
|
source_dir = target_dir
|
||
|
|
extracted_dirs: List[Path] = []
|
||
|
|
for file in files:
|
||
|
|
content = await file.read()
|
||
|
|
if not content:
|
||
|
|
continue
|
||
|
|
saved = save_uploaded_bytes(target_dir, safe_upload_name(file.filename), content)
|
||
|
|
if saved.is_dir():
|
||
|
|
extracted_dirs.append(saved)
|
||
|
|
if len(files) == 1 and extracted_dirs:
|
||
|
|
source_dir = extracted_dirs[0]
|
||
|
|
return str(source_dir.resolve())
|
||
|
|
|
||
|
|
|
||
|
|
async def _save_requirement_upload(file: UploadFile, target_dir: Path) -> str:
|
||
|
|
safe_name = safe_upload_name(file.filename)
|
||
|
|
if Path(safe_name).suffix.lower() not in {".pdf", ".docx"}:
|
||
|
|
raise HTTPException(status_code=400, detail="Requirement file must be .pdf or .docx.")
|
||
|
|
content = await file.read()
|
||
|
|
if not content:
|
||
|
|
raise HTTPException(status_code=400, detail="Requirement file is empty.")
|
||
|
|
saved = save_uploaded_bytes(target_dir, safe_name, content)
|
||
|
|
return str(saved.resolve())
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/code-kbs", response_model=CodeKnowledgeBaseResponse)
|
||
|
|
async def register_code_kb(
|
||
|
|
payload: CodeKnowledgeBaseCreate,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
try:
|
||
|
|
return create_code_kb(db, current_user.id, payload)
|
||
|
|
except (FileNotFoundError, RuntimeError, ValueError) as exc:
|
||
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/code-kbs/upload", response_model=CodeKnowledgeBaseUploadResponse)
|
||
|
|
async def upload_and_build_code_kb(
|
||
|
|
background_tasks: BackgroundTasks,
|
||
|
|
name: str = Form(...),
|
||
|
|
use_semantic: bool = Form(True),
|
||
|
|
files: List[UploadFile] = File(...),
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
if use_semantic:
|
||
|
|
try:
|
||
|
|
ModelConfigService.require_active_config(db, current_user.id)
|
||
|
|
except ValueError as exc:
|
||
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||
|
|
target_dir = CODE_UPLOAD_ROOT / str(current_user.id) / datetime_path()
|
||
|
|
code_source_dir = await _save_code_uploads(files, target_dir)
|
||
|
|
output_dir = str((target_dir / "artifacts").resolve())
|
||
|
|
code_kb = create_uploaded_code_kb(
|
||
|
|
db=db,
|
||
|
|
user_id=current_user.id,
|
||
|
|
name=name,
|
||
|
|
project_path=code_source_dir,
|
||
|
|
output_dir=output_dir,
|
||
|
|
)
|
||
|
|
background_tasks.add_task(run_code_kb_build, code_kb.id, use_semantic)
|
||
|
|
return {"id": code_kb.id, "status": code_kb.status}
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/code-kbs", response_model=List[CodeKnowledgeBaseResponse])
|
||
|
|
async def get_code_kbs(
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
return list_code_kbs(db, current_user.id)
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/code-kbs/{code_kb_id}", response_model=CodeKnowledgeBaseResponse)
|
||
|
|
async def get_code_kb(
|
||
|
|
code_kb_id: int,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
code_kb = get_owned_code_kb(db, current_user.id, code_kb_id)
|
||
|
|
if not code_kb:
|
||
|
|
raise HTTPException(status_code=404, detail="Code knowledge base not found.")
|
||
|
|
return code_kb
|
||
|
|
|
||
|
|
|
||
|
|
@router.delete("/code-kbs/{code_kb_id}")
|
||
|
|
async def delete_code_kb(
|
||
|
|
code_kb_id: int,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
code_kb = get_owned_code_kb(db, current_user.id, code_kb_id)
|
||
|
|
if not code_kb:
|
||
|
|
raise HTTPException(status_code=404, detail="Code knowledge base not found.")
|
||
|
|
db.delete(code_kb)
|
||
|
|
db.commit()
|
||
|
|
return {"message": "deleted"}
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/code-kbs/{code_kb_id}/ask", response_model=CodeQuestionResponse)
|
||
|
|
async def ask_code_kb_api(
|
||
|
|
code_kb_id: int,
|
||
|
|
payload: CodeQuestionRequest,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
code_kb = get_owned_code_kb(db, current_user.id, code_kb_id)
|
||
|
|
if not code_kb:
|
||
|
|
raise HTTPException(status_code=404, detail="Code knowledge base not found.")
|
||
|
|
try:
|
||
|
|
model_profile = ModelConfigService.require_active_config(db, current_user.id)
|
||
|
|
ModelConfigService.touch_last_used(db, model_profile)
|
||
|
|
return ask_code_kb(
|
||
|
|
code_kb=code_kb,
|
||
|
|
question=payload.question,
|
||
|
|
top_k=payload.top_k,
|
||
|
|
min_similarity=payload.min_similarity,
|
||
|
|
use_llm=payload.use_llm,
|
||
|
|
model_profile=model_profile,
|
||
|
|
)
|
||
|
|
except (RuntimeError, ValueError) as exc:
|
||
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/jobs", response_model=ConsistencyJobCreateResponse)
|
||
|
|
async def create_job(
|
||
|
|
background_tasks: BackgroundTasks,
|
||
|
|
payload: ConsistencyJobCreate,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
try:
|
||
|
|
ModelConfigService.require_active_config(db, current_user.id)
|
||
|
|
job = create_consistency_job(db, current_user.id, payload)
|
||
|
|
except ValueError as exc:
|
||
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||
|
|
background_tasks.add_task(run_consistency_job, job.id)
|
||
|
|
return {"job_id": job.id, "status": job.status}
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/auto-jobs", response_model=AutoConsistencyJobCreateResponse)
|
||
|
|
async def create_auto_job(
|
||
|
|
background_tasks: BackgroundTasks,
|
||
|
|
requirement_file: UploadFile = File(...),
|
||
|
|
code_files: List[UploadFile] = File(...),
|
||
|
|
code_kb_name: str = Form("uploaded-code-kb"),
|
||
|
|
top_k: int = Form(8),
|
||
|
|
max_call_hops: int = Form(2),
|
||
|
|
min_similarity: float = Form(0.55),
|
||
|
|
use_llm: bool = Form(True),
|
||
|
|
use_semantic: bool = Form(True),
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
try:
|
||
|
|
ModelConfigService.require_active_config(db, current_user.id)
|
||
|
|
except ValueError as exc:
|
||
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||
|
|
timestamp = datetime_path()
|
||
|
|
target_dir = AUTO_UPLOAD_ROOT / str(current_user.id) / timestamp
|
||
|
|
requirement_path = await _save_requirement_upload(requirement_file, target_dir / "requirement")
|
||
|
|
code_source_dir = await _save_code_uploads(code_files, target_dir / "code")
|
||
|
|
tool_job = create_auto_consistency_tool_job(
|
||
|
|
db=db,
|
||
|
|
user_id=current_user.id,
|
||
|
|
requirement_file_path=requirement_path,
|
||
|
|
requirement_file_name=safe_upload_name(requirement_file.filename),
|
||
|
|
code_source_dir=code_source_dir,
|
||
|
|
code_kb_name=code_kb_name,
|
||
|
|
top_k=top_k,
|
||
|
|
max_call_hops=max_call_hops,
|
||
|
|
min_similarity=min_similarity,
|
||
|
|
use_llm=use_llm,
|
||
|
|
use_semantic=use_semantic,
|
||
|
|
)
|
||
|
|
background_tasks.add_task(run_auto_consistency_job, tool_job.id)
|
||
|
|
return {"job_id": tool_job.id, "status": tool_job.status}
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/auto-jobs/{job_id}", response_model=AutoConsistencyJobStatusResponse)
|
||
|
|
async def get_auto_job(
|
||
|
|
job_id: int,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
job = get_owned_auto_job(db, current_user.id, job_id)
|
||
|
|
if not job:
|
||
|
|
raise HTTPException(status_code=404, detail="Auto consistency job not found.")
|
||
|
|
summary = job.output_summary or {}
|
||
|
|
return {
|
||
|
|
"job_id": job.id,
|
||
|
|
"status": job.status,
|
||
|
|
"error_message": job.error_message,
|
||
|
|
"current_step": summary.get("current_step"),
|
||
|
|
"srs_extraction_id": summary.get("srs_extraction_id"),
|
||
|
|
"code_kb_id": summary.get("code_kb_id"),
|
||
|
|
"consistency_job_id": summary.get("consistency_job_id"),
|
||
|
|
"created_at": job.created_at,
|
||
|
|
"updated_at": job.updated_at,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/jobs", response_model=List[ConsistencyJobResponse])
|
||
|
|
async def get_jobs(
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
return list_consistency_jobs(db, current_user.id)
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/jobs/{job_id}", response_model=ConsistencyJobResponse)
|
||
|
|
async def get_job(
|
||
|
|
job_id: int,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
job = get_owned_consistency_job(db, current_user.id, job_id)
|
||
|
|
if not job:
|
||
|
|
raise HTTPException(status_code=404, detail="Consistency job not found.")
|
||
|
|
return job
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/jobs/{job_id}/results", response_model=List[ConsistencyResultResponse])
|
||
|
|
async def get_job_results(
|
||
|
|
job_id: int,
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Any:
|
||
|
|
job = get_owned_consistency_job(db, current_user.id, job_id)
|
||
|
|
if not job:
|
||
|
|
raise HTTPException(status_code=404, detail="Consistency job not found.")
|
||
|
|
return (
|
||
|
|
db.query(ConsistencyResult)
|
||
|
|
.filter(ConsistencyResult.job_id == job.id)
|
||
|
|
.order_by(ConsistencyResult.id)
|
||
|
|
.all()
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/jobs/{job_id}/export")
|
||
|
|
async def export_job_results(
|
||
|
|
job_id: int,
|
||
|
|
format: str = Query(default="json", pattern="^(json|markdown|md|excel|xlsx)$"),
|
||
|
|
db: Session = Depends(get_db),
|
||
|
|
current_user: User = Depends(get_current_user),
|
||
|
|
) -> Response:
|
||
|
|
job = get_owned_consistency_job(db, current_user.id, job_id)
|
||
|
|
if not job:
|
||
|
|
raise HTTPException(status_code=404, detail="Consistency job not found.")
|
||
|
|
rows = (
|
||
|
|
db.query(ConsistencyResult)
|
||
|
|
.filter(ConsistencyResult.job_id == job.id)
|
||
|
|
.order_by(ConsistencyResult.id)
|
||
|
|
.all()
|
||
|
|
)
|
||
|
|
payload = [result_model_to_export_dict(row) for row in rows]
|
||
|
|
if format in {"markdown", "md"}:
|
||
|
|
content = export_markdown(payload).encode("utf-8")
|
||
|
|
return Response(
|
||
|
|
content,
|
||
|
|
media_type="text/markdown; charset=utf-8",
|
||
|
|
headers={"Content-Disposition": f'attachment; filename="consistency-job-{job.id}.md"'},
|
||
|
|
)
|
||
|
|
if format in {"excel", "xlsx"}:
|
||
|
|
try:
|
||
|
|
content = export_excel(payload)
|
||
|
|
except RuntimeError as exc:
|
||
|
|
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
||
|
|
return Response(
|
||
|
|
content,
|
||
|
|
media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||
|
|
headers={"Content-Disposition": f'attachment; filename="consistency-job-{job.id}.xlsx"'},
|
||
|
|
)
|
||
|
|
return Response(
|
||
|
|
export_json(payload),
|
||
|
|
media_type="application/json; charset=utf-8",
|
||
|
|
headers={"Content-Disposition": f'attachment; filename="consistency-job-{job.id}.json"'},
|
||
|
|
)
|