fix(code manager): parsing of called functions

[mspronesti]

Hi @gventuri,
this PR fixes the detection of called functions (skills) when parsing the generated code.
Example of failing case:

import pandas as pd
from pandasai import Agent

from pandasai.llm import OpenAI
import os
from pandasai.skills import skill

employees_data = {
    "EmployeeID": [1, 2, 3, 4, 5],
    "Name": ["John", "Emma", "Liam", "Olivia", "William"],
    "Department": ["HR", "Sales", "IT", "Marketing", "Finance"],
}

employees_df = pd.DataFrame(employees_data)
employees_df.attrs['name'] = 'foo_df'


# Add function docstring to give more context to model
@skill()
def get_dataframe_name(df: pd.DataFrame) -> str:
    """
    Returns the name of the given dataframe.
    Args:
        df (pd.Dataframe)
    """
    if df is not None and "name" in df.attrs:
        return df.attrs['name']
    else:
        return ""


os.environ["OPENAI_API_KEY"] = "sk-..."

llm = OpenAI(temperature=0)

agent = Agent([employees_df], config={"llm": llm, "enable_cache": False, "verbose": True}, memory_size=10)

agent.add_skills(get_dataframe_name)

# Chat with the agent
response = agent.chat("What is the name of the provided dataframe?")
print(response)

Generated code:

df = dfs[0]
name = get_dataframe_name(df)
result = {'type': 'string', 'value': name}

Which fails (without this fix) with the following error:

