ai-data-science-team 0.0.0.9009__tar.gz → 0.0.0.9010__tar.gz

Sign up to get free protection for your applications and to get access to all the features.
Files changed (44) hide show
  1. {ai_data_science_team-0.0.0.9009/ai_data_science_team.egg-info → ai_data_science_team-0.0.0.9010}/PKG-INFO +34 -16
  2. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/README.md +28 -13
  3. ai_data_science_team-0.0.0.9010/ai_data_science_team/_version.py +1 -0
  4. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/agents/data_cleaning_agent.py +6 -6
  5. ai_data_science_team-0.0.0.9010/ai_data_science_team/agents/data_loader_tools_agent.py +69 -0
  6. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/agents/data_visualization_agent.py +6 -7
  7. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/agents/data_wrangling_agent.py +6 -6
  8. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/agents/feature_engineering_agent.py +6 -6
  9. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/agents/sql_database_agent.py +6 -6
  10. ai_data_science_team-0.0.0.9010/ai_data_science_team/ml_agents/__init__.py +2 -0
  11. ai_data_science_team-0.0.0.9010/ai_data_science_team/ml_agents/h2o_ml_agent.py +852 -0
  12. ai_data_science_team-0.0.0.9010/ai_data_science_team/ml_agents/mlflow_tools_agent.py +327 -0
  13. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/multiagents/sql_data_analyst.py +3 -4
  14. {ai_data_science_team-0.0.0.9009/ai_data_science_team/tools → ai_data_science_team-0.0.0.9010/ai_data_science_team/parsers}/parsers.py +0 -1
  15. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/templates/agent_templates.py +6 -6
  16. ai_data_science_team-0.0.0.9010/ai_data_science_team/tools/data_loader.py +378 -0
  17. ai_data_science_team-0.0.0.9010/ai_data_science_team/tools/dataframe.py +139 -0
  18. ai_data_science_team-0.0.0.9009/ai_data_science_team/ml_agents/h2o_ml_agent.py → ai_data_science_team-0.0.0.9010/ai_data_science_team/tools/h2o.py +294 -683
  19. ai_data_science_team-0.0.0.9010/ai_data_science_team/tools/mlflow.py +961 -0
  20. ai_data_science_team-0.0.0.9009/ai_data_science_team/tools/metadata.py → ai_data_science_team-0.0.0.9010/ai_data_science_team/tools/sql.py +1 -137
  21. ai_data_science_team-0.0.0.9010/ai_data_science_team/utils/__init__.py +0 -0
  22. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010/ai_data_science_team.egg-info}/PKG-INFO +34 -16
  23. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team.egg-info/SOURCES.txt +12 -5
  24. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team.egg-info/requires.txt +4 -1
  25. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/requirements.txt +1 -0
  26. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/setup.py +2 -2
  27. ai_data_science_team-0.0.0.9009/ai_data_science_team/_version.py +0 -1
  28. ai_data_science_team-0.0.0.9009/ai_data_science_team/ml_agents/__init__.py +0 -1
  29. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/LICENSE +0 -0
  30. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/MANIFEST.in +0 -0
  31. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/__init__.py +0 -0
  32. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/agents/__init__.py +0 -0
  33. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/multiagents/__init__.py +0 -0
  34. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/multiagents/supervised_data_analyst.py +0 -0
  35. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/orchestration.py +0 -0
  36. {ai_data_science_team-0.0.0.9009/ai_data_science_team/tools → ai_data_science_team-0.0.0.9010/ai_data_science_team/parsers}/__init__.py +0 -0
  37. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/templates/__init__.py +0 -0
  38. {ai_data_science_team-0.0.0.9009/ai_data_science_team/utils → ai_data_science_team-0.0.0.9010/ai_data_science_team/tools}/__init__.py +0 -0
  39. {ai_data_science_team-0.0.0.9009/ai_data_science_team/tools → ai_data_science_team-0.0.0.9010/ai_data_science_team/utils}/logging.py +0 -0
  40. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team/utils/plotly.py +0 -0
  41. {ai_data_science_team-0.0.0.9009/ai_data_science_team/tools → ai_data_science_team-0.0.0.9010/ai_data_science_team/utils}/regex.py +0 -0
  42. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team.egg-info/dependency_links.txt +0 -0
  43. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/ai_data_science_team.egg-info/top_level.txt +0 -0
  44. {ai_data_science_team-0.0.0.9009 → ai_data_science_team-0.0.0.9010}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: ai-data-science-team
