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

feat[Agent]: add agent conversation code #584

Merged
merged 15 commits into from
Sep 24, 2023
Merged
Changes from 8 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
37 changes: 37 additions & 0 deletions examples/agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
import pandas as pd
from pandasai.agent import Agent
Copy link
Collaborator

Choose a reason for hiding this comment

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

We can change it to:
from pandasaiimport Agent


from pandasai.llm.openai import OpenAI

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

salaries_data = {
"EmployeeID": [1, 2, 3, 4, 5],
"Salary": [5000, 6000, 4500, 7000, 5500],
}

employees_df = pd.DataFrame(employees_data)
salaries_df = pd.DataFrame(salaries_data)


llm = OpenAI("OPEN_API_KEY")
agent = Agent([employees_df, salaries_df], config={"llm": llm}, memory_size=10)

# Chat with the agent
response = agent.chat("Who gets paid the most?")
print(response)


# # Get Clarification Questions
questions = agent.clarification_questions()

for question in questions:
print(question)

# Explain how the chat response is generated
response = agent.explain()
print(response)
3 changes: 2 additions & 1 deletion pandasai/__init__.py
Original file line number Diff line number Diff line change
@@ -44,6 +44,7 @@
from .callbacks.base import BaseCallback
from .schemas.df_config import Config
from .helpers.cache import Cache
from .agent import Agent

__version__ = importlib.metadata.version(__package__ or __name__)

@@ -257,4 +258,4 @@ def clear_cache(filename: str = None):
cache.clear()


__all__ = ["PandasAI", "SmartDataframe", "SmartDatalake", "clear_cache"]
__all__ = ["PandasAI", "SmartDataframe", "SmartDatalake", "Agent", "clear_cache"]
106 changes: 106 additions & 0 deletions pandasai/agent/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
import json
from typing import Union, List, Optional
from pandasai.helpers.df_info import DataFrameType
from pandasai.helpers.logger import Logger
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
_memory_size: int = None
gventuri marked this conversation as resolved.
Show resolved Hide resolved

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.
gventuri marked this conversation as resolved.
Show resolved Hide resolved
"""

if not isinstance(dfs, list):
dfs = [dfs]
gventuri marked this conversation as resolved.
Show resolved Hide resolved

gventuri marked this conversation as resolved.
Show resolved Hide resolved
Comment on lines +35 to +37
Copy link
Contributor

Choose a reason for hiding this comment

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

The code is assuming that if dfs is not a list, it should be converted into a list. This might lead to unexpected behavior if dfs is of an unsupported type. It would be better to explicitly check for supported types and raise an error if an unsupported type is provided.

-        if not isinstance(dfs, list):
-            dfs = [dfs]
+        if isinstance(dfs, DataFrameType):
+            dfs = [dfs]
+        elif not isinstance(dfs, list):
+            raise TypeError("dfs must be a DataFrameType or a list of DataFrameType")

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

def chat(self, query: str, output_type: Optional[str] = None):
"""
Simulate a chat interaction with the assistant on Dataframe.
"""
try:
result = self._lake.chat(
query,
output_type=output_type,
start_conversation=self._lake._memory.get_conversation(
self._memory_size
),
)
return result
except Exception as exception:
return (
"Unfortunately, I was not able to get your answers, "
"because of the following error:\n"
f"\n{exception}\n"
)
gventuri marked this conversation as resolved.
Show resolved Hide resolved

def clarification_questions(self) -> List[str]:
"""
Generate clarification questions based on the data
"""
try:
prompt = ClarificationQuestionPrompt(
self._lake.dfs, self._lake._memory.get_conversation(self._memory_size)
)

result = self._lake.llm.call(prompt)
self._logger.log(
f"""Clarification Questions: {result}
"""
)
questions: list[str] = json.loads(result)
return questions[:3]

except Exception as exception:
raise exception
Copy link
Contributor

@coderabbitai coderabbitai bot Sep 22, 2023

Choose a reason for hiding this comment

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

In the clarification_questions method, exceptions are caught and immediately re-raised. This is redundant and can be removed, as uncaught exceptions will automatically propagate up the call stack.

Copy link
Collaborator

Choose a reason for hiding this comment

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

@ArslanSaleem agree, I suggest we remove try/catch wrapper. In the future we might add a retry mechanism!

Copy link
Contributor

Choose a reason for hiding this comment

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

Looks like humans are talking. I'll be quiet for now. If you need me again, just tag @coderabbitai in a new comment.


def start_new_conversation(self):
gventuri marked this conversation as resolved.
Show resolved Hide resolved
"""
Clears the previous conversation
"""
self._lake._memory.clear()

gventuri marked this conversation as resolved.
Show resolved Hide resolved
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._memory_size),
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"
)
gventuri marked this conversation as resolved.
Show resolved Hide resolved
Comment on lines +92 to +97
Copy link
Contributor

Choose a reason for hiding this comment

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

Error handling in these methods only returns a string with the error message. It might be better to log the exception and re-raise it after logging. This way, the caller of the method can decide how to handle the exception. Also, this will help in debugging as the stack trace will be available in the logs.

-            return (
-                "Unfortunately, I was not able to get your answers, "
-                "because of the following error:\n"
-                f"\n{exception}\n"
-            )
+            self._logger.log(
+                "Unfortunately, I was not able to get your answers, "
+                "because of the following error:\n"
+                f"\n{exception}\n"
+            )
+            raise

...

-            return (
-                "Unfortunately, I was not able to explain, "
-                "because of the following error:\n"
-                f"\n{exception}\n"
-            )
+            self._logger.log(
+                "Unfortunately, I was not able to explain, "
-                "because of the following error:\n"
-                f"\n{exception}\n"
+            )
+            raise

51 changes: 51 additions & 0 deletions pandasai/prompts/clarification_questions_prompt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
""" Prompt to get clarification questions
You are provided with the following pandas DataFrames:

