Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(Pipelines) : Smart Data Frame Pipeline (Sourcery refactored) #740

Merged
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 5 additions & 9 deletions pandasai/smart_datalake/prompt_generation.py
Original file line number Diff line number Diff line change
@@ -37,13 +37,9 @@ def execute(self, input: Any, **kwargs) -> Any:
"last_code_generated"
)

generate_python_code_instruction = (
pipeline_context.query_exec_tracker.execute_func(
pipeline_context.get_intermediate_value("get_prompt"),
"generate_python_code",
default_prompt=GeneratePythonCodePrompt,
default_values=default_values,
)
return pipeline_context.query_exec_tracker.execute_func(
pipeline_context.get_intermediate_value("get_prompt"),
"generate_python_code",
default_prompt=GeneratePythonCodePrompt,
default_values=default_values,
)

return generate_python_code_instruction
Comment on lines -40 to -49
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function PromptGeneration.execute refactored with the following changes:

8 changes: 4 additions & 4 deletions tests/pipelines/smart_datalake/test_code_execution.py
Original file line number Diff line number Diff line change
@@ -74,8 +74,7 @@ def config(self, llm):

@pytest.fixture
def context(self, sample_df, config):
pipeline_context = PipelineContext([sample_df], config)
return pipeline_context
return PipelineContext([sample_df], config)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function TestCodeExecution.context refactored with the following changes:


@pytest.fixture
def logger(self):
@@ -127,15 +126,16 @@ def mock_intermediate_values(key : str):
return SkillsManager()
elif key == "code_manager":
return mock_code_manager

context.get_intermediate_value = Mock(side_effect=mock_intermediate_values)

assert isinstance(code_execution, CodeExecution)

result = None
try:
result = code_execution.execute(input="Test Code", context=context, logger=logger)
except Exception as e:
assert result == None
assert result is None
Comment on lines +129 to +138
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function TestCodeExecution.test_code_execution_unsuccessful_after_retries refactored with the following changes:


def test_code_execution_successful_at_retry(self, context, logger):
# Test Flow : Code Execution Successful with no exceptions
3 changes: 1 addition & 2 deletions tests/pipelines/smart_datalake/test_code_generator.py
Original file line number Diff line number Diff line change
@@ -74,8 +74,7 @@ def config(self, llm):

@pytest.fixture
def context(self, sample_df, config):
pipeline_context = PipelineContext([sample_df], config)
return pipeline_context
return PipelineContext([sample_df], config)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function TestCodeGenerator.context refactored with the following changes:


@pytest.fixture
def logger(self):
9 changes: 4 additions & 5 deletions tests/pipelines/smart_datalake/test_result_parsing.py
Original file line number Diff line number Diff line change
@@ -72,8 +72,7 @@ def config(self, llm):

@pytest.fixture
def context(self, sample_df, config):
pipeline_context = PipelineContext([sample_df], config)
return pipeline_context
return PipelineContext([sample_df], config)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function TestResultParsing.context refactored with the following changes:


@pytest.fixture
def logger(self):
@@ -111,20 +110,20 @@ def test_result_parsing_unsuccessful_with_exceptions(self, context, logger):

def mock_result_parsing(*args, **kwargs):
raise Exception("Unit test exception")

context._query_exec_tracker = Mock()
context.query_exec_tracker.execute_func = Mock(side_effect=mock_result_parsing)

def mock_intermediate_values(key : str):
if key == "response_parser" :
return mock_response_parser

context.get_intermediate_value = Mock(side_effect=mock_intermediate_values)

result = None
try:
result = result_parsing.execute(input="Test Result", context=context, logger=logger)
except Exception as e:
assert result == None
assert result is None
Comment on lines -114 to +127
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function TestResultParsing.test_result_parsing_unsuccessful_with_exceptions refactored with the following changes:

assert isinstance(result_parsing, ResultParsing)

5 changes: 2 additions & 3 deletions tests/pipelines/smart_datalake/test_result_validation.py
Original file line number Diff line number Diff line change
@@ -72,8 +72,7 @@ def config(self, llm):

@pytest.fixture
def context(self, sample_df, config):
pipeline_context = PipelineContext([sample_df], config)
return pipeline_context
return PipelineContext([sample_df], config)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function TestResultValidation.context refactored with the following changes:


@pytest.fixture
def logger(self):
@@ -96,7 +95,7 @@ def test_result_is_none(self, context, logger):

assert not context.query_exec_tracker.add_step.called
assert isinstance(result_validation, ResultValidation)
assert result == None
assert result is None
Comment on lines -99 to +98
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function TestResultValidation.test_result_is_none refactored with the following changes:


def test_result_is_not_of_dict_type(self, context, logger):
# Test Flow : Code Execution Successful with no exceptions