ronantakizawa/codereview-bench huggingface dataset:
Repository: googleapis/genai-toolbox
File: internal/util/parameters/parameters_test.go (Go)

Code before:
			want:     false,
		},
		{
			name:     "not required and no default",
			required: false,
			defaultV: nil,
			want:     false,
		},
		{
			name:     "not required and default",
			required: false,
			defaultV: "foo",
			want:     false,
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			got := parameters.CheckParamRequired(tc.required, tc.defaultV)
			if got != tc.want {
				t.Fatalf("got %v, want %v", got, tc.want)
			}
		})
	}
}

func TestParseParams_ExplicitNullForRequiredParam(t *testing.T) {
	// Define a required string parameter
	params := parameters.Parameters{
		parameters.NewStringParameter("required_param", "this is required"),
	}

	// Input map with explicit nil
	input := map[string]any{
		"required_param": nil,
	}

	// Call ParseParams
	_, err := parameters.ParseParams(params, input, nil)

	// Expect an error because the parameter is required
	if err == nil {
		t.Errorf("ParseParams allowed explicit nil for required parameter, expected error")
	}
}


Review comment (bug):
```suggestion
func TestParametersParseFail(t *testing.T) {
```

Can we move this function to right under the `TestParametersParse`?

In this function, we can do 
```
t.Run("CheckNullForRequiredParam", func(t *testing.T) {
  // code here.
}
```

Code after:
}

func TestCheckParamRequired(t *testing.T) {
	tcs := []struct {
		name     string
		required bool
		defaultV any
		want     bool
	}{
		{
			name:     "required and no default",
			required: true,
			defaultV: nil,
			want:     true,
		},
		{
			name:     "required and default",
			required: true,
			defaultV: "foo",
			want:     false,
		},
		{
			name:     "not required and no default",
			required: false,
			defaultV: nil,
			want:     false,
		},
		{
			name:     "not required and default",
			required: false,
			defaultV: "foo",
			want:     false,
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			got := parameters.CheckParamRequired(tc.required, tc.defaultV)
			if got != tc.want {
				t.Fatalf("got %v, want %v", got, tc.want)
			}
		})
	}
}
Repository: gventuri/pandas-ai
File: examples/from_snowflake.py (Python)

Code before:
"""Example of using PandasAI with a Snowflake"""

from pandasai import SmartDatalake
from pandasai.llm import OpenAI
from pandasai.connectors import SnowFlakeConnector


snowflake_connector = SnowFlakeConnector(config={
        "account": "ehxzojy-ue47135", 
        "database": "SNOWFLAKE_SAMPLE_DATA",
        "username": "test",
        "password": "*****",
        "table": "lineitem",
        "warehouse" : "COMPUTE_WH",
        "dbSchema": "tpch_sf1",
        "where": [
            # this is optional and filters the data to
            # reduce the size of the dataframe
            ["l_quantity", ">","49"]
        ],
    })

OPEN_AI_API = "Your-API-Key"
llm = OpenAI(api_token=OPEN_AI_API)
df = SmartDatalake([snowflake_connector], config={"llm": llm})

response = df.chat("Count line status is F")
print(response)


Review comment (suggestion):
Guess since it's a single connector, we can use the `SmartDatframe` here instead, like this:

```
df = SmartDatframe(snowflake_connector, config={"llm": llm})
```

Code after:
"""Example of using PandasAI with a Snowflake"""

from pandasai import SmartDataframe
from pandasai.llm import OpenAI
from pandasai.connectors import SnowFlakeConnector


snowflake_connector = SnowFlakeConnector(
    config={
        "account": "ehxzojy-ue47135",
        "database": "SNOWFLAKE_SAMPLE_DATA",
        "username": "test",
        "password": "*****",
        "table": "lineitem",
        "warehouse": "COMPUTE_WH",
        "dbSchema": "tpch_sf1",
        "where": [
            # this is optional and filters the data to
            # reduce the size of the dataframe
            ["l_quantity", ">", "49"]
        ],
    }
)

