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.

Sign up for GitHub

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

fix(QueryTracker): publish query tracker results #773

Merged
merged 3 commits into from
Nov 22, 2023
Merged

fix(QueryTracker): publish query tracker results #773

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e49ce0e
fix(QueryTracker): publish query tracker results
ArslanSaleem Nov 22, 2023
0d929b5
fix(QueryTracker): execute code is not being tracked
ArslanSaleem Nov 22, 2023
726cfaf
chore(tests): change comment
ArslanSaleem Nov 22, 2023
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
8 changes: 5 additions & 3 deletions pandasai/pipelines/smart_datalake_chat/code_execution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,11 @@ def execute(self, input: Any, **kwargs) -> Any:
pipeline_context.get_intermediate_value("last_prompt_id"),
pipeline_context.get_intermediate_value("skills"),
)
result = pipeline_context.get_intermediate_value(
"code_manager"
).execute_code(

result = pipeline_context.query_exec_tracker.execute_func(
pipeline_context.get_intermediate_value(
"code_manager"
).execute_code,
code=code_to_run,
context=code_context,
)
Expand Down
3 changes: 3 additions & 0 deletions pandasai/smart_datalake/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,9 @@ def chat(self, query: str, output_type: Optional[str] = None):

self.update_intermediate_value_post_pipeline_execution(pipeline_context)

# publish query tracker
self._query_exec_tracker.publish()

return result

def _validate_output(self, result: dict, output_type: Optional[str] = None):
Expand Down
25 changes: 17 additions & 8 deletions tests/pipelines/smart_datalake/test_code_execution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,8 @@ def mock_intermediate_values(key: str):
return mock_code_manager

context.get_intermediate_value = Mock(side_effect=mock_intermediate_values)
context._query_exec_tracker = Mock()
context.query_exec_tracker.execute_func = Mock(return_value="Mocked Result")

result = code_execution.execute(
input="Test Code", context=context, logger=logger
Expand Down Expand Up @@ -157,17 +159,24 @@ def mock_execute_code(*args, **kwargs):
raise Exception("Unit test exception")
return "Mocked Result after retry"

# Conditional return of execute_func method based arguments it is called with
def mock_execute_func(*args, **kwargs):
if isinstance(args[0], Mock) and args[0].name == "execute_code":
return mock_execute_code(*args, **kwargs)
else:
return [
"Interuppted Code",
"Exception Testing",
"Successful after Retry",
]

mock_code_manager = Mock()
mock_code_manager.execute_code = Mock(side_effect=mock_execute_code)
mock_code_manager.execute_code = Mock()
mock_code_manager.execute_code.name = "execute_code"

context._query_exec_tracker = Mock()
context.query_exec_tracker.execute_func = Mock(
return_value=[
"Interuppted Code",
"Exception Testing",
"Successful after Retry",
]
)

context.query_exec_tracker.execute_func = Mock(side_effect=mock_execute_func)

def mock_intermediate_values(key: str):
Comment on lines 159 to 181
Copy link
Contributor

Choose a reason for hiding this comment

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

The mock_execute_func is designed to simulate different behaviors based on the arguments it receives. However, the conditional logic within mock_execute_func seems to be based on the name attribute of the Mock object, which is not a standard attribute. This could lead to unexpected behavior or test failures if the name attribute is not set elsewhere in the code. Consider using a more reliable condition to differentiate the behavior of mock_execute_func.

if key == "last_prompt_id":
Expand Down
19 changes: 19 additions & 0 deletions tests/test_smartdatalake.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,25 @@ def test_last_result_is_saved(self, _mocked_method, smart_datalake: SmartDatalak
"value": "There are 10 countries in the dataframe.",
}

@patch.object(
CodeManager,
"execute_code",
return_value={
"type": "string",
"value": "There are 10 countries in the dataframe.",
},
)
@patch("pandasai.helpers.query_exec_tracker.QueryExecTracker.publish")
def test_query_tracker_publish_called_in_chat_method(
self, mock_query_tracker_publish, _mocked_method, smart_datalake: SmartDatalake
):
assert smart_datalake.last_result is None

_mocked_method.__name__ = "execute_code"

smart_datalake.chat("How many countries are in the dataframe?")
mock_query_tracker_publish.assert_called()

Comment on lines 149 to +170
Copy link
Contributor

Choose a reason for hiding this comment

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

The new test method test_query_tracker_publish_called_in_chat_method is correctly structured to test the invocation of the query_tracker_publish function. However, there are a few points to consider:

  1. The test does not verify the parameters with which mock_query_tracker_publish is called. It is important to ensure that the correct arguments are being passed to the function.
  2. The test assumes that smart_datalake.last_result is None before the chat method is called, which is good. However, it does not check the state of last_result after the chat method is called to ensure that it has been updated correctly.
  3. The test uses patching correctly, but it's not clear if the mock_query_tracker_publish is being reset after each test. This could potentially affect other tests if the mock state is not isolated.

Consider adding assertions to check the arguments passed to mock_query_tracker_publish and the state of last_result after the chat method is called. Also, ensure that mocks are reset after each test to maintain test isolation.

-        smart_datalake.chat("How many countries are in the dataframe?")
+        question = "How many countries are in the dataframe?"
+        smart_datalake.chat(question)
+        mock_query_tracker_publish.assert_called_with(question, ANY)
+        assert smart_datalake.last_result is not None
Commitable suggestion

[!IMPORTANT]
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Suggested change
"value": "There are 10 countries in the dataframe.",
}
@patch.object(
CodeManager,
"execute_code",
return_value={
"type": "string",
"value": "There are 10 countries in the dataframe.",
},
)
@patch("pandasai.helpers.query_exec_tracker.QueryExecTracker.publish")
def test_query_tracker_publish_called_in_chat_method(
self, mock_query_tracker_publish, _mocked_method, smart_datalake: SmartDatalake
):
assert smart_datalake.last_result is None
_mocked_method.__name__ = "execute_code"
smart_datalake.chat("How many countries are in the dataframe?")
mock_query_tracker_publish.assert_called()
@patch.object(
CodeManager,
"execute_code",
return_value={
"type": "string",
"value": "There are 10 countries in the dataframe.",
},
)
@patch("pandasai.helpers.query_exec_tracker.QueryExecTracker.publish")
def test_query_tracker_publish_called_in_chat_method(
self, mock_query_tracker_publish, _mocked_method, smart_datalake: SmartDatalake
):
assert smart_datalake.last_result is None
_mocked_method.__name__ = "execute_code"
question = "How many countries are in the dataframe?"
smart_datalake.chat(question)
mock_query_tracker_publish.assert_called_with(question, ANY)
assert smart_datalake.last_result is not None

def test_retry_on_error_with_single_df(
self, smart_datalake: SmartDatalake, smart_dataframe: SmartDataframe
):
Expand Down