sql agent tutorial

business-science · Dec 31, 2024 · 4fa237d · 4fa237d
1 parent c8f2cd9
commit 4fa237d
Show file tree

Hide file tree

Showing 5 changed files with 712 additions and 1 deletion.
diff --git a/ai_data_science_team/agents/data_cleaning_agent.py b/ai_data_science_team/agents/data_cleaning_agent.py
@@ -201,6 +201,8 @@ def recommend_cleaning_steps(state: GraphState):
         }
 
     def create_data_cleaner_code(state: GraphState):
+        if bypass_recommended_steps:
+            print("---DATA CLEANING AGENT----")
         print("    * CREATE DATA CLEANER CODE")
 
         data_cleaning_prompt = PromptTemplate(

diff --git a/ai_data_science_team/agents/data_wrangling_agent.py b/ai_data_science_team/agents/data_wrangling_agent.py
@@ -194,6 +194,8 @@ def recommend_wrangling_steps(state: GraphState):
 
 
     def create_data_wrangler_code(state: GraphState):
+        if bypass_recommended_steps:
+            print("---DATA WRANGLING AGENT----")
         print("    * CREATE DATA WRANGLER CODE")
 
         data_wrangling_prompt = PromptTemplate(

diff --git a/ai_data_science_team/agents/feature_engineering_agent.py b/ai_data_science_team/agents/feature_engineering_agent.py
@@ -216,6 +216,8 @@ def human_review(state: GraphState) -> Command[Literal["recommend_feature_engine
         )
 
     def create_feature_engineering_code(state: GraphState):
+        if bypass_recommended_steps:
+            print("---FEATURE ENGINEERING AGENT----")
         print("    * CREATE FEATURE ENGINEERING CODE")
 
         feature_engineering_prompt = PromptTemplate(

diff --git a/ai_data_science_team/agents/sql_database_agent.py b/ai_data_science_team/agents/sql_database_agent.py
@@ -182,7 +182,8 @@ def recommend_sql_steps(state: GraphState):
         }
 
     def create_sql_query_code(state: GraphState):
-
+        if bypass_recommended_steps:
+            print("---SQL DATABASE AGENT---")
         print("    * CREATE SQL QUERY CODE")
 
         # Prompt to get the SQL code from the LLM

diff --git a/examples/sql_database_agent.ipynb b/examples/sql_database_agent.ipynb