tinybird 0.0.1.dev301__py3-none-any.whl → 0.0.1.dev302__py3-none-any.whl

tinybird/tb/__cli__.py

@@ -4,5 +4,5 @@ __description__ = 'Tinybird Command Line Tool'
 __url__ = 'https://www.tinybird.co/docs/forward/commands'
 __author__ = 'Tinybird'
 __author_email__ = 'support@tinybird.co'
-__version__ = '0.0.1.dev301'
-__revision__ = '1ce40fd'
+__version__ = '0.0.1.dev302'
+__revision__ = 'cf4d28e'

tinybird/tb/modules/agent/agent.py

@@ -6,7 +6,6 @@ import sys
 import urllib.parse
 import uuid
 from functools import partial
-from pathlib import Path
 from typing import Any, Callable, Optional
 
 import click
@@ -27,7 +26,6 @@ from tinybird.tb.modules.agent.file_agent import FileAgent
 from tinybird.tb.modules.agent.memory import (
     clear_history,
     clear_messages,
-    get_last_messages_from_last_user_prompt,
     save_messages,
 )
 from tinybird.tb.modules.agent.mock_agent import MockAgent
@@ -72,12 +70,12 @@ from tinybird.tb.modules.common import (
 )
 from tinybird.tb.modules.config import CLIConfig
 from tinybird.tb.modules.deployment_common import create_deployment
-from tinybird.tb.modules.exceptions import CLIBuildException, CLIDeploymentException, CLIMockException
+from tinybird.tb.modules.exceptions import CLIBuildException, CLIDeploymentException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.llm import LLM
 from tinybird.tb.modules.local_common import get_tinybird_local_client
 from tinybird.tb.modules.login_common import login
-from tinybird.tb.modules.mock_common import append_mock_data, create_mock_data
+from tinybird.tb.modules.mock_common import append_mock_data
 from tinybird.tb.modules.project import Project
 from tinybird.tb.modules.test_common import run_tests as run_tests_common
 
@@ -92,6 +90,7 @@ class TinybirdAgent:
         project: Project,
         dangerously_skip_permissions: bool,
         prompt_mode: bool,
+        feature: Optional[str] = None,
     ):
         self.token = token
         self.user_token = user_token
@@ -101,14 +100,12 @@ class TinybirdAgent:
         self.project = project
         self.thinking_animation = ThinkingAnimation()
         self.confirmed_plan_id: Optional[str] = None
-        if prompt_mode:
-            self.messages: list[ModelMessage] = get_last_messages_from_last_user_prompt()
-        else:
-            self.messages = []
+        self.feature = feature
+        self.messages: list[ModelMessage] = []
         cli_config = CLIConfig.get_project_config()
         regions = get_regions(cli_config)
         self.agent = Agent(
-            model=create_model(user_token, host, workspace_id),
+            model=create_model(user_token, host, workspace_id, feature=feature),
             deps_type=TinybirdAgentContext,
             instructions=[agent_system_prompt],
             tools=[
@@ -382,7 +379,6 @@ class TinybirdAgent:
             build_project=partial(build_project, project=project, config=config),
             deploy_project=partial(deploy_project, project=project, config=config),
             deploy_check_project=partial(deploy_check_project, project=project, config=config),
-            mock_data=partial(mock_data, project=project, config=config),
             append_data_local=partial(append_data_local, config=config),
             append_data_cloud=partial(append_data_cloud, config=config),
             analyze_fixture=partial(analyze_fixture, config=config),
@@ -511,7 +507,11 @@ class TinybirdAgent:
 
 
 def run_agent(
-    config: dict[str, Any], project: Project, dangerously_skip_permissions: bool, prompt: Optional[str] = None
+    config: dict[str, Any],
+    project: Project,
+    dangerously_skip_permissions: bool,
+    prompt: Optional[str] = None,
+    feature: Optional[str] = None,
 ):
     if not prompt:
         latest_version = CheckPypi().get_latest_version()
@@ -526,7 +526,10 @@ def run_agent(
             )
             if yes:
                 update_cli()
-    click.echo(FeedbackManager.highlight(message="» Initializing Tinybird Code..."))
+
+    if not prompt:
+        click.echo(FeedbackManager.highlight(message="» Initializing Tinybird Code..."))
+
     token = config.get("token", "")
     host = config.get("host", "")
     user_token = config.get("user_token", "")
@@ -585,6 +588,7 @@ def run_agent(
             project,
             dangerously_skip_permissions,
             prompt_mode,
+            feature,
         )
 
         # Print mode: run once with the provided prompt and exit
@@ -763,36 +767,6 @@ def append_data_cloud(config: dict[str, Any], datasource_name: str, path: str) -
     append_mock_data(client, datasource_name, path)
 
 
-def mock_data(
-    config: dict[str, Any],
-    project: Project,
-    datasource_name: str,
-    data_format: str,
-    rows: int,
-    context: Optional[str] = None,
-) -> list[dict[str, Any]]:
-    client = get_tinybird_local_client(config, test=False, silent=False)
-    cli_config = CLIConfig.get_project_config()
-    datasource_path = project.get_resource_path(datasource_name, "datasource")
-
-    if not datasource_path:
-        raise CLIMockException(f"Datasource {datasource_name} not found")
-
-    datasource_content = Path(datasource_path).read_text()
-    return create_mock_data(
-        datasource_name,
-        datasource_content,
-        rows,
-        context or "",
-        cli_config,
-        config,
-        cli_config.get_user_token() or "",
-        client,
-        data_format,
-        project.folder,
-    )
-
-
 def analyze_fixture(config: dict[str, Any], fixture_path: str, format: str = "json") -> dict[str, Any]:
     local_client = get_tinybird_local_client(config, test=False, silent=True)
     meta, _data = _analyze(fixture_path, local_client, format)

tinybird/tb/modules/agent/models.py

@@ -8,7 +8,7 @@ from pydantic_ai.retries import AsyncTenacityTransport, wait_retry_after
 from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 
-def create_retrying_client(token: str, workspace_id: str):
+def create_retrying_client(token: str, workspace_id: str, feature: Optional[str] = None):
     """Create a client with smart retry handling for multiple error types."""
 
     def should_retry_status(response):
@@ -29,7 +29,10 @@ def create_retrying_client(token: str, workspace_id: str):
         ),
         validate_response=should_retry_status,
     )