OPEN_AI_API = "Your-API-Key"
llm = OpenAI(api_token=OPEN_AI_API)
df = SmartDataframe(snowflake_connector, config={"llm": llm})

response = df.chat("Count line status is F")
print(response)
Repository: gventuri/pandas-ai
File: pandasai/agent/__init__.py (Python)

Code before:
import json
from typing import Union, List, Optional
from pandasai.agent.response import ClarificationResponse
from pandasai.helpers.df_info import DataFrameType
from pandasai.helpers.logger import Logger
from pandasai.helpers.memory import Memory
from pandasai.prompts.base import Prompt
from pandasai.prompts.clarification_questions_prompt import ClarificationQuestionPrompt
from pandasai.prompts.explain_prompt import ExplainPrompt
from pandasai.schemas.df_config import Config

from pandasai.smart_datalake import SmartDatalake


class Agent:
    """
    Agent class to improve the conversational experience in PandasAI
    """

    _memory: Memory
    _lake: SmartDatalake = None
    logger: Optional[Logger] = None

    def __init__(
        self,
        dfs: Union[DataFrameType, List[DataFrameType]],
        config: Optional[Union[Config, dict]] = None,
        logger: Optional[Logger] = None,
        memory_size: int = 1,
    ):
        """
        Args:
            df (Union[SmartDataframe, SmartDatalake]): _description_
            memory_size (int, optional): _description_. Defaults to 1.
        """

        if not isinstance(dfs, list):
            dfs = [dfs]

        self._lake = SmartDatalake(dfs, config, logger)
        self.logger = self._lake.logger
        # For the conversation multiple the memory size by 2
        self._memory = Memory(memory_size * 2)

    def _get_conversation(self):

Review comment (suggestion):
Here we should rely on the memory from the `SmartDatalake` instead, since we have access to it.

Code after:
import json
from typing import Union, List, Optional
from pandasai.helpers.df_info import DataFrameType
from pandasai.helpers.logger import Logger
from pandasai.helpers.memory import Memory
from pandasai.prompts.clarification_questions_prompt import ClarificationQuestionPrompt
from pandasai.prompts.explain_prompt import ExplainPrompt
from pandasai.schemas.df_config import Config
from pandasai.smart_datalake import SmartDatalake


class Agent:
    """
    Agent class to improve the conversational experience in PandasAI
    """

    _lake: SmartDatalake = None
    _logger: Optional[Logger] = None

    def __init__(
        self,
        dfs: Union[DataFrameType, List[DataFrameType]],
        config: Optional[Union[Config, dict]] = None,
        logger: Optional[Logger] = None,
        memory_size: int = 1,
    ):
        """
        Args:
            df (Union[DataFrameType, List[DataFrameType]]): DataFrame can be Pandas,
            Polars or Database connectors
            memory_size (int, optional): Conversation history to use during chat.
            Defaults to 1.
        """

        if not isinstance(dfs, list):
            dfs = [dfs]

        self._lake = SmartDatalake(dfs, config, logger, memory=Memory(memory_size))
        self._logger = self._lake.logger

    def chat(self, query: str, output_type: Optional[str] = None):
        """
        Simulate a chat interaction with the assistant on Dataframe.
        """
        try:
Repository: gventuri/pandas-ai
File: pandasai/agent/__init__.py (Python)

Code before:
                "because of the following error:\n"
                f"\n{exception}\n"
            )

    def _get_clarification_prompt(self) -> Prompt:
        """
        Create a clarification prompt with relevant variables.
        """
        prompt = ClarificationQuestionPrompt()
        prompt.set_var("dfs", self._lake.dfs)
        prompt.set_var("conversation", self._get_conversation())
        return prompt

    def clarification_questions(self) -> ClarificationResponse:
        """
        Generate clarification questions based on the data
        """
        try:
            prompt = self._get_clarification_prompt()
            result = self._lake.llm.call(prompt)
            self.logger.log(
                f"""Clarification Questions:  {result}
                """
            )
            questions: list[str] = json.loads(result)
            return ClarificationResponse(
                success=True, questions=questions[:3], message=result
            )
        except Exception as exception:
            return ClarificationResponse(
                False,
                [],
                "Unfortunately, I was not able to get your clarification questions, "
                "because of the following error:\n"
                f"\n{exception}\n",
            )

    def start_new_conversation(self):
        """
        Clears the previous conversation
        """

        self._memory.clear()

    def explain(self) -> str:
        """
        Returns the explanation of the code how it reached to the solution
        """
        try:
            prompt = ExplainPrompt()
            prompt.set_var("code", self._lake.last_code_executed)

