Code lenses are now shown to run @query and @predict decorated me…

…thods.

docs/changelog.md

@@ -1,28 +1,30 @@
 ## Unreleased
 
-- Add `sema4ai.dropAllDataSources` command
-- Add setup support for external data sources
-  - Show a progress bar until all models are ready
 - Add Data Sources to the Packages tree
+  - Hover (tree item) action in root element to Drop all Data Sources
   - Hover (tree item) action: remove Data Source
   - Hover (tree item) action: setup Data Source
+    - Supports external Data Sources (requires `VSCode Sema4.ai Data Extension` to be installed)
+    - Supports generated Data Sources (files, models and custom SQL)
+      - A progress bar is shown until all models are actually ready
   - Hover (tree item) action: open Data Source definition location
-  - Tooltip showing Data Source information
-- Add `sema4ai.getActionsMetadata` command
-- Add `query` and `predict` actions to UI
-- Update Action Server to `2.3.0`.
-- Extension Display Name updated to `Sema4.ai SDK`.
-- New icons for the extension.
-- Aligned the release details for the extension
-- Update Agent CLI to `0.2.2`
-- Search for actions/queries/predictions considering all glob patterns supported by `sema4ai.actions`.
-- Use `package.yaml` directory as the `cwd` when searching for actions/queries/predictions if available (otherwise imports could fail).
-- Show progress when creating the input file for an action/query/prediction.
-- Integrate with the VSCode Sema4.ai Data Extension to get the data server info to run data actions.
-- No longer use `robocorp-trustore` (it's no longer needed).
-- Use `sema4ai-http-helper` to make http requests (to data server).
-- Use `Python 3.11` as the base interpreter.
+  - Tooltip shows Data Source information
+- New command to get actions metadata: `sema4ai.getActionsMetadata`
+- Branding:
+  - A new icon is used to represent the view and the extension.
+  - The extension display name is now `Sema4.ai SDK`.
+- Search for `actions`/`queries`/`predictions` considering all glob patterns supported by `sema4ai.actions`.
+- Use `package.yaml` directory as the `cwd` when searching for `actions`/`queries`/`predictions` if available (otherwise imports could fail).
+- Progress is shown when creating the input file for an `action`/`query`/`prediction`.
+- Integrate with the `VSCode Sema4.ai Data Extension` to get the Data Server info to run data actions.
+- `robocorp-trustore` is no longer used (it's not needed anymore).
+- `sema4ai-http-helper` is now used to make http requests (to the Data Server).
+- `Python 3.11` is now used as the base interpreter.
 - When launching, if the Data Sources are not properly configured, show error message to the user.
+- It's now possible to run `@query` and `@predict` decorated methods from the tree.
+- Code lenses are now shown to run `@query` and `@predict` decorated methods.
+- Action Server updated to `2.3.0`.
+- Agent CLI updated to `0.2.2`
 
 ## New in 2.8.1 (2024-11-21)
 

sema4ai/src/sema4ai_code/robo/launch_code_lens.py

@@ -102,6 +102,30 @@ def _is_sema4ai_actions_import(node: "ast.AST") -> bool:
     return False
 
 
+def _is_sema4ai_queries_import(node: "ast.AST") -> bool:
+    import ast
+
+    if isinstance(node, ast.ImportFrom):
+        # Check if the module is 'sema4ai.data' and 'query' is in the names
+        if node.module == "sema4ai.data":
+            for alias in node.names:
+                if alias.name == "query":
+                    return True
+    return False
+
+
+def _is_sema4ai_predict_import(node: "ast.AST") -> bool:
+    import ast
+
+    if isinstance(node, ast.ImportFrom):
+        # Check if the module is 'sema4ai.data' and 'predict' is in the names
+        if node.module == "sema4ai.data":
+            for alias in node.names:
+                if alias.name == "predict":
+                    return True
+    return False
+
+
 def _create_code_lens(start_line, title, command, arguments) -> CodeLensTypedDict:
     return {
         "range": {
@@ -178,7 +202,6 @@ def _collect_python_code_lenses_in_thread(
     monitor: IMonitor,
 ) -> list[CodeLensTypedDict] | None:
     import ast
-    import os
 
     from sema4ai_code.commands import (
         SEMA4AI_ROBOTS_VIEW_ACTION_DEBUG,
@@ -197,6 +220,8 @@ def _collect_python_code_lenses_in_thread(
 
     found_robocorp_tasks_import = False
     found_sema4ai_actions_import = False
+    found_sema4ai_queries_import = False
+    found_sema4ai_predict_import = False
 
     package_yaml_path = _find_package_yaml_from_path(Path(document.path))
     if not package_yaml_path:
@@ -214,10 +239,19 @@ def _collect_python_code_lenses_in_thread(
         if not found_sema4ai_actions_import and compute_action_packages_code_lenses:
             found_sema4ai_actions_import = _is_sema4ai_actions_import(node)
 
+        # Detect if there's a `from sema4ai.data import query` import.
+        if not found_sema4ai_queries_import and compute_action_packages_code_lenses:
+            found_sema4ai_queries_import = _is_sema4ai_queries_import(node)
+
+        # Detect if there's a `from sema4ai.data import predict` import.
+        if not found_sema4ai_predict_import and compute_action_packages_code_lenses:
+            found_sema4ai_predict_import = _is_sema4ai_predict_import(node)
+
         if isinstance(node, ast.FunctionDef) and node.decorator_list:
             for decorator in node.decorator_list:
                 if (
                     compute_action_packages_code_lenses
+                    and found_sema4ai_actions_import
                     and isinstance(decorator, ast.Name)
                     and decorator.id == "action"
                 ):
@@ -251,7 +285,79 @@ def _collect_python_code_lenses_in_thread(
                         )
                     )
 
-                if (
+                elif (
+                    compute_action_packages_code_lenses
+                    and found_sema4ai_queries_import
+                    and isinstance(decorator, ast.Name)
+                    and decorator.id == "query"
+                ):
+                    assert (
+                        package_yaml_path is not None
+                    ), "Expected package_yaml_path to be defined at this point."
+                    function_name = node.name
+                    start_line = decorator.lineno - 1  # AST line numbers are 1-based
+                    robot_entry = {
+                        "actionName": function_name,
+                        "robot": {
+                            "directory": str(package_yaml_path.parent),
+                            "filePath": str(package_yaml_path),
+                        },
+                        "uri": document.uri,
+                    }
+                    actions_code_lenses.append(
+                        _create_code_lens(
+                            start_line,
+                            "Run Query",
+                            SEMA4AI_ROBOTS_VIEW_ACTION_RUN,
+                            [robot_entry],
+                        )
+                    )
+                    actions_code_lenses.append(
+                        _create_code_lens(
+                            start_line,
+                            "Debug Query",
+                            SEMA4AI_ROBOTS_VIEW_ACTION_DEBUG,
+                            [robot_entry],
+                        )
+                    )
+
+                elif (
+                    compute_action_packages_code_lenses
+                    and found_sema4ai_predict_import
+                    and isinstance(decorator, ast.Name)
+                    and decorator.id == "predict"
+                ):
+                    assert (
+                        package_yaml_path is not None
+                    ), "Expected package_yaml_path to be defined at this point."
+                    function_name = node.name
+                    start_line = decorator.lineno - 1  # AST line numbers are 1-based
+                    robot_entry = {
+                        "actionName": function_name,
+                        "robot": {
+                            "directory": str(package_yaml_path.parent),
+                            "filePath": str(package_yaml_path),
+                        },
+                        "uri": document.uri,
+                    }
+                    actions_code_lenses.append(
+                        _create_code_lens(
+                            start_line,
+                            "Run Predict",
+                            SEMA4AI_ROBOTS_VIEW_ACTION_RUN,
+                            [robot_entry],
+                        )
+                    )
+                    actions_code_lenses.append(
+                        _create_code_lens(
+                            start_line,
+                            "Debug Predict",
+                            SEMA4AI_ROBOTS_VIEW_ACTION_DEBUG,
+                            [robot_entry],
+                        )
+                    )
+
+                elif (
                     compute_robo_tasks_code_lenses
                     and isinstance(decorator, ast.Name)
                     and decorator.id == "task"
@@ -292,7 +398,5 @@ def _collect_python_code_lenses_in_thread(
         # If there was no import, then @task decorators are from some other library!
         all_lenses.extend(tasks_code_lenses)
 
-    if found_sema4ai_actions_import:
-        # If there was no import, then @action decorators are from some other library!
-        all_lenses.extend(actions_code_lenses)
+    all_lenses.extend(actions_code_lenses)
     return all_lenses

sema4ai/tests/sema4ai_code_tests/robo/test_tasks_code_lenses.py → sema4ai/tests/sema4ai_code_tests/robo/test_code_lenses.py

@@ -99,6 +99,7 @@ def _fix_code_lenses(code_lenses: list | None, root):
         "simple",
         "no_import",
         "actions",
+        "queries",
     ],
 )
 def test_call_compute_code_lenses_directly(
@@ -126,6 +127,23 @@ def test_call_compute_code_lenses_directly(
 def my_entry_point():
     pass
     """
+
+    elif scenario == "queries":
+        txt = """
+from sema4ai.data import query, predict
+
+@query
+def my_query_entry_point():
+    pass
+
+@predict
+def my_predict_entry_point():
+    pass
+    """
+
+        # We must have a package.yaml file in the root directory for actions to be searched for.
+        (root / "package.yaml").write_text("")
+
     elif scenario == "actions":
         txt = """
 from sema4ai.actions import action

sema4ai/tests/sema4ai_code_tests/robo/test_code_lenses/run_task_queries.yml

@@ -0,0 +1,68 @@
+- command:
+    arguments:
+    - actionName: my_query_entry_point
+      robot:
+        directory: root
+        filePath: package.yaml
+      uri: my.py
+    command: sema4ai.robotsViewActionRun
+    title: Run Query
+  data: null
+  range:
+    end:
+      character: 0
+      line: 3
+    start:
+      character: 0
+      line: 3
+- command:
+    arguments:
+    - actionName: my_query_entry_point
+      robot:
+        directory: root
+        filePath: package.yaml
+      uri: my.py
+    command: sema4ai.robotsViewActionDebug
+    title: Debug Query
+  data: null
+  range:
+    end:
+      character: 0
+      line: 3
+    start:
+      character: 0
+      line: 3
+- command:
+    arguments:
+    - actionName: my_predict_entry_point
+      robot:
+        directory: root
+        filePath: package.yaml
+      uri: my.py
+    command: sema4ai.robotsViewActionRun
+    title: Run Predict
+  data: null
+  range:
+    end:
+      character: 0
+      line: 7
+    start:
+      character: 0
+      line: 7
+- command:
+    arguments:
+    - actionName: my_predict_entry_point
+      robot:
+        directory: root
+        filePath: package.yaml
+      uri: my.py
+    command: sema4ai.robotsViewActionDebug
+    title: Debug Predict
+  data: null
+  range:
+    end:
+      character: 0
+      line: 7
+    start:
+      character: 0
+      line: 7