3
- Version: 0.0.0.9009
3
+ Version: 0.0.0.9010
4
4
  Summary: Build and run an AI-powered data science team.
5
5
  Home-page: https://github.com/business-science/ai-data-science-team
6
6
  Author: Matt Dancho
@@ -27,10 +27,13 @@ Requires-Dist: plotly
27
27
  Requires-Dist: streamlit
28
28
  Requires-Dist: scikit-learn
29
29
  Requires-Dist: xgboost
30
- Provides-Extra: machine-learning-agent
31
- Requires-Dist: h2o; extra == "machine-learning-agent"
30
+ Requires-Dist: psutil
31
+ Provides-Extra: machine-learning
32
+ Requires-Dist: h2o; extra == "machine-learning"
33
+ Requires-Dist: mlflow; extra == "machine-learning"
32
34
  Provides-Extra: all
33
35
  Requires-Dist: h2o; extra == "all"
36
+ Requires-Dist: mlflow; extra == "all"
34
37
  Dynamic: author
35
38
  Dynamic: author-email
36
39
  Dynamic: classifier
@@ -45,7 +48,7 @@ Dynamic: summary
45
48
  <div align="center">
46
49
  <a href="https://github.com/business-science/ai-data-science-team">
47
50
  <picture>
48
- <img src="/img/ai_data_science_team_logo.jpg" alt="AI Data Science Team" width="400">
51
+ <img src="/img/ai_data_science_team_logo_small.jpg" alt="AI Data Science Team" width="400">
49
52
  </picture>
50
53
  </a>
51
54
  </div>
