AutoGPT-Code-Ability/codex/develop/compile.py

import ast
import logging
import re
from datetime import datetime
from typing import List, Set

from packaging import version
from prisma.models import (
    CompiledRoute,
    CompletedApp,
    Function,
    ObjectField,
    ObjectType,
    Package,
    Specification,
)
from prisma.types import CompiledRouteUpdateInput
from pydantic import BaseModel

from codex.common.database import INCLUDE_FIELD, INCLUDE_FUNC, get_database_schema
from codex.common.exec_external_tool import DEFAULT_DEPS
from codex.common.types import normalize_type
from codex.deploy.model import Application
from codex.develop.code_validation import CodeValidator
from codex.develop.database import get_compiled_route, get_deliverable
from codex.develop.function import generate_object_template
from codex.develop.model import Package as PackageModel

logger = logging.getLogger(__name__)


class CompiledFunction(BaseModel):
    packages: List[Package]
    imports: List[str]
    code: str
    pydantic_models: List[str] = []


class ComplicationFailure(Exception):
    pass


async def compile_route(
    compiled_route_id: str,
    route_root_func: Function,
    spec: Specification,
    available_functions: dict[str, Function] = {},
    available_objects: dict[str, ObjectType] = {},
) -> CompiledRoute:
    """
    Compiles a route by generating a CompiledRoute object.

    Args:
        compiled_route_id (str): The ID of the compiled route.
        route_root_func (Function): The root function of the route.
    Returns:
        CompiledRoute: The compiled route object.
    """
    compiled_function = await recursive_compile_route(route_root_func, set())

    unique_packages = {
        package.id: PackageModel(
            package_name=package.packageName,
            version=package.version,
            specifier=package.specifier,
        )
        for package in compiled_function.packages
    }
    code = "\n".join(compiled_function.imports)
    code += "\n\n"
    code += "\n\n".join(compiled_function.pydantic_models)
    code += "\n\n"
    code += compiled_function.code

    compiled_route = await get_compiled_route(compiled_route_id)
    database_schema = get_database_schema(spec)
    available_functions[route_root_func.functionName] = route_root_func
    # Run the auto-fixers
    formatted_code = await CodeValidator(
        compiled_route_id=compiled_route_id,
        function_name=route_root_func.functionName,
        database_schema=database_schema,
        available_functions=available_functions,
        available_objects=available_objects,
    ).reformat_code(
        code,
        list(unique_packages.values()),
    )

    data = CompiledRouteUpdateInput(
        Packages={"connect": [{"id": package_id} for package_id in unique_packages]},
        compiledCode=formatted_code,
    )
    compiled_route = await CompiledRoute.prisma().update(
        where={"id": compiled_route_id}, data=data
    )
    if compiled_route is None:
        raise ComplicationFailure(
            f"Failed to update compiled route {compiled_route_id}"
        )
    return compiled_route


def add_full_import_parth_to_custom_types(module_name: str, arg: ObjectField) -> str:
    """
    Adds the full import path to custom types in the given argument.

    Args:
        module_name (str): The name of the module.
        arg (ObjectField): The argument to process.

    Returns:
        str: The modified argument type with the full import path.

    """
    if not arg.RelatedTypes:
        return arg.typeName

    ret_type = arg.typeName
    logger.debug(f"Arg name: {arg.name}, arg type: {arg.typeName}")

    # For each related type, replace the type name with the full import path
    renamed_types: dict[str, str] = {}

    for t in arg.RelatedTypes:
        if t.isPydantic or t.isEnum:
            renamed_types[t.name] = f"{module_name}.{t.name}"

    ret_type = normalize_type(ret_type, renamed_types)

    return ret_type