Review comment (refactor):
We are not doing a lot of OOP at the moment. Let's simplify this:
- if it's successful, we only return a list of questions
- otherwise we raise an exception

Code after:
        Clears the previous conversation
        """
        self._lake._memory.clear()

    def explain(self) -> str:
        """
        Returns the explanation of the code how it reached to the solution
        """
        try:
            prompt = ExplainPrompt(
                self._lake._memory.get_conversation(),
                self._lake.last_code_executed,
            )
            response = self._lake.llm.call(prompt)
            self._logger.log(
                f"""Explaination:  {response}
                """
            )
            return response
        except Exception as exception:
            return (
                "Unfortunately, I was not able to explain, "
                "because of the following error:\n"
                f"\n{exception}\n"
            )
Repository: gventuri/pandas-ai
File: pandasai/helpers/memory.py (Python)

Code before:
""" Memory class to store the conversations """
import sys


class Memory:
    """Memory class to store the conversations"""

    _messages: list
    _max_messages: int

    def __init__(self, max_messages: int = sys.maxsize):
        self._messages = []
        self._max_messages = max_messages

    def add(self, message: str, is_user: bool):
        self._messages.append({"message": message, "is_user": is_user})

        # Delete two entry because of the conversation
        if len(self._messages) > self._max_messages:
            del self._messages[:2]

    def count(self) -> int:
        return len(self._messages)

    def all(self) -> list:
        return self._messages

    def last(self) -> dict:
        return self._messages[-1]

    def get_conversation(self, limit: int = 1) -> str:
        return "\n".join(
            [
                f"{f'User {i+1}' if message['is_user'] else f'Assistant {i}'}: "
                f"{message['message']}"
                for i, message in enumerate(self._messages[-limit:])
            ]
        )

    def clear(self):
        self._messages = []


Review comment (suggestion):
No need for this, the `get_conversation` method of memory will automatically limit based on the `limit` variable that is passed.

Code after:
""" Memory class to store the conversations """


class Memory:
    """Memory class to store the conversations"""

    _messages: list
    _memory_size: int

    def __init__(self, memory_size: int = 1):
        self._messages = []
        self._memory_size = memory_size

    def add(self, message: str, is_user: bool):
        self._messages.append({"message": message, "is_user": is_user})

    def count(self) -> int:
        return len(self._messages)

    def all(self) -> list:
        return self._messages

    def last(self) -> dict:
        return self._messages[-1]

    def get_conversation(self, limit: int = None) -> str:
        """
        Returns the conversation messages based on limit parameter
        or default memory size
        """
        limit = self._memory_size if limit is None else limit
        return "\n".join(
            [
                f"{f'User {i+1}' if message['is_user'] else f'Assistant {i}'}: "
                f"{message['message']}"
                for i, message in enumerate(self._messages[-limit:])
            ]
        )

    def clear(self):
        self._messages = []
Repository: gventuri/pandas-ai
File: pandasai/prompts/clarification_questions_prompt.py (Python)

Code before:
</dataframe>

<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions that a senior data scientist would ask? These are questions for non technical people, only ask for questions they could ask given low tech expertise and no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. If there is no clarification question, return an empty array.

Json:
"""  # noqa: E501


from .base import Prompt


class ClarificationQuestionPrompt(Prompt):
    """Prompt to get clarification questions"""

    text: str = """
You are provided with the following pandas DataFrames:

<dataframe>
{dataframes}
</dataframe>

<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions 
that a senior data scientist would ask? These are questions for non technical people, 
only ask for questions they could ask given low tech expertise and 
no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. 

If there is no clarification question, return an empty array.

Json:
"""

    def __init__(self, dataframes, conversation):
        self.set_var("dataframes", dataframes)
        self.set_var("conversation", conversation)


