Custom Prompts

[LiquidGunay]

System Info

OS: Ubuntu Python: 3.10 pandasai: v2.0.26

🐛 Describe the bug

from pandasai.prompts.base import BasePrompt
from pandasai.pipelines.chat.prompt_generation import PromptGeneration
from pandasai.pipelines.chat.generate_chat_pipeline import GenerateChatPipeline

class CustomPrompt(BasePrompt):
    template_path = "custom_template.tmpl"

    def __init__(self, **kwargs):
        self.props = kwargs

        # assumes the template file is in a templates subfolder
        path_to_templates = os.path.join(Path(__file__).parent, "templates")
        env = Environment(loader=FileSystemLoader(path_to_templates))
        self.prompt = env.get_template(self.template_path)

        self._resolved_prompt = None

class CustomPromptGeneration(PromptGeneration):
    def get_chat_prompt(self, context):
        viz_lib = "matplotlib"
        if context.config.data_viz_library:
            viz_lib = context.config.data_viz_library

        return CustomPrompt(
            context=context,
            last_code_generated=context.get("last_code_generated"),
            viz_lib=viz_lib,
            output_type=context.get("output_type"),
        )

class CustomPipeline(GenerateChatPipeline):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.pipeline._steps = [
            (
                CustomPromptGeneration(
                    skip_if=step.skip_if, on_execution=step.on_execution
                )
                if isinstance(step, PromptGeneration)
                else step
            )
            for step in self.pipeline._steps
        ]

custom_agent = Agent(dfs, config={"llm": llm}, pipeline=CustomPipeline)

I'm running this code snippet that I found in another issue and I'm getting a lot of errors with initialising my Custom Pipeline. Custom Pipeline was expecting 1 arg but was getting 3 which I fixed by modifying the initialisation of the pipeline in agent's base.py. Now I am getting errors like Pipeline object has no attribute steps. Is there a better way to use custom prompts?

[LiquidGunay]

Okay so I think self.pipeline has to be replaced with self.code_generation_pipeline (stuff has been refactored). But I think the bug in the agents/base.py should still be fixed