async def recursive_compile_route(
    in_function: Function, object_type_ids: Set[str]
) -> CompiledFunction:
    """
    Recursively compiles a function and its child functions
    into a single CompiledFunction object.

    Args:
        ids (Identifiers): The identifiers for the function.
        function (Function): The function to compile.

    Returns:
        CompiledFunction: The compiled function.

    Raises:
        ValueError: If the function code is missing.
    """
    # Can't see how to do recursive lookup with prisma, so I'm checking the next
    # layer down each time. This is a bit of a hack, can be improved later.
    function = await Function.prisma().find_unique_or_raise(
        where={"id": in_function.id},
        include={
            **INCLUDE_FUNC["include"],
            "ParentFunction": INCLUDE_FUNC,
            "ChildFunctions": INCLUDE_FUNC,
            "Packages": True,
        },  # type: ignore
    )
    logger.info(f"⚙️ Compiling function: {function.functionName}")

    if function.functionCode is None:
        raise ValueError(f"Function code is required! {function.functionName}")

    packages = []
    imports = []
    code = []
    model = []

    if function.FunctionArgs is not None:
        for arg in function.FunctionArgs:
            obj_types = await get_object_field_deps(arg, object_type_ids)
            model.extend([generate_object_template(obj_type) for obj_type in obj_types])
            for obj in obj_types:
                imports.extend(obj.importStatements)

    if function.FunctionReturn is not None:
        obj_types = await get_object_field_deps(
            function.FunctionReturn, object_type_ids
        )
        model.extend([generate_object_template(obj_type) for obj_type in obj_types])
        for obj in obj_types:
            imports.extend(obj.importStatements)

    # Child Functions
    if function.ChildFunctions is None:
        raise AssertionError("ChildFunctions should be an array")
    for child_function in function.ChildFunctions:
        compiled_function = await recursive_compile_route(
            child_function, object_type_ids
        )
        packages.extend(compiled_function.packages)
        imports.extend(compiled_function.imports)
        model.extend(compiled_function.pydantic_models)
        code.append(compiled_function.code)

    # Package
    if function.Packages:
        packages.extend(function.Packages)

    # Imports
    imports.extend(function.importStatements)
    imports = sorted(set(imports))

    # Code
    code.append(function.functionCode)

    # Check Code
    try:
        check_code = "\n".join(imports)
        check_code += "\n\n" + "\n\n".join(model)
        check_code += "\n\n" + "\n\n".join(code)
        ast.parse(check_code)
    except Exception as e:
        raise ValueError(f"Syntax error in function code: {e}, {code}")

    return CompiledFunction(
        packages=packages,
        imports=imports,
        code="\n\n".join(code),
        pydantic_models=model,
    )


async def get_object_type_deps(
    obj_type_id: str, object_type_ids: Set[str]
) -> List[ObjectType]:
    # Lookup the object getting all its subfields
    obj = await ObjectType.prisma().find_unique_or_raise(
        where={"id": obj_type_id},
        **INCLUDE_FIELD,  # type: ignore
    )
    if obj.Fields is None:
        raise ValueError(f"ObjectType {obj.name} has no fields.")

    objects: List[ObjectType] = []
    for field in obj.Fields:
        if field.RelatedTypes:
            objects.extend(await get_object_field_deps(field, object_type_ids))

    return objects + [obj]


async def get_object_field_deps(
    field: ObjectField, object_type_ids: Set[str]
) -> List[ObjectType]:
    """
    Process an object field and return the Pydantic classes
    generated from the field's type.

    Args:
        field (ObjectField): The object field to process.
        object_type_ids (Set[str]): A set of object type IDs that
                                    have already been processed.

    Returns:
        List[ObjectType]: The object types generated from the field's type.

    Raises:
        AssertionError: If the field type is None.
    """
    # Lookup the field object getting all its subfields
    field = await ObjectField.prisma().find_unique_or_raise(
        where={"id": field.id},
        include={"RelatedTypes": True},
    )

    if field.RelatedTypes is None:
        raise AssertionError("Field RelatedTypes should be an array")
    types = [t for t in field.RelatedTypes if t.id not in object_type_ids]

    if not types:
        # If the field is a primitive type or we have already processed this object,
        # we don't need to do anything
        logger.debug(
            f"Skipping field {field.name} as it's a primitive type or already processed"
        )
        return []

    logger.debug(f"Processing field {field.name} of type {field.typeName}")
    object_type_ids.update([t.id for t in types])

    # TODO: this can run in parallel
    pydantic_classes = []
    for type in types:
        pydantic_classes.extend(await get_object_type_deps(type.id, object_type_ids))

    return pydantic_classes