Review comment (suggestion):
`{dataframes}` only should be enough, the `set_var` method in prompts already takes care of wrapping each one individually.

Code after:
<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions that a senior data scientist would ask? These are questions for non technical people, only ask for questions they could ask given low tech expertise and no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. If there is no clarification question, return an empty array.

Json:
"""  # noqa: E501


from typing import List
import pandas as pd
from .base import Prompt


class ClarificationQuestionPrompt(Prompt):
    """Prompt to get clarification questions"""

    text: str = """
You are provided with the following pandas DataFrames:

{dataframes}

<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions 
that a senior data scientist would ask? These are questions for non technical people, 
only ask for questions they could ask given low tech expertise and 
no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. 

If there is no clarification question, return an empty array.

Json:
"""

    def __init__(self, dataframes: List[pd.DataFrame], conversation: str):
        self.set_var("dfs", dataframes)
        self.set_var("conversation", conversation)
Repository: gventuri/pandas-ai
File: pandasai/prompts/clarification_questions_prompt.py (Python)

Code before:

    text: str = """
You are provided with the following pandas DataFrames:

<dataframe>
{dataframes}
</dataframe>

<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions 
that a senior data scientist would ask? These are questions for non technical people, 
only ask for questions they could ask given low tech expertise and 
no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. 

If there is no clarification question, return an empty array.

Json:
"""

    def __init__(self, dataframes, conversation):
        self.set_var("dataframes", dataframes)
        self.set_var("conversation", conversation)


Review comment (suggestion):
```
def set_var(self, var, value):
        if var == "dfs":
            self._args["dataframes"] = self._generate_dataframes(value)
        self._args[var] = value
```

This method from `Prompt` automatically configures the dfs to be used in the prompt, but it requires a `dfs` key.
Let's change this to ` self.set_var("dfs", dataframes)`

Code after:

    text: str = """
You are provided with the following pandas DataFrames:

{dataframes}

<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions 
that a senior data scientist would ask? These are questions for non technical people, 
only ask for questions they could ask given low tech expertise and 
no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. 

If there is no clarification question, return an empty array.

Json:
"""

    def __init__(self, dataframes: List[pd.DataFrame], conversation: str):
        self.set_var("dfs", dataframes)
        self.set_var("conversation", conversation)
Repository: gventuri/pandas-ai
File: pandasai/prompts/clarification_questions_prompt.py (Python)

Code before:
""" Prompt to get clarification questions
You are provided with the following pandas DataFrames:

<dataframe>
{dataframe}
</dataframe>

<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions that a senior data scientist would ask? These are questions for non technical people, only ask for questions they could ask given low tech expertise and no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. If there is no clarification question, return an empty array.

Json:
"""  # noqa: E501


from .base import Prompt


class ClarificationQuestionPrompt(Prompt):
    """Prompt to get clarification questions"""

    text: str = """
You are provided with the following pandas DataFrames:

{dataframes}

<conversation>

Review comment (suggestion):
Let's also change this to

```
{dataframes}
```

for consistency!

Code after:
""" Prompt to get clarification questions
You are provided with the following pandas DataFrames:

{dataframes}

<conversation>
{conversation}
</conversation>

Based on the conversation, are there any clarification questions that a senior data scientist would ask? These are questions for non technical people, only ask for questions they could ask given low tech expertise and no knowledge about how the dataframes are structured.

Return the JSON array of the clarification questions. If there is no clarification question, return an empty array.

Json:
"""  # noqa: E501


from typing import List
import pandas as pd
from .base import Prompt


