tinybird 0.0.1.dev15__py3-none-any.whl → 0.0.1.dev17__py3-none-any.whl

tinybird/tb/modules/cli.py

@@ -36,7 +36,6 @@ from tinybird.tb.modules.common import (
     coro,
     create_tb_client,
     echo_safe_format_table,
-    folder_init,
     get_current_main_workspace,
     getenv_bool,
     is_major_semver,
@@ -63,6 +62,7 @@ from tinybird.tb.modules.datafile.exceptions import (
 from tinybird.tb.modules.datafile.parse_datasource import parse_datasource
 from tinybird.tb.modules.datafile.parse_pipe import parse_pipe
 from tinybird.tb.modules.datafile.pull import folder_pull
+from tinybird.tb.modules.local_common import get_tinybird_local_client
 from tinybird.tb.modules.telemetry import add_telemetry_event
 
 __old_click_echo = click.echo
@@ -80,7 +80,6 @@ DEFAULT_PATTERNS: List[Tuple[str, Union[str, Callable[[str], str]]]] = [
 )
 @click.option("--token", help="Use auth token, defaults to TB_TOKEN envvar, then to the .tinyb file")
 @click.option("--host", help="Use custom host, defaults to TB_HOST envvar, then to https://api.tinybird.co")
-@click.option("--semver", help="Semver of a Release to run the command. Example: 1.0.0", hidden=True)
 @click.option("--gcp-project-id", help="The Google Cloud project ID", hidden=True)
 @click.option(
     "--gcs-bucket", help="The Google Cloud Storage bucket to write temp files when using the connectors", hidden=True
@@ -108,6 +107,7 @@ DEFAULT_PATTERNS: List[Tuple[str, Union[str, Callable[[str], str]]]] = [
     "--with-headers", help="Flag to enable connector to export with headers", is_flag=True, default=False, hidden=True
 )
 @click.option("--show-tokens", is_flag=True, default=False, help="Enable the output of tokens")
+@click.option("--local", is_flag=True, default=False, help="Run in local mode")
 @click.version_option(version=VERSION)
 @click.pass_context
 @coro
@@ -116,7 +116,6 @@ async def cli(
     debug: bool,
     token: str,
     host: str,
-    semver: str,
     gcp_project_id: str,
     gcs_bucket: str,
     google_application_credentials: str,
@@ -131,13 +130,14 @@ async def cli(
     sf_stage,
     with_headers: bool,
     show_tokens: bool,
+    local: bool,
 ) -> None:
     """
     Use `OBFUSCATE_REGEX_PATTERN` and `OBFUSCATE_PATTERN_SEPARATOR` environment variables to define a regex pattern and a separator (in case of a single string with multiple regex) to obfuscate secrets in the CLI output.
     """
 
     # We need to unpatch for our tests not to break
-    if show_tokens:
+    if show_tokens or local or ctx.invoked_subcommand == "build":
         __unpatch_click_output()
     else:
         __patch_click_output()
@@ -153,9 +153,7 @@ async def cli(
         config_temp.set_token(token)
     if host:
         config_temp.set_host(host)
-    if semver:
-        config_temp.set_semver(semver)
-    if token or host or semver:
+    if token or host:
         await try_update_config_with_remote(config_temp, auto_persist=False, raise_on_errors=False)
 
     # Overwrite token and host with env vars manually, without resorting to click.
@@ -167,10 +165,8 @@ async def cli(
         token = os.environ.get("TB_TOKEN", "")
     if not host and "TB_HOST" in os.environ:
         host = os.environ.get("TB_HOST", "")
-    if not semver and "TB_SEMVER" in os.environ:
-        semver = os.environ.get("TB_SEMVER", "")
 
-    config = await get_config(host, token, semver)
+    config = await get_config(host, token)
     client = _get_tb_client(config.get("token", None), config["host"])
 
     # If they have passed a token or host as paramter and it's different that record in .tinyb, refresh the workspace id
@@ -228,85 +224,14 @@ async def cli(
 
     logging.debug("debug enabled")
 
-    ctx.ensure_object(dict)["client"] = _get_tb_client(config.get("token", None), config["host"], semver)
+    ctx.ensure_object(dict)["client"] = (
+        await get_tinybird_local_client() if local else _get_tb_client(config.get("token", None), config["host"])
+    )
 
     for connector in SUPPORTED_CONNECTORS:
         load_connector_config(ctx, connector, debug, check_uninstalled=True)
 
 
-@cli.command()
-@click.option(
-    "--generate-datasources",
-    is_flag=True,
-    default=False,
-    help="Generate datasources based on CSV, NDJSON and Parquet files in this folder",
-)
-@click.option(
-    "--folder",
-    default=None,
-    type=click.Path(exists=True, file_okay=False),
-    help="Folder where datafiles will be placed",
-)
-@click.option("-f", "--force", is_flag=True, default=False, help="Overrides existing files")
-@click.option(
-    "-ir",
-    "--ignore-remote",
-    is_flag=True,
-    default=False,
-    help="Ignores remote files not present in the local data project on git init",
-)
-@click.option(
-    "--git",
-    is_flag=True,
-    default=False,
-    help="Init workspace with git releases. Generates CI/CD files for your git provider",
-)
-@click.option(
-    "--override-commit",
-    default=None,
-    help="Use this option to manually override the reference commit of your workspace. This is useful if a commit is not recognized in your git log, such as after a force push (git push -f).",
-)
-@click.option(
-    "--cicd", is_flag=True, default=False, help="Generates only CI/CD files for your git provider", hidden=True
-)
-@click.pass_context
-@coro
-async def init(
-    ctx: Context,
-    generate_datasources: bool,
-    folder: Optional[str],
-    force: bool,
-    ignore_remote: bool,
-    git: bool,
-    override_commit: Optional[str],
-    cicd: Optional[bool],
-) -> None:
-    """Initialize folder layout."""
-    client: TinyB = ctx.ensure_object(dict)["client"]
-    config = CLIConfig.get_project_config()
-    if config.get("token") is None:
-        raise AuthNoTokenException
-    folder = folder if folder else getcwd()
-
-    workspaces: List[Dict[str, Any]] = (await client.user_workspaces_and_branches()).get("workspaces", [])
-    current_ws: Dict[str, Any] = next(
-        (workspace for workspace in workspaces if config and workspace.get("id", ".") == config.get("id", "..")), {}
-    )
-
-    if current_ws.get("is_branch"):
-        raise CLIException(FeedbackManager.error_not_allowed_in_branch())
-
-    await folder_init(client, folder, generate_datasources, generate_releases=True, force=force)
-
-    error = False
-    final_response = None
-
-    if final_response:
-        if error:
-            raise CLIException(final_response)
-        click.echo(final_response)
-
-
 @cli.command()
 @click.argument("filenames", type=click.Path(exists=True), nargs=-1, default=None)
 @click.option("--debug", is_flag=True, default=False, help="Print internal representation")
@@ -1340,9 +1265,6 @@ async def deploy(
             check_backfill_required = getenv_bool("TB_CHECK_BACKFILL_REQUIRED", True)
         try:
             tb_client = create_tb_client(ctx)
-            if dry_run:
-                config.set_semver(None)
-                tb_client.semver = None
             await folder_push(
                 tb_client=tb_client,
                 dry_run=dry_run,

tinybird/tb/modules/common.py

@@ -218,19 +218,6 @@ async def get_current_environment(client, config):
     return next((workspace for workspace in workspaces if workspace["id"] == config["id"]), None)
 
 
-async def get_current_workspace_branches(config: CLIConfig) -> List[Dict[str, Any]]:
-    current_main_workspace: Optional[Dict[str, Any]] = await get_current_main_workspace(config)
-    if not current_main_workspace:
-        raise CLIException(FeedbackManager.error_unable_to_identify_main_workspace())
-
-    client = config.get_client()
-    user_branches: List[Dict[str, Any]] = (await client.user_workspace_branches()).get("workspaces", [])
-    all_branches: List[Dict[str, Any]] = (await client.branches()).get("environments", [])
-    branches = user_branches + [branch for branch in all_branches if branch not in user_branches]
-
-    return [branch for branch in branches if branch.get("main") == current_main_workspace["id"]]
-
-
 class AliasedGroup(click.Group):
     def get_command(self, ctx, cmd_name):
         # Step one: built-in commands as normal
@@ -320,16 +307,15 @@ def getenv_bool(key: str, default: bool) -> bool:
     return v.lower() == "true" or v == "1"
 
 
-def _get_tb_client(token: str, host: str, semver: Optional[str] = None) -> TinyB:
+def _get_tb_client(token: str, host: str) -> TinyB:
     disable_ssl: bool = getenv_bool("TB_DISABLE_SSL_CHECKS", False)
-    return TinyB(token, host, version=VERSION, disable_ssl_checks=disable_ssl, send_telemetry=True, semver=semver)
+    return TinyB(token, host, version=VERSION, disable_ssl_checks=disable_ssl, send_telemetry=True)
 
 
 def create_tb_client(ctx: Context) -> TinyB:
     token = ctx.ensure_object(dict)["config"].get("token", "")
     host = ctx.ensure_object(dict)["config"].get("host", DEFAULT_API_HOST)
-    semver = ctx.ensure_object(dict)["config"].get("semver", "")
-    return _get_tb_client(token, host, semver=semver)
+    return _get_tb_client(token, host)
 
 
 async def _analyze(filename: str, client: TinyB, format: str, connector: Optional["Connector"] = None):
@@ -776,77 +762,6 @@ async def create_workspace_interactive(
     await create_workspace_non_interactive(ctx, workspace_name, starterkit, user_token, fork)  # type: ignore
 
 
-async def create_workspace_branch(
-    branch_name: Optional[str],
-    last_partition: bool,
-    all: bool,
-    ignore_datasources: Optional[List[str]],
-    wait: Optional[bool],
-) -> None:
-    """
-    Creates a workspace branch
-    """
-    config = CLIConfig.get_project_config()
-    _ = await try_update_config_with_remote(config)
-
-    try:
-        workspace = await get_current_workspace(config)
-        if not workspace:
-            raise CLIWorkspaceException(FeedbackManager.error_workspace())
-
-        if not branch_name:
-            click.echo(FeedbackManager.info_workspace_branch_create_greeting())
-            default_name = f"{workspace['name']}_{uuid.uuid4().hex[0:4]}"
-            branch_name = click.prompt("\Branch name", default=default_name, err=True, type=str)
-        assert isinstance(branch_name, str)
-
-        response = await config.get_client().create_workspace_branch(
-            branch_name,
-            last_partition,
-            all,
-            ignore_datasources,
-        )
-        assert isinstance(response, dict)
-
-        is_job: bool = "job" in response
-        is_summary: bool = "partitions" in response
-
-        if not is_job and not is_summary:
-            raise CLIException(str(response))
-
-        if all and not is_job:
-            raise CLIException(str(response))
-
-        click.echo(
-            FeedbackManager.success_workspace_branch_created(workspace_name=workspace["name"], branch_name=branch_name)
-        )
-
-        job_id: Optional[str] = None
-
-        if is_job:
-            job_id = response["job"]["job_id"]
-            job_url = response["job"]["job_url"]
-            click.echo(FeedbackManager.info_data_branch_job_url(url=job_url))
-
-        if wait and is_job:
-            assert isinstance(job_id, str)
-
-            # Await the job to finish and get the result dict
-            job_response = await wait_job(config.get_client(), job_id, job_url, "Branch creation")
-            if job_response is None:
-                raise CLIException(f"Empty job API response (job_id: {job_id}, job_url: {job_url})")
-            else:
-                response = job_response.get("result", {})
-                is_summary = "partitions" in response
-
-        await switch_workspace(config, branch_name, only_environments=True)
-        if is_summary and (bool(last_partition) or bool(all)):
-            await print_data_branch_summary(config.get_client(), None, response)
-
-    except Exception as e:
-        raise CLIException(FeedbackManager.error_exception(error=str(e)))
-
-
 async def print_data_branch_summary(client, job_id, response=None):
     response = await client.job(job_id) if job_id else response or {"partitions": []}
     columns = ["Data Source", "Partition", "Status", "Error"]
@@ -1327,13 +1242,10 @@ def _get_setting_value(connection, setting, sensitive_settings):
     return connection.get(setting, "")
 
 
-async def switch_workspace(config: CLIConfig, workspace_name_or_id: str, only_environments: bool = False) -> None:
+async def switch_workspace(config: CLIConfig, workspace_name_or_id: str) -> None:
     try:
-        if only_environments:
-            workspaces = await get_current_workspace_branches(config)
-        else:
-            response = await config.get_client().user_workspaces()
-            workspaces = response["workspaces"]
+        response = await config.get_client().user_workspaces()
+        workspaces = response["workspaces"]
 
         workspace = next(
             (
@@ -1345,10 +1257,7 @@ async def switch_workspace(config: CLIConfig, workspace_name_or_id: str, only_en
         )
 
         if not workspace:
-            if only_environments:
-                raise CLIException(FeedbackManager.error_branch(branch=workspace_name_or_id))
-            else:
-                raise CLIException(FeedbackManager.error_workspace(workspace=workspace_name_or_id))
+            raise CLIException(FeedbackManager.error_workspace(workspace=workspace_name_or_id))
 
         config.set_token(workspace["token"])
         config.set_token_for_host(workspace["token"], config.get_host())

tinybird/tb/modules/config.py

@@ -72,7 +72,6 @@ class CLIConfig:
         "token": "TB_TOKEN",
         "user_token": "TB_USER_TOKEN",
         "host": "TB_HOST",
-        "semver": "TB_SEMVER",
     }
 
     DEFAULTS: Dict[str, str] = {"host": DEFAULT_API_HOST if not FeatureFlags.is_localhost() else DEFAULT_LOCALHOST}
@@ -178,15 +177,6 @@ class CLIConfig:
         except KeyError:
             return None
 
-    def set_semver(self, semver: Optional[str]) -> None:
-        self["semver"] = semver
-
-    def get_semver(self) -> Optional[str]:
-        try:
-            return self["semver"]
-        except KeyError:
-            return None
-
     def set_token_for_host(self, token: Optional[str], host: Optional[str]) -> None:
         """Sets the token for the specified host.
 

tinybird/tb/modules/create.py

@@ -5,13 +5,13 @@ from pathlib import Path
 from typing import Optional
 
 import click
-from click import Context
+import requests
 
 from tinybird.client import TinyB
 from tinybird.feedback_manager import FeedbackManager
 from tinybird.tb.modules.cicd import init_cicd
 from tinybird.tb.modules.cli import cli
-from tinybird.tb.modules.common import _generate_datafile, coro, generate_datafile
+from tinybird.tb.modules.common import _generate_datafile, check_user_token, coro, generate_datafile
 from tinybird.tb.modules.config import CLIConfig
 from tinybird.tb.modules.datafile.fixture import build_fixture_name, persist_fixture
 from tinybird.tb.modules.exceptions import CLIException
@@ -19,6 +19,11 @@ from tinybird.tb.modules.llm import LLM
 
 
 @cli.command()
+@click.option(
+    "--demo",
+    is_flag=True,
+    help="Demo data and files to get started",
+)
 @click.option(
     "--data",
     type=click.Path(exists=True),
@@ -41,7 +46,8 @@ from tinybird.tb.modules.llm import LLM
 @click.pass_context
 @coro
 async def create(
-    ctx: Context,
+    ctx: click.Context,
+    demo: bool,
     data: Optional[str],
     prompt: Optional[str],
     folder: Optional[str],
@@ -50,7 +56,16 @@ async def create(
     """Initialize a new project."""
     folder = folder or getcwd()
     try:
-        config = CLIConfig.get_project_config()
+        config = CLIConfig.get_project_config(folder)
+
+        if prompt:
+            user_token = config.get_user_token()
+            try:
+                await check_user_token(ctx, token=user_token)
+            except Exception:
+                click.echo(FeedbackManager.error(message="This action requires authentication. Run 'tb login' first."))
+                return
+
         tb_client = config.get_client()
         click.echo(FeedbackManager.gray(message="Creating new project structure..."))
         await project_create(tb_client, data, prompt, folder)
@@ -63,7 +78,60 @@ async def create(
 
         click.echo(FeedbackManager.gray(message="Building fixtures..."))
 
-        if data:
+        if demo:
+            # Users datasource
+            ds_name = "users"
+            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
+            datasource_content = fetch_gist_content(
+                "https://gist.githubusercontent.com/gnzjgo/b48fb9c92825ed27c04e3104b9e871e1/raw/1f33c20eefbabc4903f38e234329e028d8ef9def/users.datasource"
+            )
+            datasource_path.write_text(datasource_content)
+            click.echo(FeedbackManager.info(message=f"✓ /datasources/{ds_name}.datasource"))
+
+            # Users fixtures
+            fixture_content = fetch_gist_content(
+                "https://gist.githubusercontent.com/gnzjgo/8e8f66a39d7576ce3a2529bf773334a8/raw/9cab636767990e97d44a141867e5f226e992de8c/users.ndjson"
+            )
+            fixture_name = build_fixture_name(datasource_path.absolute(), ds_name, datasource_path.read_text())
+            persist_fixture(fixture_name, fixture_content)
+            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
+
+            # Events datasource
+            ds_name = "events"
+            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
+            datasource_content = fetch_gist_content(
+                "https://gist.githubusercontent.com/gnzjgo/f8ca37b5b1f6707c75206b618de26bc9/raw/cd625da0dcd1ba8de29f12bc1c8600b9ff7c809c/events.datasource"
+            )
+            datasource_path.write_text(datasource_content)
+            click.echo(FeedbackManager.info(message=f"✓ /datasources/{ds_name}.datasource"))
+
+            # Events fixtures
+            fixture_content = fetch_gist_content(
+                "https://gist.githubusercontent.com/gnzjgo/859ab9439c17e77241d0c14a5a532809/raw/251f2f3f00a968f8759ec4068cebde915256b054/events.ndjson"
+            )
+            fixture_name = build_fixture_name(datasource_path.absolute(), ds_name, datasource_path.read_text())
+            persist_fixture(fixture_name, fixture_content)
+            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
+
+            # Create sample endpoint
+            pipe_name = "api_token_usage"
+            pipe_path = Path(folder) / "endpoints" / f"{pipe_name}.pipe"
+            pipe_content = fetch_gist_content(
+                "https://gist.githubusercontent.com/gnzjgo/68ecc47472c2b754b0ae0c1187022963/raw/52cc3aa3afdf939e58d43355bfe4ddc739989ddd/api_token_usage.pipe"
+            )
+            pipe_path.write_text(pipe_content)
+            click.echo(FeedbackManager.info(message=f"✓ /endpoints/{pipe_name}.pipe"))
+
+            # Create sample test
+            test_name = "api_token_usage"
+            test_path = Path(folder) / "tests" / f"{test_name}.yaml"
+            test_content = fetch_gist_content(
+                "https://gist.githubusercontent.com/gnzjgo/e58620bbb977d6f42f1d0c2a7b46ac8f/raw/a5f61d5019111f937484f941111829dfce69f648/api_token_usage.yaml"
+            )
+            test_path.write_text(test_content)
+            click.echo(FeedbackManager.info(message=f"✓ /tests/{test_name}.yaml"))
+
+        elif data:
             ds_name = os.path.basename(data.split(".")[0])
             data_content = Path(data).read_text()
             datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
@@ -79,7 +147,7 @@ async def create(
                 datasource_content = datasource_path.read_text()
                 has_json_path = "`json:" in datasource_content
                 if has_json_path:
-                    sql = await llm.generate_sql_sample_data(schema=datasource_content, rows=rows)
+                    sql = await llm.generate_sql_sample_data(schema=datasource_content, rows=rows, context=prompt)
                     result = await tb_client.query(f"{sql} FORMAT JSON")
                     data = result.get("data", [])
                     fixture_name = build_fixture_name(datasource_path.absolute(), datasource_name, datasource_content)
@@ -97,7 +165,7 @@ async def project_create(
     prompt: Optional[str],
     folder: str,
 ):
-    project_paths = ["datasources", "endpoints", "materializations", "copies", "sinks", "fixtures"]
+    project_paths = ["datasources", "endpoints", "materializations", "copies", "sinks", "fixtures", "tests"]
     force = True
     for x in project_paths:
         try:
@@ -113,7 +181,7 @@ async def project_create(
         try:
             await _generate_datafile(str(path), client, format=format, force=force)
         except Exception as e:
-            click.echo(FeedbackManager.error(message=f"Ersssssror: {str(e)}"))
+            click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
         name = data.split(".")[0]
         generate_pipe_file(
             f"{name}_endpoint",
@@ -170,3 +238,9 @@ def generate_pipe_file(name: str, content: str, folder: str):
     with open(f"{f}", "w") as file:
         file.write(content)
     click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder)))
+
+
+def fetch_gist_content(url: str) -> str:  # TODO: replace this with a function that fetches the content from a repo
+    response = requests.get(url)
+    response.raise_for_status()
+    return response.text

tinybird/tb/modules/datafile/build_pipe.py

@@ -266,7 +266,7 @@ async def new_pipe(
 
     if data.get("type") == "endpoint":
         token = tb_client.token
-        print(f"""** => Test endpoint with:\n** $ curl {host}/v0/pipes/{p["name"]}.json?token={token}""")  # noqa: T201
+        click.echo(f"""** => Test endpoint with:\n** $ curl {host}/v0/pipes/{p["name"]}.json?token={token}""")
 
 
 async def get_token_from_main_branch(branch_tb_client: TinyB) -> Optional[str]:

tinybird/tb/modules/datasource.py

@@ -21,7 +21,6 @@ if TYPE_CHECKING:
     from tinybird.connectors import Connector
 
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb.modules.branch import warn_if_in_live
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import (
     _analyze,
@@ -358,9 +357,6 @@ async def datasource_delete(ctx: Context, datasource_name: str, yes: bool):
                 FeedbackManager.error_datasource_can_not_be_deleted(datasource=datasource_name, error=e)
             )
 
-    semver: str = ctx.ensure_object(dict)["config"]["semver"]
-    await warn_if_in_live(semver)
-
     if yes or click.confirm(
         FeedbackManager.warning_confirm_delete_datasource(
             warning_message=warning_message, dependencies_information=dependencies_information
@@ -391,9 +387,6 @@ async def datasource_delete(ctx: Context, datasource_name: str, yes: bool):
 async def datasource_truncate(ctx, datasource_name, yes, cascade):
     """Truncate a data source"""
 
-    semver: str = ctx.ensure_object(dict)["config"]["semver"]
-    await warn_if_in_live(semver)
-
     client = ctx.obj["client"]
     if yes or click.confirm(FeedbackManager.warning_confirm_truncate_datasource(datasource=datasource_name)):
         try:
@@ -453,9 +446,6 @@ async def datasource_delete_rows(ctx, datasource_name, sql_condition, yes, wait,
     - Delete rows with SQL condition and wait for the job to finish: `tb datasource delete [datasource_name] --sql-condition "country='ES'" --wait`
     """
 
-    semver: str = ctx.ensure_object(dict)["config"]["semver"]
-    await warn_if_in_live(semver)
-
     client: TinyB = ctx.ensure_object(dict)["client"]
     if (
         dry_run

tinybird/tb/modules/llm.py

@@ -1,11 +1,14 @@
 import asyncio
+import json
 import urllib.parse
 from copy import deepcopy
-from typing import Awaitable, Callable, List
+from typing import Awaitable, Callable, List, Optional
 
+from openai import OpenAI
 from pydantic import BaseModel
 
 from tinybird.client import TinyB
+from tinybird.prompts import create_test_calls_prompt
 from tinybird.tb.modules.config import CLIConfig
 
 
@@ -19,13 +22,25 @@ class DataProject(BaseModel):
     pipes: List[DataFile]
 
 
+class TestExpectation(BaseModel):
+    name: str
+    description: str
+    parameters: str
+    expected_result: str
+
+
+class TestExpectations(BaseModel):
+    tests: List[TestExpectation]
+
+
 class LLM:
-    def __init__(self, client: TinyB):
+    def __init__(self, client: TinyB, api_key: Optional[str] = None):
         self.client = client
         user_token = CLIConfig.get_project_config().get_user_token()
         user_client = deepcopy(client)
         user_client.token = user_token
         self.user_client = user_client
+        self.openai = OpenAI(api_key=api_key) if api_key else None
 
     async def _execute(self, action_fn: Callable[[], Awaitable[str]], checker_fn: Callable[[str], bool]):
         is_valid = False
@@ -46,9 +61,10 @@ class LLM:
             response = await self.user_client._req(
                 "/v0/llm/create",
                 method="POST",
-                data=f'{{"prompt": "{prompt}"}}',
+                data=f'{{"prompt": {json.dumps(prompt)}}}',
                 headers={"Content-Type": "application/json"},
             )
+
             return DataProject.model_validate(response.get("result", {}))
         except Exception:
             return DataProject(datasources=[], pipes=[])
@@ -61,3 +77,19 @@ class LLM:
             headers={"Content-Type": "application/json"},
         )
         return response.get("result", "")
+
+    async def create_test_commands(
+        self, pipe_content: str, pipe_params: set[str], context: str = ""
+    ) -> TestExpectations:
+        if not self.openai:
+            raise ValueError("OpenAI API key is not set")
+
+        completion = self.openai.beta.chat.completions.parse(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": create_test_calls_prompt.format(context=context)},
+                {"role": "user", "content": f"Pipe content: {pipe_content}\nPipe params: {pipe_params}"},
+            ],
+            response_format=TestExpectations,
+        )
+        return completion.choices[0].message.parsed or TestExpectations(tests=[])