Generate AsyncAPI spec based on the user description

fastkafka_gen/_code_generator/plan_generator.py → fastkafka_gen/_code_generator/asyncapi_spec_generator.py

@@ -1,32 +1,32 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/Plan_Generator.ipynb.
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/AsyncAPI_Spec_Generator.ipynb.
 
 # %% auto 0
 __all__ = ['logger', 'ENTITY_ERROR_MSG', 'APPS_ERROR_MSG', 'CONSUME_FUNCTIONS_ERROR_MSG', 'PRODUCE_FUNCTIONS_ERROR_MSG',
-           'EXPECTED_FUNCTION_KEYS', 'EXPECTED_APP_KEYS', 'generate_plan']
+           'EXPECTED_FUNCTION_KEYS', 'EXPECTED_APP_KEYS', 'generate_asyncapi_spec']
 
-# %% ../../nbs/Plan_Generator.ipynb 1
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 1
 from typing import *
 import time
-import json
+import yaml
 from pathlib import Path
 
 from yaspin import yaspin
 
 from .._components.logger import get_logger
 from .helper import CustomAIChat, ValidateAndFixResponse
-from .prompts import PLAN_GENERATION_PROMPT
+from .prompts import ASYNCAPI_SPEC_GENERATION_PROMPT
 
-# %% ../../nbs/Plan_Generator.ipynb 3
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 3
 logger = get_logger(__name__)
 
-# %% ../../nbs/Plan_Generator.ipynb 5
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 5
 ENTITY_ERROR_MSG = {
     "invalid_entity": "The entities should be a list and cannot be empty in the generated plan. Please read the ==== APP DESCRIPTION: ==== and generate valid entities",
     "invalid_name": "The name of the entity should be defined and cannot be empty. Please read the ==== APP DESCRIPTION: ==== and add a valid value to the 'name' key",
     "invalid_arguments": "The arguments of the entity should be a dictionary with key, value pairs and cannot be empty or any other datatype. Please read the ==== APP DESCRIPTION: ==== and generate valid arguments",
 }
 
-# %% ../../nbs/Plan_Generator.ipynb 6
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 6
 def _validate_entities(plan: Dict[str, List[Dict[str, Any]]]) -> List[str]:
     """Validate the entities in the given plan and returns a list of any error messages encountered.
 
@@ -51,7 +51,7 @@ def _validate_entities(plan: Dict[str, List[Dict[str, Any]]]) -> List[str]:
             errors.append(ENTITY_ERROR_MSG["invalid_arguments"])
     return errors
 
-# %% ../../nbs/Plan_Generator.ipynb 12
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 12
 APPS_ERROR_MSG = {
     "invalid_app": "The apps should be a list and cannot be empty in the generated plan. Please read the ==== APP DESCRIPTION: ==== and generate valid apps",
     "missing_app_keys": "The below keys are missing from the apps. Please read the ==== APP DESCRIPTION: ==== and add the missing keys",
@@ -78,7 +78,7 @@ def _validate_entities(plan: Dict[str, List[Dict[str, Any]]]) -> List[str]:
     "parameters",
 ]
 
-# %% ../../nbs/Plan_Generator.ipynb 13
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 13
 def _validate_for_missing_keys(
     key: str, missing_keys: List[str], errors: List[str], error_msgs: Dict[str, str]
 ) -> List[str]:
@@ -99,7 +99,7 @@ def _validate_for_missing_keys(
 
     return errors
 
-# %% ../../nbs/Plan_Generator.ipynb 16
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 16
 def _validate_prefix(
     key: str,
     params: Dict[str, Union[str, List[Dict[str, str]]]],
@@ -121,7 +121,7 @@ def _validate_prefix(
         errors.append(error_msgs["invalid_prefix"].format(key))
     return errors
 
-# %% ../../nbs/Plan_Generator.ipynb 20
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 20
 def _get_error_msgs_and_expected_keys(
     is_producer_function: bool,
 ) -> Tuple[Dict[str, str], List[str]]:
@@ -138,7 +138,7 @@ def _get_error_msgs_and_expected_keys(
     else:
         return CONSUME_FUNCTIONS_ERROR_MSG, EXPECTED_FUNCTION_KEYS
 
-# %% ../../nbs/Plan_Generator.ipynb 23
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 23
 def _validate_functions(
     functions: Dict[str, Dict[str, Union[str, List[Dict[str, str]]]]],
     errors: List[str],
@@ -174,7 +174,7 @@ def _validate_functions(
                     errors.append(error_msgs["missing_return"].format(key))
     return errors
 
-# %% ../../nbs/Plan_Generator.ipynb 30
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 30
 EXPECTED_APP_KEYS = [
     "app_name",
     "kafka_brokers",
@@ -221,7 +221,7 @@ def _validate_apps(plan: Dict[str, List[Dict[str, Any]]]) -> List[str]:
                 errors = _validate_functions(func_details, errors, flag)
     return errors
 
-# %% ../../nbs/Plan_Generator.ipynb 38
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 38
 def _vaidate_plan(plan: Dict[str, List[Dict[str, Any]]]) -> List[str]:
     """Validates the generated plan
 