class ClarificationQuestionPrompt(Prompt):
    """Prompt to get clarification questions"""

    text: str = """
You are provided with the following pandas DataFrames:

{dataframes}

<conversation>
Repository: gventuri/pandas-ai
File: pandasai/llm/google_vertexai.py (Python)

Code before:

        from vertexai.preview.language_models import (
            CodeGenerationModel,
            TextGenerationModel,
        )

        if self.model in self._supported_code_models:
            code_generation = CodeGenerationModel.from_pretrained(self.model)

            completion = code_generation.predict(
                prefix=prompt,
                temperature=self.temperature,
                max_output_tokens=self.max_output_tokens,
            )
        elif self.model in self._supported_text_models:
            text_generation = TextGenerationModel.from_pretrained(self.model)

            completion = text_generation.predict(
                prompt=prompt,
                temperature=self.temperature,
                top_p=self.top_p,
                top_k=self.top_k,
                max_output_tokens=self.max_output_tokens,
            )
        else:
            raise UnsupportedOpenAIModelError("Unsupported model")

        return str(completion)

    @property
    def type(self) -> str:
        return "google-vertexai"


Review comment (suggestion):
Shouldn't we generalize this exception - e.g. `UnsupportedModelError`? Or just throwing a `ValueError` would suffice, I suppose.

Code after:

        from vertexai.preview.language_models import (
            CodeGenerationModel,
            TextGenerationModel,
        )

        if self.model in self._supported_code_models:
            code_generation = CodeGenerationModel.from_pretrained(self.model)

            completion = code_generation.predict(
                prefix=prompt,
                temperature=self.temperature,
                max_output_tokens=self.max_output_tokens,
            )
        elif self.model in self._supported_text_models:
            text_generation = TextGenerationModel.from_pretrained(self.model)

            completion = text_generation.predict(
                prompt=prompt,
                temperature=self.temperature,
                top_p=self.top_p,
                top_k=self.top_k,
                max_output_tokens=self.max_output_tokens,
            )
        else:
            raise UnsupportedModelError("Unsupported model")

        return str(completion)

    @property
    def type(self) -> str:
        return "google-vertexai"
Repository: gventuri/pandas-ai
File: pandasai/llm/google_vertexai.py (Python)

Code before:

        from vertexai.preview.language_models import (
            CodeGenerationModel,
            TextGenerationModel,
        )

        if self.model in self._supported_code_models:
            code_generation = CodeGenerationModel.from_pretrained(self.model)

            completion = code_generation.predict(
                prefix=prompt,
                temperature=self.temperature,
                max_output_tokens=self.max_output_tokens,
            )
        elif self.model in self._supported_text_models:
            text_generation = TextGenerationModel.from_pretrained(self.model)

            completion = text_generation.predict(
                prompt=prompt,
                temperature=self.temperature,
                top_p=self.top_p,
                top_k=self.top_k,
                max_output_tokens=self.max_output_tokens,
            )
        else:
            raise UnsupportedModelError("Unsupported model")

        return str(completion)

    @property
    def type(self) -> str:
        return "google-vertexai"


Review comment (suggestion):
@arian81 also please consider using the exactly name of a class instead of it's representation, make it to be something like this:
```python
raise UnsupportedModelError(self.model.__name__)
```
Just because we've got already the same like [here](https://github.com/gventuri/pandas-ai/blob/720811f9ebd5632a20aeff01ea7b73f125ba3a48/pandasai/prompts/file_based_prompt.py#L34).

Of course, this isn't important at all, it's on your choice, both of the ways look good as for me.

Code after:

        from vertexai.preview.language_models import (
            CodeGenerationModel,
            TextGenerationModel,
        )

        if self.model in self._supported_code_models:
            code_generation = CodeGenerationModel.from_pretrained(self.model)

            completion = code_generation.predict(
                prefix=prompt,
                temperature=self.temperature,
                max_output_tokens=self.max_output_tokens,
            )
        elif self.model in self._supported_text_models:
            text_generation = TextGenerationModel.from_pretrained(self.model)

            completion = text_generation.predict(
                prompt=prompt,
                temperature=self.temperature,
                top_p=self.top_p,
                top_k=self.top_k,
                max_output_tokens=self.max_output_tokens,
            )
        else:
            raise UnsupportedModelError(self.model)


        return str(completion)

    @property
    def type(self) -> str:
        return "google-vertexai"