diff --git a/pandasai/prompts/__init__.py b/pandasai/prompts/__init__.py index 28bede48c..4e06b108f 100644 --- a/pandasai/prompts/__init__.py +++ b/pandasai/prompts/__init__.py @@ -1,4 +1,5 @@ -from .base import AbstractPrompt, FileBasedPrompt +from .base import AbstractPrompt +from .file_based_prompt import FileBasedPrompt from .correct_error_prompt import CorrectErrorPrompt from .generate_python_code import GeneratePythonCodePrompt diff --git a/pandasai/prompts/base.py b/pandasai/prompts/base.py index 29c30498c..93c8d1986 100644 --- a/pandasai/prompts/base.py +++ b/pandasai/prompts/base.py @@ -2,10 +2,6 @@ In order to better handle the instructions, this prompt module is written. """ from abc import ABC, abstractmethod -import os -from pathlib import Path - -from ..exceptions import TemplateFileNotFoundError class AbstractPrompt(ABC): @@ -78,37 +74,3 @@ def __str__(self): def validate(self, output: str) -> bool: return isinstance(output, str) - - -class FileBasedPrompt(AbstractPrompt): - """Base class for prompts supposed to read template content from a file. - - `_path_to_template` attribute has to be specified. - """ - - _path_to_template: str - - def __init__(self, **kwargs): - if (template_path := kwargs.pop("path_to_template", None)) is not None: - self._path_to_template = template_path - else: - current_dir_path = Path(__file__).parent - self._path_to_template = os.path.join( - current_dir_path, "..", self._path_to_template - ) - - super().__init__(**kwargs) - - @property - def template(self) -> str: - try: - with open(self._path_to_template) as fp: - return fp.read() - except FileNotFoundError: - raise TemplateFileNotFoundError( - self._path_to_template, self.__class__.__name__ - ) - except IOError as exc: - raise RuntimeError( - f"Failed to read template file '{self._path_to_template}': {exc}" - ) diff --git a/pandasai/prompts/clarification_questions_prompt.py b/pandasai/prompts/clarification_questions_prompt.py index 8299d4ef3..19b43b02d 100644 --- a/pandasai/prompts/clarification_questions_prompt.py +++ b/pandasai/prompts/clarification_questions_prompt.py @@ -18,7 +18,7 @@ import json from typing import List import pandas as pd -from .base import FileBasedPrompt +from .file_based_prompt import FileBasedPrompt class ClarificationQuestionPrompt(FileBasedPrompt): diff --git a/pandasai/prompts/correct_error_prompt.py b/pandasai/prompts/correct_error_prompt.py index 20969b6bb..ffdbb57c0 100644 --- a/pandasai/prompts/correct_error_prompt.py +++ b/pandasai/prompts/correct_error_prompt.py @@ -16,7 +16,7 @@ Correct the python code and return a new python code that fixes the above mentioned error. Do not generate the same code again. """ # noqa: E501 -from .base import FileBasedPrompt +from .file_based_prompt import FileBasedPrompt class CorrectErrorPrompt(FileBasedPrompt): diff --git a/pandasai/prompts/explain_prompt.py b/pandasai/prompts/explain_prompt.py index 733151564..b0c3be027 100644 --- a/pandasai/prompts/explain_prompt.py +++ b/pandasai/prompts/explain_prompt.py @@ -15,7 +15,7 @@ mentioning technical details or mentioning the libraries used? """ -from .base import FileBasedPrompt +from .file_based_prompt import FileBasedPrompt class ExplainPrompt(FileBasedPrompt): diff --git a/pandasai/prompts/file_based_prompt.py b/pandasai/prompts/file_based_prompt.py new file mode 100644 index 000000000..80f48a0b5 --- /dev/null +++ b/pandasai/prompts/file_based_prompt.py @@ -0,0 +1,39 @@ +import os +from pathlib import Path + +from ..exceptions import TemplateFileNotFoundError +from .base import AbstractPrompt + + +class FileBasedPrompt(AbstractPrompt): + """Base class for prompts supposed to read template content from a file. + + `_path_to_template` attribute has to be specified. + """ + + _path_to_template: str + + def __init__(self, **kwargs): + if (template_path := kwargs.pop("path_to_template", None)) is not None: + self._path_to_template = template_path + else: + current_dir_path = Path(__file__).parent + self._path_to_template = os.path.join( + current_dir_path, "..", self._path_to_template + ) + + super().__init__(**kwargs) + + @property + def template(self) -> str: + try: + with open(self._path_to_template) as fp: + return fp.read() + except FileNotFoundError: + raise TemplateFileNotFoundError( + self._path_to_template, self.__class__.__name__ + ) + except IOError as exc: + raise RuntimeError( + f"Failed to read template file '{self._path_to_template}': {exc}" + ) diff --git a/pandasai/prompts/generate_python_code.py b/pandasai/prompts/generate_python_code.py index a262c585d..24b16cd7d 100644 --- a/pandasai/prompts/generate_python_code.py +++ b/pandasai/prompts/generate_python_code.py @@ -33,7 +33,7 @@ def analyze_data(dfs: list[{engine_df_name}]) -> dict: """ # noqa: E501 -from .base import FileBasedPrompt +from .file_based_prompt import FileBasedPrompt class GeneratePythonCodePrompt(FileBasedPrompt): diff --git a/pandasai/prompts/rephase_query_prompt.py b/pandasai/prompts/rephase_query_prompt.py index bad84ed8c..c5f953af4 100644 --- a/pandasai/prompts/rephase_query_prompt.py +++ b/pandasai/prompts/rephase_query_prompt.py @@ -11,7 +11,7 @@ from typing import List import pandas as pd -from .base import FileBasedPrompt +from .file_based_prompt import FileBasedPrompt class RephraseQueryPrompt(FileBasedPrompt):