@@ -235,60 +235,65 @@ def _vaidate_plan(plan: Dict[str, List[Dict[str, Any]]]) -> List[str]:
     app_error = _validate_apps(plan)
     return entity_error + app_error
 
-# %% ../../nbs/Plan_Generator.ipynb 41
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 41
 def _validate_response(response: str) -> List[str]:
-    """Validate the plan response generated by OpenAI
+    """Validate the AsyncAPI spec generated by OpenAI
 
     Args:
-        response: The JSON plan response generated by OpenAI in string format.
+        response: The AsyncAPI spec generated by OpenAI in string format.
 
     Returns:
-        Returns a list of errors if any found during the validation of the plan.
+        Returns a list of errors if any found during the validation of the spec.
 
     Raises:
         json.JSONDecodeError: If the response is not a valid JSON.
     """
     try:
-        response_dict = json.loads(response)
-        errors_list = _vaidate_plan(response_dict)
-        return errors_list
-    except json.JSONDecodeError as e:
-        return ["JSON decoding failed. Please send JSON response only."]
+        response_dict = yaml.safe_load(response)
+        # TODO: validate the spec using a compiler
+        # errors_list = _vaidate_plan(response_dict)
+        return []
+    except yaml.YAMLError as exc:
+        return ["Failed to parse AsyncAPI spec. Please generate a valid YAML response only."]
 
-# %% ../../nbs/Plan_Generator.ipynb 44
-def _save_plan_as_json(contents: str, output_file: str) -> None:
-    """Save the contents as JSON to the specified output file.
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 44
+def _save_async_api_spec(contents: str, output_path: str) -> str:
+    """Save the YAML-formatted asyncapi spec in the specified output path.
 
     Args:
-        contents: A JSON-formatted string to be saved as a JSON file.
-        output_file: The path to the output file where the JSON content will be saved.
+        contents: A YAML-formatted asyncapi spec.
+        output_file: The path to save the asyncapi spec.
     """
-    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
+    Path(output_path).mkdir(parents=True, exist_ok=True)
+
+    output_file = f"{output_path}/asyncapi.yml"
     with open(output_file, "w", encoding="utf-8") as f:
-        json.dump(json.loads(contents), f, indent=4)
+        f.write(contents)
+
+    return output_file
 
-# %% ../../nbs/Plan_Generator.ipynb 46
-def generate_plan(description: str, output_file: str) -> str:
-    """Generate a plan from user's application description
+# %% ../../nbs/AsyncAPI_Spec_Generator.ipynb 46
+def generate_asyncapi_spec(description: str, output_path: str) -> str:
+    """Generate a AsyncAPI spec from the user's application description
 
     Args:
         description: Validated User application description
-        output_path: The path to the output file where the generated plan JSON will be saved.
+        output_path: The path to the output file where the generated AsyncAPI spec will be saved.
 
     Returns:
-        The total token used to generate the plan
+        The total token used to generate the AsyncAPI spec
     """
     with yaspin(
-        text="Generating plan", #  (slowest step, usually takes 30 to 90 seconds)...
+        text="Generating plan (usually takes 30 to 60 seconds)...",
         color="cyan",
         spinner="clock",
     ) as sp:
-        plan_generator = CustomAIChat(user_prompt=PLAN_GENERATION_PROMPT)
-        plan_validator = ValidateAndFixResponse(plan_generator, _validate_response)
-        validated_plan, total_tokens = plan_validator.fix(description)
+        async_spec_generator = CustomAIChat(user_prompt=ASYNCAPI_SPEC_GENERATION_PROMPT)
+        async_spec_validator = ValidateAndFixResponse(async_spec_generator, _validate_response)
+        validated_async_spec, total_tokens = async_spec_validator.fix(description)
 
-        _save_plan_as_json(validated_plan, output_file)
+        output_file = _save_async_api_spec(validated_async_spec, output_path)
 
         sp.text = ""
-        sp.ok(f" ✔ Plan generated and saved at: {output_file}")
+        sp.ok(f" ✔ AsyncAPI specification generated and saved to: {output_file}")
         return total_tokens