-
Notifications
You must be signed in to change notification settings - Fork 1.4k
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
Changes from 8 commits
70b0da8
1b51727
70244c3
f715035
2da890a
6736c44
cdeec68
9025f4e
d1b8e61
49d8720
b92fb39
7f17af8
2e4c902
7a554a5
f7e4d98
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
import pandas as pd | ||
from pandasai.agent import 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) |
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
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The code is assuming that if - 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 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In the There was a problem hiding this comment. Choose a reason for hiding this commentThe 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! There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
|
||
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
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
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> | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's also change this to
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> | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
|
||
<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) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
This method from |
||
self.set_var("conversation", conversation) | ||
gventuri marked this conversation as resolved.
Show resolved
Hide resolved
|
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""" | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The |
||
|
||
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
|
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, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's remove this There was a problem hiding this comment. Choose a reason for hiding this commentThe 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
|
There was a problem hiding this comment.
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