pr-agent/pr_agent/servers/gerrit_server.py

import copy
import logging
import sys
from enum import Enum
from json import JSONDecodeError

import uvicorn
from fastapi import APIRouter, FastAPI, HTTPException
from pydantic import BaseModel
from starlette.middleware import Middleware
from starlette_context import context
from starlette_context.middleware import RawContextMiddleware

from pr_agent.agent.pr_agent import PRAgent
from pr_agent.config_loader import global_settings, get_settings

logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
router = APIRouter()


class Action(str, Enum):
    review = "review"
    describe = "describe"
    ask = "ask"
    improve = "improve"


class Item(BaseModel):
    refspec: str
    project: str
    msg: str = None


@router.post("/api/v1/gerrit/{action}")
async def handle_gerrit_request(action: Action, item: Item):
    logging.debug("Received a Gerrit request")
    context["settings"] = copy.deepcopy(global_settings)

    agent = PRAgent()
    pr_url = f"{item.project}:{item.refspec}"
    if action == Action.review:
        await agent.handle_request(pr_url, "/review")
    elif action == Action.describe:
        await agent.handle_request(pr_url, "/describe")
    elif action == Action.improve:
        await agent.handle_request(pr_url, "/improve")
    elif action == Action.ask:
        if not item.msg:
            return HTTPException(
                status_code=400,
                detail="msg is required for ask command"
            )
        await agent.handle_request(pr_url, f"/ask {item.msg.strip()}")


async def get_body(request):
    try:
        body = await request.json()
    except JSONDecodeError as e:
        logging.error("Error parsing request body", e)
        return {}
    return body


@router.get("/")
async def root():
    return {"status": "ok"}


def start():
    # to prevent adding help messages with the output
    get_settings().set("CONFIG.CLI_MODE", True)
    middleware = [Middleware(RawContextMiddleware)]
    app = FastAPI(middleware=middleware)
    app.include_router(router)

    uvicorn.run(app, host="0.0.0.0", port=3000)


if __name__ == '__main__':
    start()