2024-01-20 15:44:53 [ERROR] Failed with error: Traceback (most recent call last):
  File "/home/mspronesti/Desktop/pandas-ai/pandasai/pipelines/smart_datalake_chat/code_execution.py", line 52, in execute
    result = pipeline_context.query_exec_tracker.execute_func(
  File "/home/mspronesti/Desktop/pandas-ai/pandasai/helpers/query_exec_tracker.py", line 128, in execute_func
    result = function(*args, **kwargs)
  File "/home/mspronesti/Desktop/pandas-ai/pandasai/helpers/code_manager.py", line 212, in execute_code
    exec(code_to_run, environment)
  File "<string>", line 2, in <module>
NameError: name 'get_dataframe_name' is not defined

Summary by CodeRabbit

• Refactor
  • Improved code indentation for better readability and maintainability.
  • Streamlined the identification of unsafe methods to enhance security checks.
  • Optimized the function call detection within the code execution context.
  • Refined the logic for identifying unauthorized table use in SQL queries.
  • Enhanced the detection of dataframe overwrites to prevent unintended data loss.
  • Updated the import checks to align with whitelisted libraries and custom dependencies.
  • Improved utility functions for dataframe identification and comparison extraction.

[coderabbitai]

Warning

Rate Limit Exceeded

@mspronesti has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 21 minutes and 58 seconds before requesting another review.

How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.
Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.
Please see our FAQ for further information.

Commits

Files that changed from the base of the PR and between 81a7173 and 49a9fd9.

Walkthrough

The changes primarily involve formatting adjustments, such as aligning indentation within class definitions and method arguments. Additionally, there's a simplification in the list comprehension within a security-related method. A few logical structures have been condensed for clarity, and there are minor changes to variable assignments and function calls to enhance readability and maintainability.

Changes

File Path	Change Summary
.../helpers/code_manager.py	Adjusted indentation in class definitions and methods.
.../helpers/code_manager.py	Simplified list comprehension in security checks.
.../helpers/code_manager.py	Streamlined logical conditions and assignment structures.

🐰✨
In code's vast garden, indents align,
A rabbit hops through, making it fine.
Security tightens, with lists so sleek,
Bugs scatter away, with every tweak. 🌟🐾

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share

• X
• Mastodon
• Reddit
• LinkedIn

---

Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>.
  • Generate unit-tests for this file.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit tests for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai generate interesting stats about this repository from git and render them as a table.
  • @coderabbitai show all the console.log statements in this repository.
  • @coderabbitai read src/utils.ts and generate unit tests.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (invoked as PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger a review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai help to get help.

Additionally, you can add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.

CodeRabbit Configration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• The JSON schema for the configuration file is available here.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/coderabbit-overrides.v2.json

CodeRabbit Discord Community

Join our Discord Community to get help, request features, and share feedback.

[sourcery-ai]

PR Type: Bug fix

PR Summary: The pull request addresses an issue with the detection of called functions within the code execution context. It modifies the logic for identifying function calls in the abstract syntax tree, ensuring that the skills manager correctly recognizes and adds used skills. Additionally, the PR includes updates to the openai_info module, adding new fine-tuned model identifiers and their associated costs.

Decision: Comment

📝 Type: 'Bug fix' - not supported yet.

• Sourcery currently only approves 'Typo fix' PRs.

✅ Issue addressed: this change correctly addresses the issue or implements the desired feature.

No details provided.

✅ Small diff: the diff is small enough to approve with confidence.

No details provided.

General suggestions:

• Ensure that the refactoring of the function call detection logic does not introduce any regressions by thoroughly testing all possible cases where a function call might occur.
• Review the changes to ensure that they adhere to the project's coding standards, particularly with respect to indentation and formatting.
• Verify that the addition of new model identifiers in the openai_info module is accompanied by the necessary updates elsewhere in the codebase to handle these new models.

Thanks for using Sourcery. We offer it for free for open source projects and would be very grateful if you could help us grow. If you like it, would you consider sharing Sourcery on your favourite social media? ✨

Share Sourcery

• X
• Mastodon
• LinkedIn
• Facebook

---

Help me be more useful! Please click 👍 or 👎 on each comment to tell me if it was helpful.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 3f4b1e5 and 81a7173.

Files selected for processing (2)

• pandasai/helpers/code_manager.py (9 hunks)
• pandasai/helpers/openai_info.py (1 hunks)

Additional comments: 13

pandasai/helpers/openai_info.py (2)

• 47-47: The addition of the new model "gpt-3.5-turbo-1106-finetuned" with its associated cost is consistent with the existing pattern of model cost definitions.
• 50-50: The addition of the new model "gpt-3.5-turbo-1106-finetuned-completion" with its associated cost is consistent with the existing pattern of model cost definitions.

pandasai/helpers/code_manager.py (11)

• 32-34: The indentation of the constructor parameters has been increased. This is a style change and does not affect the functionality of the code.
• 75-78: The indentation of the constructor parameters has been increased. This is a style change and does not affect the functionality of the code.
• 296-311: The streamlining of the list of methods checked by the _is_unsafe function is a formatting change that improves readability without altering functionality.
• 316-322: The refactoring of the find_function_calls method improves the detection of function calls by considering both ast.Name and ast.Attribute types. This change is crucial for recognizing functions decorated with @skill() and making them available during code execution.
• 328-330: The indentation change in the check_direct_sql_func_def_exists method is a style change and does not affect the functionality of the code.
• 336-339: The indentation change in the _get_sql_irrelevant_tables method is a style change and does not affect the functionality of the code.
• 375-377: The indentation change in the _clean_code method is a style change and does not affect the functionality of the code.
• 388-390: The indentation change in the _clean_code method is a style change and does not affect the functionality of the code.
• 415-417: The indentation change in the _is_df_overwrite method is a style change and does not affect the functionality of the code.
• 438-439: The indentation change in the _check_imports method is a style change and does not affect the functionality of the code.
• 611-614: The indentation change in the _extract_comparisons method is a style change and does not affect the functionality of the code.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (470b70a) 85.21% compared to head (49a9fd9) 85.22%.
Report is 3 commits behind head on main.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #883      +/-   ##
==========================================
+ Coverage   85.21%   85.22%   +0.01%     
==========================================
  Files          90       90              
  Lines        3881     3885       +4     
==========================================
+ Hits         3307     3311       +4     
  Misses        574      574              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[gventuri]

Thanks a lot for the great fix @mspronesti!!