@@ -86,8 +89,11 @@ The AI Data Science Team of Copilots includes Agents that specialize data cleani
86
89
  - [Generative AI for Data Scientists Workshop](#generative-ai-for-data-scientists-workshop)
87
90
  - [Data Science Agents](#data-science-agents)
88
91
  - [NEW: Multi-Agents](#new-multi-agents)
89
- - [Coming Soon: Data Science Apps](#coming-soon-data-science-apps)
92
+ - [Data Science Apps](#data-science-apps)
93
+ - [Apps Available Now](#apps-available-now)
94
+ - [🔥 Agentic Applications](#-agentic-applications)
90
95
  - [Agents Available Now](#agents-available-now)
96
+ - [🔥🔥 NEW! Machine Learning Agents](#-new-machine-learning-agents)
91
97
  - [Data Science Agents](#data-science-agents-1)
92
98
  - [Multi-Agents](#multi-agents)
93
99
  - [Agents Coming Soon](#agents-coming-soon)
@@ -124,32 +130,44 @@ This is the internals of the SQL Data Analyst Agent that connects to SQL databas
124
130
 
125
131
  ![Business Intelligence SQL Agent](/img/multi_agent_sql_data_visualization.jpg)
126
132
 
127
- ### Coming Soon: Data Science Apps
133
+ ### Data Science Apps
128
134
 
129
135
  This is a top secret project I'm working on. It's a multi-agent data science app that performs time series forecasting.
130
136
 
131
- ![Multi-Agent Data Science App](/img/ai_powered_apps.jpg)
137
+ ![Multi-Agent Data Science App](/img/ai_powered_apps.jpg)
138
+
139
+ ### Apps Available Now
140
+
141
+ [See all available apps here](/apps)
142
+
143
+ #### 🔥 Agentic Applications
144
+
145
+ 1. **SQL Database Agent App:** Connects any SQL Database, generates SQL queries from natural language, and returns data as a downloadable table. [See Application](/apps/sql-database-agent-app/)
132
146
 
133
147
  ### Agents Available Now
134
148
 
149
+ #### 🔥🔥 NEW! Machine Learning Agents
150
+
151
+ 1. **🔥 H2O Machine Learning Agent:** Builds and logs 100's of high-performance machine learning models. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/ml_agents/h2o_machine_learning_agent.ipynb)
152
+ 2. **🔥 MLflow Tools Agent (MLOps):** This agent has 11+ tools for managing models, ML projects, and making production ML predictions with MLflow. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/ml_agents/mlflow_tools_agent.ipynb)
153
+
135
154
  #### Data Science Agents
136
155
 
137
- 1. **Data Wrangling Agent:** Merges, Joins, Preps and Wrangles data into a format that is ready for data analysis.
138
- 2. **Data Visualization Agent:** Creates visualizations to help you understand your data. Returns JSON serializable plotly visualizations.
139
- 3. **Data Cleaning Agent:** Performs Data Preparation steps including handling missing values, outliers, and data type conversions.
140
- 4. **Feature Engineering Agent:** Converts the prepared data into ML-ready data. Adds features to increase predictive accuracy of ML models.
141
- 5. **SQL Database Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations.
156
+ 1. **Data Wrangling Agent:** Merges, Joins, Preps and Wrangles data into a format that is ready for data analysis. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/data_wrangling_agent.ipynb)
157
+ 2. **Data Visualization Agent:** Creates visualizations to help you understand your data. Returns JSON serializable plotly visualizations. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/data_visualization_agent.ipynb)
158
+ 3. **Data Cleaning Agent:** Performs Data Preparation steps including handling missing values, outliers, and data type conversions. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/data_cleaning_agent.ipynb)
159
+ 4. **Feature Engineering Agent:** Converts the prepared data into ML-ready data. Adds features to increase predictive accuracy of ML models. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/feature_engineering_agent.ipynb)
160
+ 5. **SQL Database Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/sql_database_agent.ipynb)
142
161
 
143
162
  #### Multi-Agents
144
163
 
145
- 1. **SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data.
164
+ 1. **SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/sql_data_analyst.ipynb)
146
165
 
147
166
  ### Agents Coming Soon
148
167
 
149
168
  1. **Data Analyst:** Analyzes data structure, creates exploratory visualizations, and performs correlation analysis to identify relationships.
150
- 2. **Machine Learning Agent:** Builds and logs the machine learning models.
151
- 3. **Interpretability Agent:** Performs Interpretable ML to explain why the model returned predictions including which features were the most important to the model.
152
- 4. **Supervisor:** Forms task list. Moderates sub-agents. Returns completed assignment.
169
+ 2. **Interpretability Agent:** Performs Interpretable ML to explain why the model returned predictions including which features were the most important to the model.
170
+ 3. **Supervisor:** Forms task list. Moderates sub-agents. Returns completed assignment.
153
171
 
154
172
  ## Disclaimer
155
173
 
@@ -1,7 +1,7 @@
1
1
  <div align="center">
2
2
  <a href="https://github.com/business-science/ai-data-science-team">
3
3
  <picture>
4
- <img src="/img/ai_data_science_team_logo.jpg" alt="AI Data Science Team" width="400">
4
+ <img src="/img/ai_data_science_team_logo_small.jpg" alt="AI Data Science Team" width="400">
5
5
  </picture>
6
6
  </a>
7
7
  </div>
@@ -42,8 +42,11 @@ The AI Data Science Team of Copilots includes Agents that specialize data cleani
42
42
  - [Generative AI for Data Scientists Workshop](#generative-ai-for-data-scientists-workshop)
43
43
  - [Data Science Agents](#data-science-agents)
44
44
  - [NEW: Multi-Agents](#new-multi-agents)
45
- - [Coming Soon: Data Science Apps](#coming-soon-data-science-apps)
45
+ - [Data Science Apps](#data-science-apps)
46
+ - [Apps Available Now](#apps-available-now)
47
+ - [🔥 Agentic Applications](#-agentic-applications)
46
48
  - [Agents Available Now](#agents-available-now)
49
+ - [🔥🔥 NEW! Machine Learning Agents](#-new-machine-learning-agents)
47
50
  - [Data Science Agents](#data-science-agents-1)
48
51
  - [Multi-Agents](#multi-agents)
49
52
  - [Agents Coming Soon](#agents-coming-soon)
@@ -80,32 +83,44 @@ This is the internals of the SQL Data Analyst Agent that connects to SQL databas
80
83
 
81
84
  ![Business Intelligence SQL Agent](/img/multi_agent_sql_data_visualization.jpg)
82
85
 
83
- ### Coming Soon: Data Science Apps
86
+ ### Data Science Apps
84
87
 
85
88
  This is a top secret project I'm working on. It's a multi-agent data science app that performs time series forecasting.
86
89
 
87
- ![Multi-Agent Data Science App](/img/ai_powered_apps.jpg)
90
+ ![Multi-Agent Data Science App](/img/ai_powered_apps.jpg)
91
+
92
+ ### Apps Available Now
93
+
94
+ [See all available apps here](/apps)
95
+
96
+ #### 🔥 Agentic Applications
97
+
98
+ 1. **SQL Database Agent App:** Connects any SQL Database, generates SQL queries from natural language, and returns data as a downloadable table. [See Application](/apps/sql-database-agent-app/)
88
99
 
89
100
  ### Agents Available Now
90
101
 
102
+ #### 🔥🔥 NEW! Machine Learning Agents
103
+
104
+ 1. **🔥 H2O Machine Learning Agent:** Builds and logs 100's of high-performance machine learning models. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/ml_agents/h2o_machine_learning_agent.ipynb)
105
+ 2. **🔥 MLflow Tools Agent (MLOps):** This agent has 11+ tools for managing models, ML projects, and making production ML predictions with MLflow. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/ml_agents/mlflow_tools_agent.ipynb)
106
+
91
107
  #### Data Science Agents
92
108
 
93
- 1. **Data Wrangling Agent:** Merges, Joins, Preps and Wrangles data into a format that is ready for data analysis.
94
- 2. **Data Visualization Agent:** Creates visualizations to help you understand your data. Returns JSON serializable plotly visualizations.
95
- 3. **Data Cleaning Agent:** Performs Data Preparation steps including handling missing values, outliers, and data type conversions.
96
- 4. **Feature Engineering Agent:** Converts the prepared data into ML-ready data. Adds features to increase predictive accuracy of ML models.
97
- 5. **SQL Database Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations.
109
+ 1. **Data Wrangling Agent:** Merges, Joins, Preps and Wrangles data into a format that is ready for data analysis. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/data_wrangling_agent.ipynb)
110
+ 2. **Data Visualization Agent:** Creates visualizations to help you understand your data. Returns JSON serializable plotly visualizations. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/data_visualization_agent.ipynb)
111
+ 3. **Data Cleaning Agent:** Performs Data Preparation steps including handling missing values, outliers, and data type conversions. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/data_cleaning_agent.ipynb)
112
+ 4. **Feature Engineering Agent:** Converts the prepared data into ML-ready data. Adds features to increase predictive accuracy of ML models. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/feature_engineering_agent.ipynb)
113
+ 5. **SQL Database Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/sql_database_agent.ipynb)
98
114
 
99
115
  #### Multi-Agents
100
116
 
101
- 1. **SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data.
117
+ 1. **SQL Data Analyst Agent:** Connects to SQL databases to pull data into the data science environment. Creates pipelines to automate data extraction. Performs Joins, Aggregations, and other SQL Query operations. Includes a Data Visualization Agent that creates visualizations to help you understand your data. [See Example](https://github.com/business-science/ai-data-science-team/blob/master/examples/multiagents/sql_data_analyst.ipynb)
102
118
 
103
119
  ### Agents Coming Soon
104
120
 
105
121
  1. **Data Analyst:** Analyzes data structure, creates exploratory visualizations, and performs correlation analysis to identify relationships.
106
- 2. **Machine Learning Agent:** Builds and logs the machine learning models.
107
- 3. **Interpretability Agent:** Performs Interpretable ML to explain why the model returned predictions including which features were the most important to the model.
108
- 4. **Supervisor:** Forms task list. Moderates sub-agents. Returns completed assignment.
122
+ 2. **Interpretability Agent:** Performs Interpretable ML to explain why the model returned predictions including which features were the most important to the model.
123
+ 3. **Supervisor:** Forms task list. Moderates sub-agents. Returns completed assignment.
109
124
 
110
125
  ## Disclaimer
111
126
 
@@ -0,0 +1 @@
1
+ __version__ = "0.0.0.9010"
@@ -27,16 +27,16 @@ from ai_data_science_team.templates import(
27
27
  create_coding_agent_graph,
28
28
  BaseAgent,
29
29
  )
30
- from ai_data_science_team.tools.parsers import PythonOutputParser
31
- from ai_data_science_team.tools.regex import (
30
+ from ai_data_science_team.parsers.parsers import PythonOutputParser
31
+ from ai_data_science_team.utils.regex import (
32
32
  relocate_imports_inside_function,
33
33
  add_comments_to_top,
34
34
  format_agent_name,
35
35
  format_recommended_steps,
36
36
  get_generic_summary,
37
37
  )
38
- from ai_data_science_team.tools.metadata import get_dataframe_summary
39
- from ai_data_science_team.tools.logging import log_ai_function
38
+ from ai_data_science_team.tools.dataframe import get_dataframe_summary
39
+ from ai_data_science_team.utils.logging import log_ai_function
40
40
 
41
41
  # Setup
42
42
  AGENT_NAME = "data_cleaning_agent"
@@ -183,7 +183,7 @@ class DataCleaningAgent(BaseAgent):
183
183
  self.response=None
184
184
  return make_data_cleaning_agent(**self._params)
185
185
 
186
- def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
186
+ async def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
187
187
  """
188
188
  Asynchronously invokes the agent. The response is stored in the response attribute.
189
189
 
@@ -204,7 +204,7 @@ class DataCleaningAgent(BaseAgent):
204
204
  --------
205
205
  None. The response is stored in the response attribute.
206
206
  """
207
- response = self._compiled_graph.ainvoke({
207
+ response = await self._compiled_graph.ainvoke({
208
208
  "user_instructions": user_instructions,
209
209
  "data_raw": data_raw.to_dict(),
210
210
  "max_retries": max_retries,
@@ -0,0 +1,69 @@
1
+
2
+
3
+
4
+ from typing import Any, Optional, Annotated, Sequence, List, Dict
5
+ import operator
6
+
7
+ import pandas as pd
8
+ import os
9
+
10
+ from IPython.display import Markdown
11
+
12
+ from langchain_core.messages import BaseMessage, AIMessage
13
+
14
+ from langgraph.prebuilt import create_react_agent, ToolNode
15
+ from langgraph.prebuilt.chat_agent_executor import AgentState
16
+ from langgraph.graph import START, END, StateGraph
17
+
18
+ from ai_data_science_team.templates import BaseAgent
19
+ from ai_data_science_team.utils.regex import format_agent_name
20
+ from ai_data_science_team.tools.data_loader import (
21
+ load_directory,
22
+ load_file,
23
+ list_directory_contents,
24
+ list_directory_recursive,
25
+ get_file_info,
26
+ search_files_by_pattern,
27
+ )
28
+
29
+ AGENT_NAME = "data_loader_tools_agent"
30
+
31
+ tools = [
32
+ load_directory,
33
+ load_file,
34
+ list_directory_contents,
35
+ list_directory_recursive,
36
+ get_file_info,
37
+ search_files_by_pattern,
38
+ ]
39
+
40
+
41
+
42
+ def make_data_loader_tools_agent(
43
+ model: Any,
44
+ directory: Optional[str] = os.getcwd(),
45
+ ):
46
+ """
47
+ Creates a Data Loader Agent that can interact with data loading tools.
48
+
49
+ Parameters:
50
+ ----------
51
+ model : langchain.llms.base.LLM
52
+ The language model used to generate the tool calling agent.
53
+ directory : str, optional
54
+ The directory to search for files. Defaults to the current working directory.
55
+
56
+ Returns:
57
+ --------
58
+ Data Loader Agent
59
+ An agent that can interact with data loading tools.
60
+ """
61
+
62
+ class GraphState(AgentState):
63
+ internal_messages: Annotated[Sequence[BaseMessage], operator.add]
64
+ directory: str
65
+ user_instructions: str
66
+ data_artifacts: dict
67
+
68
+ pass
69
+
@@ -10,7 +10,6 @@ from typing import TypedDict, Annotated, Sequence, Literal
10
10
  import operator
11
11
 
12
12
  from langchain.prompts import PromptTemplate
13
- from langchain_core.output_parsers import StrOutputParser
14
13
  from langchain_core.messages import BaseMessage
15
14
 
16
15
  from langgraph.types import Command
@@ -30,16 +29,16 @@ from ai_data_science_team.templates import(
30
29
  create_coding_agent_graph,
31
30
  BaseAgent,
32
31
  )
33
- from ai_data_science_team.tools.parsers import PythonOutputParser
34
- from ai_data_science_team.tools.regex import (
32
+ from ai_data_science_team.parsers.parsers import PythonOutputParser
33
+ from ai_data_science_team.utils.regex import (
35
34
  relocate_imports_inside_function,
36
35
  add_comments_to_top,
37
36
  format_agent_name,
38
37
  format_recommended_steps,
39
38
  get_generic_summary,
40
39
  )
41
- from ai_data_science_team.tools.metadata import get_dataframe_summary
42
- from ai_data_science_team.tools.logging import log_ai_function
40
+ from ai_data_science_team.tools.dataframe import get_dataframe_summary
41
+ from ai_data_science_team.utils.logging import log_ai_function
43
42
  from ai_data_science_team.utils.plotly import plotly_from_dict
44
43
 
45
44
  # Setup
@@ -197,7 +196,7 @@ class DataVisualizationAgent(BaseAgent):
197
196
  # Rebuild the compiled graph
198
197
  self._compiled_graph = self._make_compiled_graph()
199
198
 
200
- def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
199
+ async def ainvoke_agent(self, data_raw: pd.DataFrame, user_instructions: str=None, max_retries:int=3, retry_count:int=0, **kwargs):
201
200
  """
202
201
  Asynchronously invokes the agent to generate a visualization.
203
202
  The response is stored in the 'response' attribute.
@@ -219,7 +218,7 @@ class DataVisualizationAgent(BaseAgent):
219
218
  -------
220
219
  None
221
220
  """
222
- response = self._compiled_graph.ainvoke({
221
+ response = await self._compiled_graph.ainvoke({
223
222
  "user_instructions": user_instructions,
224
223
  "data_raw": data_raw.to_dict(),
225
224
  "max_retries": max_retries,
@@ -24,16 +24,16 @@ from ai_data_science_team.templates import(
24
24
  create_coding_agent_graph,
25
25
  BaseAgent,
26
26
  )
27
- from ai_data_science_team.tools.parsers import PythonOutputParser
28
- from ai_data_science_team.tools.regex import (
27
+ from ai_data_science_team.parsers.parsers import PythonOutputParser
28
+ from ai_data_science_team.utils.regex import (
29
29
  relocate_imports_inside_function,
30
30
  add_comments_to_top,
31
31
  format_agent_name,
32
32
  format_recommended_steps,
33
33
  get_generic_summary,
34
34
  )
35
- from ai_data_science_team.tools.metadata import get_dataframe_summary
36
- from ai_data_science_team.tools.logging import log_ai_function
35
+ from ai_data_science_team.tools.dataframe import get_dataframe_summary
36
+ from ai_data_science_team.utils.logging import log_ai_function
37
37
 
38
38
  # Setup Logging Path
39
39
  AGENT_NAME = "data_wrangling_agent"
@@ -213,7 +213,7 @@ class DataWranglingAgent(BaseAgent):
213
213
  self._params[k] = v
214
214
  self._compiled_graph = self._make_compiled_graph()
215
215
 
216
- def ainvoke_agent(
216
+ async def ainvoke_agent(
217
217
  self,
218
218
  data_raw: Union[pd.DataFrame, dict, list],
219
219
  user_instructions: str=None,
@@ -245,7 +245,7 @@ class DataWranglingAgent(BaseAgent):
245
245
  None
246
246
  """
247
247
  data_input = self._convert_data_input(data_raw)
248
- response = self._compiled_graph.ainvoke({
248
+ response = await self._compiled_graph.ainvoke({
249
249
  "user_instructions": user_instructions,
250
250
  "data_raw": data_input,
251
251
  "max_retries": max_retries,
@@ -27,16 +27,16 @@ from ai_data_science_team.templates import(
27
27
  create_coding_agent_graph,
28
28
  BaseAgent,
29
29
  )
30
- from ai_data_science_team.tools.parsers import PythonOutputParser
31
- from ai_data_science_team.tools.regex import (
30
+ from ai_data_science_team.parsers.parsers import PythonOutputParser
31
+ from ai_data_science_team.utils.regex import (
32
32
  relocate_imports_inside_function,
33
33
  add_comments_to_top,
34
34
  format_agent_name,
35
35
  format_recommended_steps,
36
36
  get_generic_summary,
37
37
  )
38
- from ai_data_science_team.tools.metadata import get_dataframe_summary
39
- from ai_data_science_team.tools.logging import log_ai_function
38
+ from ai_data_science_team.tools.dataframe import get_dataframe_summary
39
+ from ai_data_science_team.utils.logging import log_ai_function
40
40
 
41
41
  # Setup
42
42
  AGENT_NAME = "feature_engineering_agent"
@@ -203,7 +203,7 @@ class FeatureEngineeringAgent(BaseAgent):
203
203
  self._params[k] = v
204
204
  self._compiled_graph = self._make_compiled_graph()
205
205
 
206
- def ainvoke_agent(
206
+ async def ainvoke_agent(
207
207
  self,
208
208
  data_raw: pd.DataFrame,
209
209
  user_instructions: str=None,
@@ -235,7 +235,7 @@ class FeatureEngineeringAgent(BaseAgent):
235
235
  -------
236
236
  None
237
237
  """
238
- response = self._compiled_graph.ainvoke({
238
+ response = await self._compiled_graph.ainvoke({
239
239
  "user_instructions": user_instructions,
240
240
  "data_raw": data_raw.to_dict(),
241
241
  "target_variable": target_variable,
@@ -25,15 +25,15 @@ from ai_data_science_team.templates import(
25
25
  create_coding_agent_graph,
26
26
  BaseAgent,
27
27
  )
28
- from ai_data_science_team.tools.parsers import SQLOutputParser
29
- from ai_data_science_team.tools.regex import (
28
+ from ai_data_science_team.parsers.parsers import SQLOutputParser
29
+ from ai_data_science_team.utils.regex import (
30
30
  add_comments_to_top,
31
31
  format_agent_name,
32
32
  format_recommended_steps,
33
33
  get_generic_summary,
34
34
  )
35
- from ai_data_science_team.tools.metadata import get_database_metadata
36
- from ai_data_science_team.tools.logging import log_ai_function
35
+ from ai_data_science_team.tools.sql import get_database_metadata
36
+ from ai_data_science_team.utils.logging import log_ai_function
37
37
 
38
38
  # Setup
39
39
  AGENT_NAME = "sql_database_agent"
@@ -193,7 +193,7 @@ class SQLDatabaseAgent(BaseAgent):
193
193
  self._params[k] = v
194
194
  self._compiled_graph = self._make_compiled_graph()
195
195
 
196
- def ainvoke_agent(self, user_instructions: str=None, max_retries=3, retry_count=0, **kwargs):
196
+ async def ainvoke_agent(self, user_instructions: str=None, max_retries=3, retry_count=0, **kwargs):
197
197
  """
198
198
  Asynchronously runs the SQL Database Agent based on user instructions.
199
199
 
@@ -212,7 +212,7 @@ class SQLDatabaseAgent(BaseAgent):
212
212
  -------
213
213
  None
214
214
  """
215
- response = self._compiled_graph.ainvoke({
215
+ response = await self._compiled_graph.ainvoke({
216
216
  "user_instructions": user_instructions,
217
217
  "max_retries": max_retries,
218
218
  "retry_count": retry_count
@@ -0,0 +1,2 @@
1
+ from ai_data_science_team.ml_agents.h2o_ml_agent import make_h2o_ml_agent, H2OMLAgent
2
+ from ai_data_science_team.ml_agents.mlflow_tools_agent import make_mlflow_tools_agent, MLflowToolsAgent