<dataframe>
{dataframe}
</dataframe>
Copy link
Collaborator

Choose a reason for hiding this comment

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

Let's also change this to

{dataframes}

for consistency!


<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
gventuri marked this conversation as resolved.
Show resolved Hide resolved


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>
Copy link
Collaborator

Choose a reason for hiding this comment

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

{dataframes} only should be enough, the set_var method in prompts already takes care of wrapping each one individually.


<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:
"""
gventuri marked this conversation as resolved.
Show resolved Hide resolved
gventuri marked this conversation as resolved.
Show resolved Hide resolved

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

Choose a reason for hiding this comment

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

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)

self.set_var("conversation", conversation)
gventuri marked this conversation as resolved.
Show resolved Hide resolved
34 changes: 34 additions & 0 deletions pandasai/prompts/explain_prompt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
""" Prompt to explain solution generated
Based on the last conversation you generated the code.
Can you explain briefly for non technical person on how you came up with code
without explaining pandas library?
"""


from .base import Prompt


class ExplainPrompt(Prompt):
"""Prompt to get clarification questions"""
Copy link
Contributor

Choose a reason for hiding this comment

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

The class docstring is misleading. This class is for generating explanation prompts, not clarification questions. Please correct it to avoid confusion.

-     """Prompt to get clarification questions"""
+     """Prompt to generate explanation for the code"""


text: str = """
The previous conversation we had

<Conversation>
{conversation}
</Conversation>

Based on the last conversation you generated the following code:

gventuri marked this conversation as resolved.
Show resolved Hide resolved
<Code>
{code}
</Code

Explain how you came up with code for non-technical people without
mentioning technical details or mentioning the libraries used?

"""
Comment on lines +24 to +40
Copy link
Contributor

Choose a reason for hiding this comment

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

The text attribute in the ExplainPrompt class is a class variable, which means it's shared across all instances of this class. If you modify it in one instance, it will affect all other instances. This could lead to unexpected behavior if multiple instances of ExplainPrompt are used concurrently. Consider moving this attribute to the instance level (inside __init__) to avoid potential issues.


def __init__(self, conversation, code):
self.set_var("conversation", conversation)
self.set_var("code", code)
gventuri marked this conversation as resolved.
Show resolved Hide resolved
20 changes: 18 additions & 2 deletions pandasai/smart_datalake/__init__.py
Original file line number Diff line number Diff line change
@@ -255,7 +255,12 @@ def _get_cache_key(self) -> str:

return cache_key

def chat(self, query: str, output_type: Optional[str] = None):
def chat(
self,
query: str,
output_type: Optional[str] = None,
start_conversation: Optional[str] = None,
Copy link
Collaborator

Choose a reason for hiding this comment

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

Let's remove this

Copy link
Collaborator Author

@ArslanSaleem ArslanSaleem Sep 22, 2023

Choose a reason for hiding this comment

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

@gventuri Currently the get_conversation method is called with default which has limit=1 means returning the last message only. We need to then pass the memory size to SmartLake constructor to do so or let's use memory for that.

):
"""
Run a query on the dataframe.

@@ -305,6 +310,9 @@ def chat(self, query: str, output_type: Optional[str] = None):
"save_charts_path": self._config.save_charts_path.rstrip("/"),
"output_type_hint": output_type_helper.template_hint,
}
if start_conversation is not None:
default_values["conversation"] = start_conversation
gventuri marked this conversation as resolved.
Show resolved Hide resolved

generate_python_code_instruction = self._get_prompt(
"generate_python_code",
default_prompt=GeneratePythonCodePrompt,
@@ -623,7 +631,7 @@ def last_code_generated(self):

@last_code_generated.setter
def last_code_generated(self, last_code_generated: str):
self._code_manager._last_code_generated = last_code_generated
self._last_code_generated = last_code_generated
gventuri marked this conversation as resolved.
Show resolved Hide resolved

@property
def last_code_executed(self):
@@ -644,3 +652,11 @@ def last_error(self):
@last_error.setter
def last_error(self, last_error: str):
self._last_error = last_error

@property
def dfs(self):
return self._dfs

@property
def memory(self):
return self._memory
gventuri marked this conversation as resolved.
Show resolved Hide resolved
Loading