Add more information when collecting data source information using AST.

sema4ai-python-ls-core/src/sema4ai_ls_core/protocols.py

@@ -1007,6 +1007,10 @@ class DatasourceInfoTypedDict(TypedDict):
     engine: str
     model_name: str | None
     created_table: str | None
+    setup_sql: str | list[str] | None
+    setup_sql_files: list[str] | str | None
+    description: str | None
+    file: str | None
 
 
 class DataSourceStateDict(TypedDict):

sema4ai/src/sema4ai_code/robo/collect_actions_ast.py

@@ -99,6 +99,10 @@ class _DatasourceInfo(TypedDict, total=False):
     model_name: str | None
     created_table: str | None
     python_variable_name: str
+    description: str | None
+    file: str | None
+    setup_sql: str | None
+    setup_sql_files: list[str] | None
 
 
 def _extract_datasource_info(
@@ -116,14 +120,33 @@ def _extract_datasource_info(
     """
     import typing
 
-    collect_info = ("name", "model_name", "engine", "created_table")
+    collect_info = (
+        "name",
+        "model_name",
+        "engine",
+        "description",
+        "file",
+        "created_table",
+        "setup_sql",
+        "setup_sql_files",
+    )
 
     info: _DatasourceInfo = {"kind": "datasource", "node": target_node}
 
     for keyword in call_node.keywords:
         if keyword.arg in collect_info:
             key = typing.cast(
-                Literal["name", "model_name", "engine", "created_table"], keyword.arg
+                Literal[
+                    "name",
+                    "model_name",
+                    "engine",
+                    "description",
+                    "file",
+                    "created_table",
+                    "setup_sql",
+                    "setup_sql_files",
+                ],
+                keyword.arg,
             )
             name = _resolve_value(keyword.value, variable_values)
             info[key] = name
@@ -308,6 +331,12 @@ def iter_actions_and_datasources(
                                     python_variable_name=node_info_datasource.get(
                                         "python_variable_name"
                                     ),
+                                    setup_sql=node_info_datasource.get("setup_sql"),
+                                    setup_sql_files=node_info_datasource.get(
+                                        "setup_sql_files"
+                                    ),
+                                    description=node_info_datasource.get("description"),
+                                    file=node_info_datasource.get("file"),
                                 )
                 except Exception as e:
                     log.error(

sema4ai/tests/sema4ai_code_tests/robo/test_list_actions/test_list_actions.yml

@@ -1,5 +1,7 @@
 - created_table: churn
+  description: Ddts.
   engine: files
+  file: files/customer-churn.csv
   kind: datasource
   model_name: null
   name: files
@@ -11,6 +13,8 @@
     start:
       character: 0
       line: 5
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 - kind: query
   name: get_churn_data
@@ -23,7 +27,9 @@
       line: 8
   uri: data_actions.py
 - created_table: null
+  description: Datasource which provides along with a table named `customer_churn_predictor`.
   engine: prediction:lightwood
+  file: null
   kind: datasource
   model_name: customer_churn_predictor
   name: models
@@ -35,6 +41,10 @@
     start:
       character: 0
       line: 15
+  setup_sql: "\n            CREATE MODEL IF NOT EXISTS models.customer_churn_predictor\n\
+    \            FROM files\n            (SELECT * FROM churn)\n            PREDICT\
+    \ Churn;\n    "
+  setup_sql_files: null
   uri: data_sources.py
 - kind: action
   name: my_action

sema4ai/tests/sema4ai_code_tests/robo/test_list_actions/test_list_actions_and_datasources_mutiple_data_sources_.yml

@@ -1,5 +1,8 @@
 - created_table: churn
+  description: Datasource which provides a table named 'churn' with customer churn
+    data.
   engine: files
+  file: files/customer-churn.csv
   kind: datasource
   model_name: null
   name: files
@@ -11,9 +14,13 @@
     start:
       character: 0
       line: 5
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Datasource which provides along with a table named `customer_churn_predictor`.
   engine: prediction:lightwood
+  file: null
   kind: datasource
   model_name: customer_churn_predictor
   name: models
@@ -25,4 +32,7 @@
     start:
       character: 0
       line: 15
+  setup_sql: CREATE MODEL IF NOT EXISTS customer_churn_predictor FROM files (SELECT
+    * FROM churn) PREDICT Churn;
+  setup_sql_files: null
   uri: data_sources.py

sema4ai/tests/sema4ai_code_tests/robo/test_list_actions/test_list_actions_and_datasources_mutiple_data_sources_2_.yml

@@ -1,5 +1,8 @@
 - created_table: churn
+  description: Datasource which provides a table named 'churn' with customer churn
+    data.
   engine: files
+  file: files/customer-churn.csv
   kind: datasource
   model_name: null
   name: files
@@ -11,9 +14,13 @@
     start:
       character: 0
       line: 5
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Datasource which provides along with a table named `customer_churn_predictor`.
   engine: prediction:lightwood
+  file: null
   kind: datasource
   model_name: customer_churn_predictor
   name: models
@@ -25,4 +32,7 @@
     start:
       character: 0
       line: 15
+  setup_sql: CREATE MODEL IF NOT EXISTS customer_churn_predictor FROM files (SELECT
+    * FROM churn) PREDICT Churn;
+  setup_sql_files: null
   uri: data_sources.py

sema4ai/tests/sema4ai_code_tests/robo/test_list_actions/test_list_actions_and_datasources_simple.yml

@@ -1,5 +1,7 @@
 - created_table: churn
+  description: Ddts.
   engine: files
+  file: files/customer-churn.csv
   kind: datasource
   model_name: null
   name: files
@@ -11,6 +13,8 @@
     start:
       character: 0
       line: 5
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 - kind: query
   name: get_churn_data
@@ -23,7 +27,9 @@
       line: 8
   uri: data_actions.py
 - created_table: null
+  description: Datasource which provides along with a table named `customer_churn_predictor`.
   engine: prediction:lightwood
+  file: null
   kind: datasource
   model_name: customer_churn_predictor
   name: models
@@ -35,6 +41,10 @@
     start:
       character: 0
       line: 15
+  setup_sql: "\n            CREATE MODEL IF NOT EXISTS models.customer_churn_predictor\n\
+    \            FROM files\n            (SELECT * FROM churn)\n            PREDICT\
+    \ Churn;\n    "
+  setup_sql_files: null
   uri: data_sources.py
 - kind: action
   name: my_action

sema4ai/tests/sema4ai_code_tests/test_compute_data_sources_state/missing_data_source_all.yml

@@ -1,7 +1,9 @@
 data_sources_in_data_server: []
 required_data_sources:
 - created_table: customers_in_test_compute_data_sources_state
+  description: Data source for customers.
   engine: files
+  file: files/customers.csv
   kind: datasource
   model_name: null
   name: files
@@ -13,9 +15,13 @@ required_data_sources:
     start:
       character: 0
       line: 6
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Predict something.
   engine: prediction:lightwood
+  file: null
   kind: datasource
   model_name: predict_compute_data_sources_state
   name: models
@@ -27,9 +33,25 @@ required_data_sources:
     start:
       character: 0
       line: 17
+  setup_sql: '
+
+    CREATE MODEL models.predict_compute_data_sources_state
+
+    (SELECT * FROM files.customers)
+
+    PREDICT country
+
+    ORDER BY first_name
+
+    WINDOW 8
+
+    HORIZON 4;'
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Data source for tests.
   engine: sqlite
+  file: null
   kind: datasource
   model_name: null
   name: test_compute_data_sources_state
@@ -41,10 +63,14 @@ required_data_sources:
     start:
       character: 0
       line: 34
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 unconfigured_data_sources:
 - created_table: customers_in_test_compute_data_sources_state
+  description: Data source for customers.
   engine: files
+  file: files/customers.csv
   kind: datasource
   model_name: null
   name: files
@@ -56,9 +82,13 @@ unconfigured_data_sources:
     start:
       character: 0
       line: 6
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Predict something.
   engine: prediction:lightwood
+  file: null
   kind: datasource
   model_name: predict_compute_data_sources_state
   name: models
@@ -70,9 +100,25 @@ unconfigured_data_sources:
     start:
       character: 0
       line: 17
+  setup_sql: '
+
+    CREATE MODEL models.predict_compute_data_sources_state
+
+    (SELECT * FROM files.customers)
+
+    PREDICT country
+
+    ORDER BY first_name
+
+    WINDOW 8
+
+    HORIZON 4;'
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Data source for tests.
   engine: sqlite
+  file: null
   kind: datasource
   model_name: null
   name: test_compute_data_sources_state
@@ -84,5 +130,7 @@ unconfigured_data_sources:
     start:
       character: 0
       line: 34
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 uri_to_error_messages: {}

sema4ai/tests/sema4ai_code_tests/test_compute_data_sources_state/missing_data_source_none.yml

@@ -1,7 +1,9 @@
 data_sources_in_data_server: []
 required_data_sources:
 - created_table: customers_in_test_compute_data_sources_state
+  description: Data source for customers.
   engine: files
+  file: files/customers.csv
   kind: datasource
   model_name: null
   name: files
@@ -13,9 +15,13 @@ required_data_sources:
     start:
       character: 0
       line: 6
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Predict something.
   engine: prediction:lightwood
+  file: null
   kind: datasource
   model_name: predict_compute_data_sources_state
   name: models
@@ -27,9 +33,25 @@ required_data_sources:
     start:
       character: 0
       line: 17
+  setup_sql: '
+
+    CREATE MODEL models.predict_compute_data_sources_state
+
+    (SELECT * FROM files.customers)
+
+    PREDICT country
+
+    ORDER BY first_name
+
+    WINDOW 8
+
+    HORIZON 4;'
+  setup_sql_files: null
   uri: data_sources.py
 - created_table: null
+  description: Data source for tests.
   engine: sqlite
+  file: null
   kind: datasource
   model_name: null
   name: test_compute_data_sources_state
@@ -41,6 +63,8 @@ required_data_sources:
     start:
       character: 0
       line: 34
+  setup_sql: null
+  setup_sql_files: null
   uri: data_sources.py
 unconfigured_data_sources: []
 uri_to_error_messages: {}