-    return AsyncClient(transport=transport, params={"token": token, "workspace_id": workspace_id})
+    params = {"token": token, "workspace_id": workspace_id}
+    if feature:
+        params["feature"] = feature
+    return AsyncClient(transport=transport, params=params)
 
 
 def create_model(
@@ -38,6 +41,7 @@ def create_model(
     workspace_id: str,
     model: AnthropicModelName = "claude-4-sonnet-20250514",
     run_id: Optional[str] = None,
+    feature: Optional[str] = None,
 ):
     default_headers = {}
     if run_id:

tinybird/tb/modules/agent/utils.py

@@ -44,7 +44,6 @@ class TinybirdAgentContext(BaseModel):
     build_project_test: Callable[..., None]
     deploy_project: Callable[..., None]
     deploy_check_project: Callable[[], None]
-    mock_data: Callable[..., list[dict[str, Any]]]
     append_data_local: Callable[..., None]
     append_data_cloud: Callable[..., None]
     analyze_fixture: Callable[..., dict[str, Any]]

tinybird/tb/modules/cli.py

@@ -488,7 +488,6 @@ def create_ctx_client(
         "logout",
         "update",
         "upgrade",
-        "create",
         "info",
         "tag",
         "push",
@@ -502,7 +501,7 @@ def create_ctx_client(
         return None
 
     commands_always_cloud = ["infra"]
-    commands_always_local = ["build", "dev"]
+    commands_always_local = ["build", "dev", "create"]
     command_always_test = ["test"]
 
     if (

tinybird/tb/modules/create.py

@@ -3,25 +3,22 @@ import os
 import re
 import shutil
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Set
 from urllib.parse import urlparse
 
 import click
 import requests
 
-from tinybird.prompts import claude_rules_prompt, create_prompt, readme_prompt, rules_prompt
-from tinybird.tb.client import TinyB
+from tinybird.prompts import claude_rules_prompt, rules_prompt
+from tinybird.tb.modules.agent import run_agent
 from tinybird.tb.modules.cicd import init_cicd
 from tinybird.tb.modules.cli import cli
-from tinybird.tb.modules.common import _generate_datafile, generate_datafile
+from tinybird.tb.modules.common import _generate_datafile
 from tinybird.tb.modules.config import CLIConfig
 from tinybird.tb.modules.datafile.fixture import persist_fixture
 from tinybird.tb.modules.exceptions import CLICreateException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
-from tinybird.tb.modules.llm import LLM
-from tinybird.tb.modules.llm_utils import extract_xml, parse_xml
 from tinybird.tb.modules.local_common import get_tinybird_local_client
-from tinybird.tb.modules.mock_common import create_mock_data
 from tinybird.tb.modules.project import Project
 
 
@@ -67,13 +64,9 @@ def create(
         folder_path.mkdir()
 
     try:
-        tb_client = config.get_client(token=ctx_config.get("token"), host=ctx_config.get("host"))
-        user_token: str = ""
         created_something = False
-        if prompt:
-            user_token = ctx_config.get("user_token")
-            if not user_token:
-                raise Exception("This action requires authentication. Run 'tb login' first.")
+        if prompt and not ctx_config.get("user_token"):
+            raise Exception("This action requires authentication. Run 'tb login' first.")
 
         if not validate_project_structure(project):
             click.echo(FeedbackManager.highlight(message="\n» Creating new project structure..."))
@@ -99,28 +92,20 @@ def create(
 
         prompt_result: List[Path] = []
         if prompt:
-            prompt_result = create_resources_from_prompt(tb_client, user_token, prompt, project)
-            result.extend(prompt_result)
-            readme_path = folder_path / "README.md"
-            if readme_path.exists():
-                click.echo(FeedbackManager.highlight(message="\n» Updating project description..."))
-            else:
-                click.echo(FeedbackManager.highlight(message="\n» Creating project description..."))
-                readme_path.touch()
-            llm = LLM(user_token=str(user_token), host=tb_client.host)
-            readme_user_prompt = prompt or ""
-            all_resources_xml = get_resources_xml(project)
-            readme_response = llm.ask(
-                system_prompt=readme_prompt(
-                    readme_path.read_text(), tb_client.host, "$TB_ADMIN_TOKEN", all_resources_xml
-                ),
-                prompt=readme_user_prompt,
-                feature="tb_create_readme",
+            prompt_instructions = (
+                "Create or update the Tinybird datasources, pipes, and connections required to satisfy the following request. "
+                "Do not generate mock data or append data; those steps will run later programmatically."
             )
-            readme_result = extract_xml(readme_response, "readme")
-            readme_path.write_text(readme_result)
-            click.echo(FeedbackManager.info_file_created(file="README.md"))
-            created_something = True
+            prompt_result = create_resources_from_prompt(
+                ctx_config,
+                project,
+                prompt,
+                feature="tb_create",
+                instructions=prompt_instructions,
+            )
+            result.extend(prompt_result)
+            if prompt_result:
+                created_something = True
 
         if data or prompt:
             click.echo(FeedbackManager.success(message="✓ Resources created!\n"))
@@ -180,24 +165,30 @@ def create(
                     datasource_name = datasource_path.stem
                     datasource_content = datasource_path.read_text()
                     has_json_path = "`json:" in datasource_content
-                    if has_json_path:
-                        mock_data = create_mock_data(
-                            datasource_name,
-                            datasource_content,
-                            rows,
-                            prompt,
-                            config,
-                            ctx_config,
-                            user_token,
-                            tb_client,
-                            format_="ndjson",
-                            folder=project.folder,
-                        )
-                        if mock_data:
-                            persist_fixture(datasource_name, mock_data, folder, format="ndjson")
-                            click.echo(FeedbackManager.info_file_created(file=f"fixtures/{datasource_name}.ndjson"))
-                            click.echo(FeedbackManager.success(message="✓ Done!"))
-                            created_something = True
+                    if not has_json_path:
+                        continue
+
+                    fixture_path = Path(folder) / "fixtures" / f"{datasource_name}.ndjson"
+                    fixture_existed = fixture_path.exists()
+                    fixture_prompt = (
+                        f"Generate {rows} rows of representative sample data for the Tinybird datasource defined in {datasource_path}. "
+                        f"Store the data in ndjson format at fixtures/{datasource_name}.ndjson."
+                    )
+                    if prompt.strip():
+                        fixture_prompt += f"\n\nOriginal project request:\n{prompt.strip()}"
+
+                    run_agent(
+                        ctx_config,
+                        project,
+                        True,
+                        prompt=fixture_prompt,
+                        feature="tb_mock",
+                    )
+
+                    if fixture_path.exists() and not fixture_existed:
+                        click.echo(FeedbackManager.info_file_created(file=f"fixtures/{datasource_name}.ndjson"))
+                        click.echo(FeedbackManager.success(message="✓ Done!"))
+                        created_something = True
 
         if not created_something and not len(result) > 0:
             click.echo(FeedbackManager.warning(message="△ No resources created\n"))
@@ -291,86 +282,39 @@ def create_project_structure(folder: str):
 
 
 def create_resources_from_prompt(
-    tb_client: TinyB,
-    user_token: str,
-    prompt: str,
+    config: Dict[str, Any],
     project: Project,
-    feature: str = "tb_create_resources",
+    prompt: str,
+    feature: str = "tb_create",
+    instructions: Optional[str] = None,
 ) -> List[Path]:
-    result: List[Path] = []
-    datasource_paths = [Path(ds_file) for ds_file in project.get_datasource_files()]
-    pipes_paths = [Path(pipe_file) for pipe_file in project.get_pipe_files()]
-    connections_paths = [Path(conn_file) for conn_file in project.get_connection_files()]
-    resources_xml = "\n".join(
-        [
-            f"<resource><type>{resource_type}</type><name>{resource_name}</name><content>{resource_content}</content></resource>"
-            for resource_type, resource_name, resource_content in [
-                ("datasource", ds.stem, ds.read_text()) for ds in datasource_paths
-            ]
-            + [
-                (
-                    "pipe",
-                    pipe.stem,
-                    pipe.read_text(),
-                )
-                for pipe in pipes_paths
-            ]
-            + [
-                (
-                    "connection",
-                    conn.stem,
-                    conn.read_text(),
-                )
-                for conn in connections_paths
-            ]
-        ]
-    )
-    llm = LLM(user_token=user_token, host=tb_client.host)
-    prompt_result = llm.ask(system_prompt=create_prompt(resources_xml), prompt=prompt, feature=feature)
-    prompt_result = extract_xml(prompt_result, "response")
-    resources = parse_xml(prompt_result, "resource")
-    datasources = []
-    pipes = []
-    connections = []
-    for resource_xml in resources:
-        resource_type = extract_xml(resource_xml, "type")
-        name = extract_xml(resource_xml, "name")
-        content = extract_xml(resource_xml, "content")
-        resource = {
-            "name": name,
-            "content": content,
-        }
-        if resource_type.lower() == "datasource":
-            datasources.append(resource)
-        elif resource_type.lower() == "pipe":
-            pipes.append(resource)
-        elif resource_type.lower() == "connection":
-            connections.append(resource)
-
-    for ds in datasources:
-        content = ds["content"].replace("```", "")
-        filename = f"{ds['name']}.datasource"
-        ds_file = generate_datafile(
-            content,
-            filename=filename,
-            data=None,
-            _format="ndjson",
-            force=True,
-            folder=project.folder,
-        )
-        result.append(ds_file)
-    for pipe in pipes:
-        content = pipe["content"].replace("```", "")
-        pipe_file = generate_pipe_file(pipe["name"], content, project.folder)
-        result.append(pipe_file)
+    """Run the agent in prompt mode and report newly created project resources."""
 
-    for conn in connections:
-        content = conn["content"].replace("```", "")
-        filename = f"{conn['name']}.connection"
-        conn_file = generate_connection_file(conn["name"], content, project.folder)
-        result.append(conn_file)
+    agent_prompt = prompt.strip()
+    if instructions:
+        instructions = instructions.strip()
+        if agent_prompt:
+            agent_prompt = f"{instructions}\n\n{agent_prompt}"
+        else:
+            agent_prompt = instructions
 
-    return result
+    if not agent_prompt:
+        return []
+
+    resources_before = _collect_project_resource_paths(project)
+    run_agent(config, project, True, prompt=agent_prompt, feature=feature)
+    resources_after = _collect_project_resource_paths(project)
+
+    created_resources = [Path(path) for path in sorted(resources_after - resources_before)]
+    return created_resources
+
+
+def _collect_project_resource_paths(project: Project) -> Set[Path]:
+    resources: Set[Path] = set()
+    resources.update(Path(path) for path in project.get_datasource_files())
+    resources.update(Path(path) for path in project.get_pipe_files())
+    resources.update(Path(path) for path in project.get_connection_files())
+    return resources
 
 
 def init_git(folder: str):
@@ -513,28 +457,6 @@ def save_context(prompt: str, feedback: str):
     context_file.write_text(f"- {prompt}\n{feedback}")
 
 
-def get_resources_xml(project: Project) -> str:
-    datasource_paths = [Path(f) for f in project.get_datasource_files()]
-    pipes_paths = [Path(f) for f in project.get_pipe_files()]
-    resources_xml = "\n".join(
-        [
-            f"<resource><type>{resource_type}</type><name>{resource_name}</name><content>{resource_content}</content></resource>"
-            for resource_type, resource_name, resource_content in [
-                ("datasource", ds.stem, ds.read_text()) for ds in datasource_paths
-            ]
-            + [
-                (
-                    "pipe",
-                    pipe.stem,
-                    pipe.read_text(),
-                )
-                for pipe in pipes_paths
-            ]
-        ]
-    )
-    return resources_xml
-
-
 def create_resources_from_data(
     data: str,
     project: Project,

tinybird/tb/modules/datasource.py

@@ -768,13 +768,31 @@ def datasource_create(
 
     if datasource_type == "prompt":
         click.echo(FeedbackManager.gray(message="\n» Creating .datasource file..."))
-        user_token = config.get("user_token")
-        if not user_token:
+        if not config.get("user_token"):
             raise Exception("This action requires authentication. Run 'tb login' first.")
-        project_config = CLIConfig.get_project_config()
-        tb_client: TinyB = project_config.get_client(token=config.get("token"), host=config.get("host"))
-        create_resources_from_prompt(tb_client, user_token, prompt, project, feature="tb_datasource_create")
-        click.echo(FeedbackManager.success(message="✓ .datasource created!"))
+
+        instructions = (
+            "Create or update a Tinybird datasource (.datasource file) for this project. "
+            "Do not generate mock data or append data; those steps will run later programmatically."
+        )
+        if name:
+            instructions += f" Name the datasource '{name}'."
+
+        created_resources = create_resources_from_prompt(
+            config,
+            project,
+            prompt,
+            feature="tb_datasource_create",
+            instructions=instructions,
+        )
+        if any(path.suffix == ".datasource" for path in created_resources):
+            click.echo(FeedbackManager.success(message="✓ .datasource created!"))
+        else:
+            click.echo(
+                FeedbackManager.gray(
+                    message="△ No new datasource file detected. Existing resources may have been updated instead."
+                )
+            )
         return
 
     connection_required = datasource_type in ("kafka", "s3", "gcs")

tinybird/tb/modules/mock.py

@@ -1,15 +1,9 @@
-import glob
-from pathlib import Path
-
 import click
 
-from tinybird.tb.client import TinyB
+from tinybird.tb.modules.agent import run_agent
 from tinybird.tb.modules.cli import cli
-from tinybird.tb.modules.config import CLIConfig
-from tinybird.tb.modules.datafile.fixture import persist_fixture
 from tinybird.tb.modules.exceptions import CLIMockException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
-from tinybird.tb.modules.mock_common import append_mock_data, create_mock_data
 from tinybird.tb.modules.project import Project
 
 
@@ -31,60 +25,17 @@ from tinybird.tb.modules.project import Project
 )
 @click.pass_context
 def mock(ctx: click.Context, datasource: str, rows: int, prompt: str, format_: str) -> None:
-    """Generate sample data for a data source.
-
-    Args:
-        datasource: Path to the datasource file to load sample data into
-        rows: Number of events to send
-        prompt: Extra context to use for data generation
-        skip: Skip following up on the generated data
-    """
+    """Generate sample data for a data source."""
 
     try:
-        tb_client: TinyB = ctx.ensure_object(dict)["client"]
         project: Project = ctx.ensure_object(dict)["project"]
         ctx_config = ctx.ensure_object(dict)["config"]
+        prompt = f"""Generate mock data for the following datasource: {datasource} with {rows} rows and {format_} format. Extra context: {prompt}"""
         env = ctx.ensure_object(dict)["env"]
-        datasource_path = Path(datasource)
-        datasource_name = datasource
-        folder = project.folder
-        click.echo(FeedbackManager.highlight(message=f"\n» Creating fixture for {datasource_name}..."))
-        if datasource_path.suffix == ".datasource":
-            datasource_name = datasource_path.stem
-        else:
-            datasource_from_glob = glob.glob(f"{folder}/**/{datasource}.datasource")
-            if datasource_from_glob:
-                datasource_path = Path(datasource_from_glob[0])
-
-        if not datasource_path.exists():
-            raise Exception(f"Datasource '{datasource_path.stem}' not found")
-
-        datasource_content = datasource_path.read_text()
-        config = CLIConfig.get_project_config()
-        user_token = ctx_config.get("user_token")
-
-        if not user_token:
-            raise Exception("This action requires authentication. Run 'tb login' first.")
-
-        data = create_mock_data(
-            datasource_name,
-            datasource_content,
-            rows,
-            prompt,
-            config,
-            ctx_config,
-            user_token,
-            tb_client,
-            format_,
-            folder,
-        )
-
-        fixture_path = persist_fixture(datasource_name, data, folder, format=format_)
-        click.echo(FeedbackManager.info(message=f"✓ /fixtures/{datasource_name}.{format_} created"))
         if env == "cloud":
-            append_mock_data(tb_client, datasource_name, str(fixture_path))
+            prompt += "Append the fixture data to the datasource in Tinybird Cloud."
 
-        click.echo(FeedbackManager.success(message=f"✓ Sample data for {datasource_name} created with {rows} rows"))
+        run_agent(ctx_config, project, True, prompt=prompt, feature="tb_mock")
 
     except Exception as e:
         raise CLIMockException(FeedbackManager.error(message=str(e)))

tinybird/tb/modules/mock_common.py

@@ -1,12 +1,5 @@
-from typing import Any, Dict, List
-
-from tinybird.prompts import mock_prompt
 from tinybird.tb.client import TinyB
 from tinybird.tb.modules.common import push_data
-from tinybird.tb.modules.config import CLIConfig
-from tinybird.tb.modules.datafile.fixture import persist_fixture_sql
-from tinybird.tb.modules.llm import LLM
-from tinybird.tb.modules.llm_utils import extract_xml
 
 
 def append_mock_data(
@@ -22,50 +15,3 @@ def append_mock_data(
         concurrency=1,
         silent=True,
     )
-
-
-def create_mock_data(
-    datasource_name: str,
-    datasource_content: str,
-    rows: int,
-    prompt: str,
-    config: CLIConfig,
-    ctx_config: Dict[str, Any],
-    user_token: str,
-    tb_client: TinyB,
-    format_: str,
-    folder: str,
-) -> List[Dict[str, Any]]:
-    user_client = config.get_client(token=ctx_config.get("token"), host=ctx_config.get("host"))
-    llm = LLM(user_token=user_token, host=user_client.host)
-    prompt = f"<datasource_schema>{datasource_content}</datasource_schema>\n<user_input>{prompt}</user_input>"
-    sql = ""
-    attempts = 0
-    data = []
-    error = ""
-    sql_path = None
-    while True:
-        try:
-            response = llm.ask(system_prompt=mock_prompt(rows, error), prompt=prompt, feature="tb_mock")
-            sql = extract_xml(response, "sql")
-            sql_path = persist_fixture_sql(datasource_name, sql, folder)
-            sql_format = "JSON" if format_ == "ndjson" else "CSV"
-            result = tb_client.query(f"SELECT * FROM ({sql}) LIMIT {rows} FORMAT {sql_format}")
-            if sql_format == "JSON":
-                data = result.get("data", [])[:rows]
-                error_response = result.get("error", None)
-                if error_response:
-                    raise Exception(error_response)
-            else:
-                data = result
-            break
-        except Exception as e:
-            error = str(e)
-            attempts += 1
-            if attempts > 5:
-                raise Exception(
-                    f"Failed to generate a valid solution. Check {str(sql_path or '.sql path')} and try again."
-                )
-            else:
-                continue
-    return data

tinybird/tb/modules/test.py

@@ -8,9 +8,10 @@ from typing import Any, Tuple
 import click
 
 from tinybird.tb.client import TinyB
+from tinybird.tb.modules.agent import run_agent
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.project import Project
-from tinybird.tb.modules.test_common import create_test, run_tests, update_test
+from tinybird.tb.modules.test_common import run_tests, update_test
 
 
 @cli.group()
@@ -24,18 +25,18 @@ def test(ctx: click.Context) -> None:
     help="Create a test for an existing pipe",
 )
 @click.argument("name_or_filename", type=str)
-@click.option(
-    "--prompt", type=str, default="Create a test for the selected pipe", help="Prompt to be used to create the test"
-)
+@click.option("--prompt", type=str, default="", help="Prompt to be used to create the test")
 @click.pass_context
 def test_create(ctx: click.Context, name_or_filename: str, prompt: str) -> None:
     """
     Create a test for an existing pipe
     """
     project: Project = ctx.ensure_object(dict)["project"]
-    client: TinyB = ctx.ensure_object(dict)["client"]
     config: dict[str, Any] = ctx.ensure_object(dict)["config"]
-    create_test(name_or_filename, prompt, project, client, config=config)
+    prompt = (
+        f"""Create tests for the following pipe: {name_or_filename}. Extra context: {prompt or "No extra context."}"""
+    )
+    run_agent(config, project, True, prompt=prompt, feature="tb_test_create")
 
 
 @test.command(

tinybird/tb/modules/test_common.py

@@ -14,14 +14,11 @@ import click
 import yaml
 from requests import Response
 
-from tinybird.prompts import test_create_prompt
 from tinybird.tb.client import TinyB
 from tinybird.tb.modules.build_common import process as build_project
 from tinybird.tb.modules.common import sys_exit
 from tinybird.tb.modules.exceptions import CLITestException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
-from tinybird.tb.modules.llm import LLM
-from tinybird.tb.modules.llm_utils import extract_xml, parse_xml
 from tinybird.tb.modules.local_common import get_local_tokens, get_test_workspace_name
 from tinybird.tb.modules.project import Project
 
@@ -53,92 +50,6 @@ def generate_test_file(pipe_name: str, tests: List[Dict[str, Any]], folder: Opti
     return path
 
 
-def create_test(
-    name_or_filename: str, prompt: str, project: Project, client: TinyB, config: dict[str, Any], preview: bool = False
-) -> list[dict[str, Any]]:
-    """
-    Create a test for an existing pipe
-    """
-    tests: List[Dict[str, Any]] = []
-
-    try:
-        click.echo(FeedbackManager.highlight(message="\n» Building test environment"))
-        build_error = build_project(
-            project=project,
-            tb_client=client,
-            watch=False,
-            silent=True,
-            exit_on_error=False,
-            config=config,
-            load_fixtures=True,
-        )
-        if build_error:
-            raise Exception(build_error)
-        click.echo(FeedbackManager.info(message="✓ Done!\n"))
-        folder = project.folder
-        pipe_path = get_pipe_path(name_or_filename, folder)
-        pipe_name = pipe_path.stem
-        click.echo(FeedbackManager.highlight(message=f"» Creating tests for {pipe_name} endpoint..."))
-        pipe_content = pipe_path.read_text()
-        pipe = client._req(f"/v0/pipes/{pipe_name}")
-        parameters = set([param["name"] for node in pipe["nodes"] for param in node["params"]])
-
-        system_prompt = test_create_prompt.format(
-            name=pipe_name,
-            content=pipe_content,
-            parameters=parameters or "No parameters",
-        )
-        user_token = config.get("user_token")
-        if not user_token:
-            raise Exception("No user token found")
-
-        llm = LLM(user_token=user_token, host=config.get("host") or "")
-        response_llm = llm.ask(system_prompt=system_prompt, prompt=prompt, feature="tb_test_create")
-        response_xml = extract_xml(response_llm, "response")
-        tests_content = parse_xml(response_xml, "test")
-
-        for test_content in tests_content:
-            test: Dict[str, Any] = {}
-            test["name"] = extract_xml(test_content, "name")
-            test["description"] = extract_xml(test_content, "description")
-            parameters_api = extract_xml(test_content, "parameters")
-            test["parameters"] = parameters_api.split("?")[1] if "?" in parameters_api else parameters_api
-            test["expected_result"] = ""
-
-            response = None
-            try:
-                response = get_pipe_data(client, pipe_name=pipe_name, test_params=test["parameters"])
-            except Exception:
-                pass
-
-            if response:
-                if response.status_code >= 400:
-                    test["expected_http_status"] = response.status_code
-                    test["expected_result"] = response.json()["error"]
-                else:
-                    test.pop("expected_http_status", None)
-                    test["expected_result"] = response.text or ""
-
-            tests.append(test)
-
-        if not preview:
-            if len(tests) > 0:
-                generate_test_file(pipe_name, tests, folder, mode="a")
-                for test in tests:
-                    test_name = test["name"]
-                    click.echo(FeedbackManager.info(message=f"✓ {test_name} created"))
-            else:
-                click.echo(FeedbackManager.info(message="* No tests created"))
-
-            click.echo(FeedbackManager.success(message="✓ Done!\n"))
-    except Exception as e:
-        raise CLITestException(FeedbackManager.error(message=str(e)))
-    finally:
-        cleanup_test_workspace(client, project.folder)
-
-    return tests
-
-
 def parse_tests(tests_content: str) -> List[Dict[str, Any]]:
     return yaml.safe_load(tests_content)
 

{tinybird-0.0.1.dev301.dist-info → tinybird-0.0.1.dev302.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tinybird
-Version: 0.0.1.dev301
+Version: 0.0.1.dev302
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/forward/commands
 Author: Tinybird

{tinybird-0.0.1.dev301.dist-info → tinybird-0.0.1.dev302.dist-info}/RECORD

@@ -18,7 +18,7 @@ tinybird/datafile/exceptions.py,sha256=8rw2umdZjtby85QbuRKFO5ETz_eRHwUY5l7eHsy1w
 tinybird/datafile/parse_connection.py,sha256=tRyn2Rpr1TeWet5BXmMoQgaotbGdYep1qiTak_OqC5E,1825
 tinybird/datafile/parse_datasource.py,sha256=ssW8QeFSgglVFi3sDZj_HgkJiTJ2069v2JgqnH3CkDE,1825
 tinybird/datafile/parse_pipe.py,sha256=8e9LMecSQVWHC4AXf8cdxoQ1nxUR4fTObYxTctO_EXQ,3816
-tinybird/tb/__cli__.py,sha256=O5MC9cmzrDP6k84cWTDY7avcJMRgZjYRkLdaV70DGbI,247
+tinybird/tb/__cli__.py,sha256=e05Aj_qnRzTfnmbKAbMjhN8CmcAsslxku5omAGI0Ou4,247
 tinybird/tb/check_pypi.py,sha256=Gp0HkHHDFMSDL6nxKlOY51z7z1Uv-2LRexNTZSHHGmM,552
 tinybird/tb/cli.py,sha256=FdDFEIayjmsZEVsVSSvRiVYn_FHOVg_zWQzchnzfWho,1008
 tinybird/tb/client.py,sha256=Z27O6kXbQ1ZVi_K6qzdUJoS877nt7rqQDRckOycjNmI,53544
@@ -26,13 +26,13 @@ tinybird/tb/config.py,sha256=XIKju3xxpo40oQ48zDtSv0vgUrBZBHSmHJi4SHksjXE,5447
 tinybird/tb/modules/build.py,sha256=q3f_CM-_yt7R24VCmQqe4yat1aPdLqmP_hW5WoqdxBg,7907
 tinybird/tb/modules/build_common.py,sha256=aBpumDCZ5FYfBErBa1zGbJMsxdInr68Z21K40W7-DLE,20151
 tinybird/tb/modules/cicd.py,sha256=0KLKccha9IP749QvlXBmzdWv1On3mFwMY4DUcJlBxiE,7326
-tinybird/tb/modules/cli.py,sha256=kPmGW71difgANCRVzKlbpAwGwldFRH-YryRBejM8ff4,20511
+tinybird/tb/modules/cli.py,sha256=geWO5FRlpI9Xuv-sAjfnO-jFwCvaOZ49gAkEvwi-2aY,20503
 tinybird/tb/modules/common.py,sha256=1OwsFHjVPDAuxFctzGjXoqhbKzc_jb1HHq96RaPcaeU,83611
 tinybird/tb/modules/config.py,sha256=gK7rgaWTDd4ZKCrNEg_Uemr26EQjqWt6TjyQKujxOws,11462
 tinybird/tb/modules/connection.py,sha256=axp8Fny1_4PSLJGN4UF6WygyRbQtM3Lbt6thxHKTxzw,17790
 tinybird/tb/modules/copy.py,sha256=dPZkcIDvxjJrlQUIvToO0vsEEEs4EYumbNV77-BzNoU,4404
-tinybird/tb/modules/create.py,sha256=pJxHXG69c9Z_21s-7VuJ3RZOF_nJU51LEwiAkvI3dZY,23251
-tinybird/tb/modules/datasource.py,sha256=7Ugj62Nyj6eBB7v9Jyny44KenYSce7rVs2h09N_0hyE,41812
+tinybird/tb/modules/create.py,sha256=j0WvO6LxuGW1ZOYVuvR1FxFls2oMmKePYNTRJMSXfuU,20135
+tinybird/tb/modules/datasource.py,sha256=O-w4qwhAIitD_SLVdxkpknckliRp_xMswdB8n8Ohxo0,42335
 tinybird/tb/modules/deployment.py,sha256=v0layOmG0IMnuXc3RT39mpGfa5M8yPlrL9F089fJFCo,15964
 tinybird/tb/modules/deployment_common.py,sha256=ZX27oUw7u1ld20qiXH69j7m0DImOc5nckABZ37BPAn4,20406
 tinybird/tb/modules/deprecations.py,sha256=rrszC1f_JJeJ8mUxGoCxckQTJFBCR8wREf4XXXN-PRc,4507
@@ -51,8 +51,8 @@ tinybird/tb/modules/login.py,sha256=zerXZqIv15pbFk5XRt746xGcVnp01YmL_403byBf4jQ,
 tinybird/tb/modules/login_common.py,sha256=CKXD_BjivhGUmBtNbLTJwzQDr6885C72XPZjo0GLLvI,12010
 tinybird/tb/modules/logout.py,sha256=sniI4JNxpTrVeRCp0oGJuQ3yRerG4hH5uz6oBmjv724,1009
 tinybird/tb/modules/materialization.py,sha256=8fdbTTxFSk0Sjet3pcEk_x-cs4RGpgl6toN8vrMLXJE,5630
-tinybird/tb/modules/mock.py,sha256=ET8sRpmXnQsd2sSJXH_KCdREU1_XQgkORru6T357Akc,3260
-tinybird/tb/modules/mock_common.py,sha256=72yKp--Zo40hrycUtiajSRW2BojOsgOZFqUorQ_KQ3w,2279
+tinybird/tb/modules/mock.py,sha256=zxSBvB_Bd3dhVnQUvZEhBREg0tqskzu5dQxF1YldOaw,1479
+tinybird/tb/modules/mock_common.py,sha256=LaqBtbnjpMZBLMpVSwaWdKVCNA485-Gz_96ZfCqgvco,320
 tinybird/tb/modules/open.py,sha256=LYiuO8Z1I9O_v6pv58qpUCWFD6BT00BdeO21fRa4I4Y,1315
 tinybird/tb/modules/pipe.py,sha256=xPKtezhnWZ6k_g82r4XpgKslofhuIxb_PvynH4gdUzI,2393
 tinybird/tb/modules/project.py,sha256=tP6XuG8dR8QCA4pojlXZ7Am6EfS2h3Xnzw5lA3sdtlY,7747
@@ -63,14 +63,14 @@ tinybird/tb/modules/shell.py,sha256=ExOAY2SwywCc5D6X-WCbl7tAV73bYWcGAan4E6u02vI,
 tinybird/tb/modules/sink.py,sha256=dK2s__my0ePIUYrqBzhPSgdWN9rbpvP1G4dT7DJzz80,3865
 tinybird/tb/modules/table.py,sha256=4XrtjM-N0zfNtxVkbvLDQQazno1EPXnxTyo7llivfXk,11035
 tinybird/tb/modules/telemetry.py,sha256=T9gtsQffWqG_4hRBaUJPzOfMkPwz7mH-R6Bn1XRYViA,11482
-tinybird/tb/modules/test.py,sha256=gi7fH6zavFp1jDs3SoRWdf9uBiKjp1hdOPAEleg7H-Y,2145
-tinybird/tb/modules/test_common.py,sha256=501Fo6oDI-ZzJv8bqcLz7-6ulS8-FGViwi1tHq_A5h0,14766
+tinybird/tb/modules/test.py,sha256=NcpRCEDA2itnGRrkao8UNYDZlJ4d35N55OfO_UyJR_U,2229
+tinybird/tb/modules/test_common.py,sha256=7iaHMUSxWsXd98IRMkYezcTa-x1uNEUmG1T48wQenKk,11185
 tinybird/tb/modules/token.py,sha256=ZhW_o7XCr90wJRhMN6816vyo_TVfnzPXyIhrhzQ7oZ0,13807
 tinybird/tb/modules/watch.py,sha256=R0wERRyL-PrwxKOLSk-T-EvkIczHpeoGirrl3JZxXa8,8935
 tinybird/tb/modules/workspace.py,sha256=tCP1zZMwBhLRGm22TGfpSd4cHvQLAS1o_azIXv_r6uw,11172
 tinybird/tb/modules/workspace_members.py,sha256=5JdkJgfuEwbq-t6vxkBhYwgsiTDxF790wsa6Xfif9nk,8608
 tinybird/tb/modules/agent/__init__.py,sha256=i3oe3vDIWWPaicdCM0zs7D7BJ1W0k7th93ooskHAV00,54
-tinybird/tb/modules/agent/agent.py,sha256=R-njI3JNM5_qO5tdXatwV6KGDuyNa93Due9PKkPTShw,39988
+tinybird/tb/modules/agent/agent.py,sha256=_PRCkFgpLCzo-5EdUphLKZ0pG8YMmBfCNFuHdTQKKO8,39057
 tinybird/tb/modules/agent/animations.py,sha256=4WOC5_2BracttmMCrV0H91tXfWcUzQHBUaIJc5FA7tE,3490
 tinybird/tb/modules/agent/banner.py,sha256=l6cO5Fi7lbVKp-GsBP8jf3IkjOWxg2jpAt9NBCy0WR8,4085
 tinybird/tb/modules/agent/command_agent.py,sha256=0Z08rQsir59zQAr-kkOvsKIFpIBsBSTGJJ1VgqqF5WA,3654
@@ -79,10 +79,10 @@ tinybird/tb/modules/agent/explore_agent.py,sha256=QVMPmX6GDW8bh65_X4yXepvAgJmNRQ
 tinybird/tb/modules/agent/file_agent.py,sha256=vE8VChVbG_-uyeBtp3I6QrXeNaPaVlVoNn4c9cGTAJA,2466
 tinybird/tb/modules/agent/memory.py,sha256=vBewB_64L_wHoT4tLT6UX2uxcHwSY880QZ26F9rPqXs,3793
 tinybird/tb/modules/agent/mock_agent.py,sha256=9_w5-MJDC3vp0U5SkM-ucqhxftfpeUM4FPLl5BAr1_0,8200
-tinybird/tb/modules/agent/models.py,sha256=eokO8XlY-kVJOsbqiVporGUAOCyKAXCO5xgTEK9SM6Y,2208
+tinybird/tb/modules/agent/models.py,sha256=QSTwxz_-U4_iQxXT_t3uJZRzZmIsuxVrdmV4APsr1ms,2346
 tinybird/tb/modules/agent/prompts.py,sha256=eua-QawthfG5X6Kz1dRzOZSjXU8tj0nr4k7onmvBZnk,45792
 tinybird/tb/modules/agent/testing_agent.py,sha256=AtwtJViH7805i7djyBgDb7SSUtDyJnw0TWJu6lBFsrg,2953
-tinybird/tb/modules/agent/utils.py,sha256=F4ZuaYPuDCHligFJT_265V8mKyPrsm591DTAzCAIMYk,31998
+tinybird/tb/modules/agent/utils.py,sha256=KMSqOGLTOeyUxQfXES1maROMIKle2uz88z45sFAQXAs,31947
 tinybird/tb/modules/agent/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tinybird/tb/modules/agent/tools/analyze.py,sha256=CR5LXg4fou-zYEksqnjpJ0icvxJVoKnTctoI1NRvqCM,3873
 tinybird/tb/modules/agent/tools/append.py,sha256=cF-1XYmgTXI3dfMQ_nEpOKo6duJGe2GOg64rXNnS9gw,8284
@@ -121,8 +121,8 @@ tinybird/tb_cli_modules/config.py,sha256=IsgdtFRnUrkY8-Zo32lmk6O7u3bHie1QCxLwgp4
 tinybird/tb_cli_modules/exceptions.py,sha256=pmucP4kTF4irIt7dXiG-FcnI-o3mvDusPmch1L8RCWk,3367
 tinybird/tb_cli_modules/regions.py,sha256=QjsL5H6Kg-qr0aYVLrvb1STeJ5Sx_sjvbOYO0LrEGMk,166
 tinybird/tb_cli_modules/telemetry.py,sha256=Hh2Io8ZPROSunbOLuMvuIFU4TqwWPmQTqal4WS09K1A,10449
-tinybird-0.0.1.dev301.dist-info/METADATA,sha256=N4Ot_u5Ky4TsdsOj9LEFHpqY1i3VnaDBicMwZT306dA,1845
-tinybird-0.0.1.dev301.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-tinybird-0.0.1.dev301.dist-info/entry_points.txt,sha256=LwdHU6TfKx4Qs7BqqtaczEZbImgU7Abe9Lp920zb_fo,43
-tinybird-0.0.1.dev301.dist-info/top_level.txt,sha256=VqqqEmkAy7UNaD8-V51FCoMMWXjLUlR0IstvK7tJYVY,54
-tinybird-0.0.1.dev301.dist-info/RECORD,,
+tinybird-0.0.1.dev302.dist-info/METADATA,sha256=OUZdrAs2qYQQfpCENDpvzQeMu8KR8Eomfgwayia55Uo,1845
+tinybird-0.0.1.dev302.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+tinybird-0.0.1.dev302.dist-info/entry_points.txt,sha256=LwdHU6TfKx4Qs7BqqtaczEZbImgU7Abe9Lp920zb_fo,43
+tinybird-0.0.1.dev302.dist-info/top_level.txt,sha256=VqqqEmkAy7UNaD8-V51FCoMMWXjLUlR0IstvK7tJYVY,54
+tinybird-0.0.1.dev302.dist-info/RECORD,,