code_engine.py

import time

from schemas import CodeXRequest, CodeXResponse, CodeTaskType
from code_router import detect_task_type
from prompt_builder import build_prompt
from model_client import model_client
from response_formatter import build_response, build_error_response
from code_retriever import retrieve_code_evidence
from config import settings


def should_use_retrieval(task_type: CodeTaskType, request: CodeXRequest) -> bool:
    if not request.use_retrieval:
        return False

    if task_type == CodeTaskType.FIX and settings.ENABLE_RETRIEVAL_FOR_FIX:
        return True

    if task_type == CodeTaskType.GENERATE and settings.ENABLE_RETRIEVAL_FOR_GENERATE:
        return True

    if task_type == CodeTaskType.EXPLAIN and settings.ENABLE_RETRIEVAL_FOR_EXPLAIN:
        return True

    return False


def get_retrieved_evidence(task_type: CodeTaskType, request: CodeXRequest):
    return retrieve_code_evidence(task_type, request)


def process_codex_request(request: CodeXRequest) -> CodeXResponse:
    start_time = time.perf_counter()

    try:
        task_type = detect_task_type(
            message=request.message,
            code=request.code,
            error_message=request.error_message,
            mode_hint=request.mode,
        )

        evidence_list = []
        retrieval_used = False

        if should_use_retrieval(task_type, request):
            evidence_list = get_retrieved_evidence(task_type, request)
            retrieval_used = len(evidence_list) > 0

        prompt = build_prompt(
            task_type=task_type,
            message=request.message,
            code=request.code,
            error_message=request.error_message,
            language=request.language,
            framework=request.framework,
            previous_context=request.previous_context,
            evidence_list=evidence_list,
        )

        model_output, model_used, used_fallback = model_client.generate(prompt)

        processing_time_ms = int((time.perf_counter() - start_time) * 1000)

        return build_response(
            task_type=task_type,
            model_output=model_output,
            model_used=model_used,
            used_fallback=used_fallback,
            retrieval_used=retrieval_used,
            source_count=len(evidence_list),
            processing_time_ms=processing_time_ms,
            original_code=request.code,
            sources=evidence_list,
        )

    except Exception as e:
        processing_time_ms = int((time.perf_counter() - start_time) * 1000)
        fallback_task = request.mode if request.mode else CodeTaskType.UNKNOWN

        return build_error_response(
            task_type=fallback_task,
            error_message=str(e),
            processing_time_ms=processing_time_ms,
        )