def create_server_route_code(compiled_route: CompiledRoute) -> str:
    """
    Create the server route code for a compiled route.

    Args:
        compiled_route (CompiledRoute): The compiled route to create the
                                        server route code for.

    Returns:
        str: The server route code.
    """
    main_function = compiled_route.RootFunction

    if main_function is None:
        raise ValueError("Compiled route must have a root function.")

    return_type = main_function.FunctionReturn
    # if return_type is None:
    #     raise AssertionError("Compiled route must have a return type.")

    args = main_function.FunctionArgs
    if args is None:
        raise AssertionError("Compiled route must have function arguments.")

    route_spec = compiled_route.ApiRouteSpec
    if route_spec is None:
        raise AssertionError("Compiled route must have an API route spec.")

    module_name = f"project.{compiled_route.fileName.replace('.py', '')}"

    is_file_response = False
    response_model = "Response"
    route_response_annotation = "Response"
    if (
        return_type
        and return_type.RelatedTypes
        and return_type.RelatedTypes[0]
        and return_type.RelatedTypes[0].Fields
        and return_type.RelatedTypes[0].Fields[0].typeName == "bytes"
    ):
        is_file_response = True
    elif return_type:
        if return_type.typeName is not None:
            response_model = f"{module_name}.{return_type.typeName} | Response"
            route_response_annotation = return_type.typeName

    # 4. Determine path parameters
    path_params = extract_path_params(route_spec.path)

    func_args_names = set([arg.name for arg in args])
    if not set(path_params).issubset(func_args_names):
        logger.warning(
            f"Path parameters {path_params} not in function arguments {func_args_names}"
        )

    http_verb = str(route_spec.method)
    route_decorator = f"@app.{http_verb.lower()}('{route_spec.path}'"
    if not is_file_response:
        route_decorator += f", response_model={module_name}.{route_response_annotation}"

    route_decorator += ")\n"

    # TODO(SwiftyOS): consider replacing the prefix with a proper import and func call.
    route_function_def = (
        f"async def api_{http_verb.lower()}_{main_function.functionName}("
    )
    route_function_def += ", ".join(
        [
            f"{arg.name}: {add_full_import_parth_to_custom_types(module_name, arg)}"
            for arg in args
            if arg.name and arg.typeName
        ]
    )
    route_function_def += ")"

    return_response = ""
    if is_file_response:
        route_function_def += " -> StreamingResponse:"
        headers = """
        headers = {
        "Content-Disposition": f"attachment; filename={res.file_name}"
    }
        """
        return_response = f"""
        {headers}
        return StreamingResponse(
            content=res.data,
            media_type=res.media_type,
            headers=headers)
    """
    else:
        route_function_def += f" -> {response_model}:"
        return_response = "return res"

    await_str = "await " if "async def" in main_function.template else ""

    function_body = f"""
    \"\"\"
    {route_spec.description}
    \"\"\"
    try:
        res = {await_str}{module_name}.{main_function.functionName}({", ".join([arg.name for arg in args])})
        {return_response}
    except Exception as e:
        logger.exception("Error processing request")
        res = dict()
        res["error"] =  str(e)
        return Response(
            content=jsonable_encoder(res),
            status_code=500,
            media_type="application/json",
        )
    """
    route_code = route_decorator
    route_code += route_function_def
    route_code += function_body
    route_code += "\n\n"

    try:
        ast.parse(route_code)
    except Exception as e:
        raise ComplicationFailure(f"Syntax error in route code: {e}, {route_code}")

    return route_code


def extract_path_params(path: str) -> List[str]:
    """
    Extracts path parameters from a given path.

    Args:
        path (str): The path string containing path parameters.

    Returns:
        List[str]: A list of path parameters extracted from the path.
    """
    matches = re.findall(r"\{(.*?)\}", path)
    return matches


DEFAULT_LIBRARIES = [
    Package(
        packageName=lib,
        specifier="",
        version="",
        id=lib,
        createdAt=datetime.now(),
    )
    for lib in ["fastapi", "pydantic", "uvicorn", "prisma"]
]


