app/main.py

# entry point aka building FAST_API here

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from app.predictor import classifier, guide_generator

app = FastAPI(title="GitGud AI Service")


# Data Model: Matches what NestJS (server-side[refer to visualization.services.ts]) sends
class FileRequest(BaseModel):
    fileName: str
    content: str | None = None


class GuideRequest(BaseModel):
    repoName: str
    filePaths: list[str]


@app.get("/")
def health_check():
    """
    Simple check to see if the server is alive and which GPU it's using.
    """
    return {
        "status": "online",
        "model": "microsoft/codebert-base",
        "device": classifier.device,
    }


# first FAST_API with endpoint('/classify') called in [visualization.services.ts]
# @param {*} file
# @return {*} layerd based classified_info along with file-name
@app.post("/classify")
async def classify_file(request: FileRequest):
    try:
        # calling the predict function of our classifier to determine which layer it belongs
        # returns { label, confidence, embedding }
        result = classifier.predict(request.fileName, request.content)
        return {
            "fileName": request.fileName, 
            "layer": result["label"],
            "confidence": result["confidence"],
            "embedding": result["embedding"]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/generate-guide")
async def generate_guide(request: GuideRequest):
    try:
        markdown = guide_generator.generate_markdown(request.repoName, request.filePaths)
        return {"markdown": markdown}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn

    # Runs on localhost:8000
    uvicorn.run(app, host="0.0.0.0", port=8000)