Files
rag_agent/rag-web-ui/backend/app/api/api_v1/consistency.py

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"'},
)