ishantvivek/codegen

1import asyncio2from fastapi_router_controller import Controller3from fastapi import APIRouter, Depends, HTTPException4from services.ibm_model.ibm_extract_code_block import IbmExtractCodeblock5from services.ibm_model.ibm_text_generator import IbmTextGenerator6from services.model_generator import ModelGenerator7 8from utils.logger import Logger9 10logger = Logger.get_logger(__name__)11 12router = APIRouter(prefix='/v1')13controller = Controller(router, openapi_tag={14    'name': 'Generate the Code Snippets',15})16 17model = ModelGenerator()18 19 20@controller.use()21@controller.resource()22class GenerateController():23    def __init__(24            self,25            service: ModelGenerator = Depends()) -> None:26        self.model_generator = service27 28    @controller.route.get(29        '/generate',30        tags=['generate-code'],31        summary='Generates the code for the given input')32    async def generate_code(self, input: str):33        try:34            if not input:35                logger.error('Input is required.')36                raise HTTPException(37                    status_code=500, detail='Input is required.')38 39            ibm_generate_text_visitor = IbmTextGenerator()40            generated_text = await self.model_generator.acceptTextGenerator(ibm_generate_text_visitor, input)41 42            ibm_extract_code_block_visitor = IbmExtractCodeblock()43            code_block = self.model_generator.acceptExtractCodeBlock(44                ibm_extract_code_block_visitor, generated_text)45 46            logger.info('Output: {}'.format(generated_text))47 48            return {"data": code_block}49        except asyncio.CancelledError:50            logger.error(51                'Canceling network request due to disconnect in client.')52        except Exception as error:53            logger.error('Error {}'.format(error))54