async def create_bundle_code(
    completed_app: CompletedApp,
    spec: Specification,
) -> Application:
    """
    Args:
        application (Application): _description_

    Returns:
        Application: _description_
    """
    if not completed_app.companionCompletedAppId:
        return await create_server_code(completed_app, spec)

    server_completed_app = await get_deliverable(completed_app.companionCompletedAppId)
    libraries = DEFAULT_LIBRARIES + [
        Package(
            packageName="nicegui",
            specifier=">=",
            version="1.4.22",
            id="nicegui",
            createdAt=datetime.now(),
        )
    ]

    app = await create_server_code(server_completed_app, spec, libraries)
    app.app_code = await create_ui_code(completed_app)
    app.companion_completed_app = completed_app
    return app


async def create_ui_code(completed_app: CompletedApp) -> str:
    code_imports = [
        "import asyncio",
        "from nicegui import ui",
        "from prisma import Prisma",
    ] + [
        f"from project.{compiled_route.fileName.replace('.py', '')} import *"
        for compiled_route in completed_app.CompiledRoutes or []
    ]

    code_content = [
        "db_client = Prisma(auto_register=True)",
        "asyncio.get_event_loop().run_until_complete(db_client.connect())",
        "ui.run(dark=True)",
    ]

    return "\n".join(code_imports) + "\n\n" + "\n\n".join(code_content)


async def create_server_code(
    completed_app: CompletedApp,
    spec: Specification,
    libraries: List[Package] = DEFAULT_LIBRARIES,
) -> Application:
    name = completed_app.name
    desc = completed_app.description

    server_code_imports = [
        "from fastapi import FastAPI",
        "from fastapi.responses import Response, StreamingResponse",
        "from fastapi.encoders import jsonable_encoder",
        "from prisma import Prisma",
        "from contextlib import asynccontextmanager",
        "import logging",
        "import prisma",
        "import io",
    ]
    server_code_header = f"""logger = logging.getLogger(__name__)

db_client = Prisma(auto_register=True)

@asynccontextmanager
async def lifespan(app: FastAPI):
    await db_client.connect()
    yield
    await db_client.disconnect()

app = FastAPI(title="{name}", lifespan=lifespan, description='''{desc}''')

"""

    service_routes_code = []
    if completed_app.CompiledRoutes is None:
        raise ValueError("Application must have at least one compiled route.")

    for compiled_route in completed_app.CompiledRoutes:
        if compiled_route.ApiRouteSpec is None:
            raise ValueError(f"Compiled route {compiled_route.id} has no APIRouteSpec")

        if compiled_route.Packages:
            libraries.extend(compiled_route.Packages)

        if compiled_route.RootFunction is None:
            raise ValueError(f"Compiled route {compiled_route.id} has no root function")

        server_code_imports.append(
            f"import project.{compiled_route.fileName.replace('.py', '')}"
        )

        service_routes_code.append(create_server_route_code(compiled_route))

    packages = resolve_package_requirements(libraries)

    # Compile the server code
    server_code = "\n".join(server_code_imports)
    server_code += "\n\n"
    server_code += server_code_header
    server_code += "\n\n"
    server_code += "\n\n".join(service_routes_code)

    db_schema: str = get_database_schema(spec)

    # Update the application with the server code
    formatted_code = await CodeValidator(
        database_schema=db_schema,
        compiled_route_id=f"CompletedAppID-{completed_app.id}",
    ).reformat_code(
        code=server_code,
        packages=[
            PackageModel(
                package_name=package.packageName,
                version=package.version,
                specifier=package.specifier,
            )
            for package in packages
        ],
    )

    return Application(
        name=name,
        description=desc,
        server_code=formatted_code,
        completed_app=completed_app,
        packages=packages,
    )


def resolve_package_requirements(packages: list[Package]) -> list[Package]:
    """
    Resolve duplicates and version conflicts in a list of package requirements.

    - Resolves multiple requirements for the same package to the highest specified version
    - Resolves requirements that match an entry of `DEFAULT_DEPS` to the hardcoded entry
    """
    resolved_packages: dict[str, Package] = {}
    for package in packages:
        name = package.packageName
        if (
            package.id == name
            and name in DEFAULT_DEPS
            or (
                name not in resolved_packages
                or not resolved_packages[name].version
                or package.version
                and version.parse(package.version)
                > version.parse(resolved_packages[name].version)
            )
            and name not in DEFAULT_DEPS
        ):
            resolved_packages[name] = package
    return list(resolved_packages.values())