AutoGPT-Code-Ability/codex/requirements/database.py

from prisma.models import Specification
from prisma.types import SpecificationCreateInput

from codex.api_model import (
    Identifiers,
    Pagination,
    SpecificationResponse,
    SpecificationsListResponse,
)
from codex.requirements.model import ApplicationRequirements


async def create_spec(ids: Identifiers, spec: ApplicationRequirements) -> Specification:
    routes = []

    if not spec.api_routes:
        raise ValueError("No routes found in the specification")
    for route in spec.api_routes:
        create_request = {
            "name": route.request_model.name,
            "description": route.request_model.description,
            "Params": {
                "create": [
                    {
                        "name": param.name,
                        "description": param.description,
                        "paramType": param.param_type,
                    }
                    for param in route.request_model.params
                ],
            },
        }
        create_response = {
            "name": route.response_model.name,
            "description": route.response_model.description,
            "Params": {
                "create": [
                    {
                        "name": param.name,
                        "description": param.description,
                        "paramType": param.param_type,
                    }
                    for param in route.response_model.params
                ],
            },
        }

        create_db = None
        if route.database_schema:
            create_db = {
                "description": route.database_schema.description,
                "DatabaseTables": {
                    "create": [
                        {
                            "description": table.description,
                            "definition": table.definition,
                        }
                        for table in route.database_schema.tables
                    ],
                },
            }
        create_route = {
            "method": route.method,
            "path": route.path,
            "functionName": route.function_name,
            "AccessLevel": route.access_level.value,
            "description": route.description,
            "RequestObject": {"create": create_request},
            "ResponseObject": {"create": create_response},
        }
        if create_db:
            create_route["DatabaseSchema"] = {"create": create_db}
        routes.append(create_route)

    create_spec = SpecificationCreateInput(
        name=spec.name,
        context=spec.context,
        User={"connect": {"id": ids.user_id}},
        Application={"connect": {"id": ids.app_id}},
        ApiRouteSpecs={"create": routes},
    )

    new_spec = await Specification.prisma().create(
        data=create_spec,
        include={
            "ApiRouteSpecs": {
                "include": {
                    "RequestObject": {"include": {"Params": True}},
                    "ResponseObject": {"include": {"Params": True}},
                    "DatabaseSchema": {"include": {"DatabaseTables": True}},
                }
            }
        },
    )
    return new_spec


async def get_specification(user_id: str, app_id: str, spec_id: str) -> Specification:
    specification = await Specification.prisma().find_first_or_raise(
        where={
            "id": spec_id,
            "userId": user_id,
            "applicationId": app_id,
        },
        include={
            "ApiRouteSpecs": {
                "include": {
                    "RequestObject": {"include": {"Params": True}},
                    "ResponseObject": {"include": {"Params": True}},
                }
            }
        },
    )

    return specification


async def get_latest_specification(user_id: str, app_id: str) -> Specification:
    """
    Gets the latest specification for the given user and app

    Args:
        user_id (str): uuid of the user
        app_id (str): uuid of the application

    Returns:
        str: uuid of the latest specification for the user and app
    """
    specification = await Specification.prisma().find_many(
        where={
            "userId": user_id,
            "applicationId": app_id,
            "deleted": False,
        },
        include={
            "ApiRouteSpecs": {
                "include": {
                    "RequestObject": {"include": {"Params": True}},
                    "ResponseObject": {"include": {"Params": True}},
                }
            }
        },
        order={"createdAt": "desc"},
    )
    if not specification:
        raise ValueError("No specifications found for the user and app")

    return specification[0]


async def delete_specification(spec_id: str) -> None:
    await Specification.prisma().update(
        where={"id": spec_id},
        data={"deleted": True},
    )


async def list_specifications(
    user_id: str, app_id: str, page: int, page_size: int
) -> SpecificationsListResponse:
    skip = (page - 1) * page_size
    total_items = await Specification.prisma().count(
        where={"userId": user_id, "applicationId": app_id, "deleted": False}
    )
    if total_items > 0:
        specs = await Specification.prisma().find_many(
            where={"userId": user_id, "applicationId": app_id, "deleted": False},
            include={
                "ApiRouteSpecs": {
                    "include": {
                        "RequestObject": {"include": {"Params": True}},
                        "ResponseObject": {"include": {"Params": True}},
                    }
                }
            },
            skip=skip,
            take=page_size,
        )

        total_pages = (total_items + page_size - 1) // page_size

        specs_response = [
            SpecificationResponse.from_specification(spec) for spec in specs
        ]

        pagination = Pagination(
            total_items=total_items,
            total_pages=total_pages,
            current_page=page,
            page_size=page_size,
        )

        return SpecificationsListResponse(specs=specs_response, pagination=pagination)
    else:
        return SpecificationsListResponse(
            specs=[],
            pagination=Pagination(
                total_items=0, total_pages=0, current_page=0, page_size=0
            ),
        )