tinybird 0.0.1.dev291__py3-none-any.whl → 1.0.5__py3-none-any.whl

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
+
+    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
 
-    return result
+
+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,
@@ -602,6 +524,9 @@ def generate_kafka_connection_with_secrets(
     tb_secret_secret: Optional[str],
     security_protocol: str,
     sasl_mechanism: str,
+    ssl_ca_pem: Optional[str],
+    tb_secret_ssl_ca_pem: Optional[str],
+    schema_registry_url: Optional[str],
     folder: str,
 ) -> Path:
     kafka_bootstrap_servers = (
@@ -609,6 +534,7 @@ def generate_kafka_connection_with_secrets(
     )
     kafka_key = inject_tb_secret(tb_secret_key) if tb_secret_key else key
     kafka_secret = inject_tb_secret(tb_secret_secret) if tb_secret_secret else secret
+    kafka_ssl_ca_pem = inject_tb_secret(tb_secret_ssl_ca_pem) if tb_secret_ssl_ca_pem else ssl_ca_pem
     content = f"""TYPE kafka
 KAFKA_BOOTSTRAP_SERVERS {kafka_bootstrap_servers}
 KAFKA_SECURITY_PROTOCOL {security_protocol or "SASL_SSL"}
@@ -616,6 +542,13 @@ KAFKA_SASL_MECHANISM {sasl_mechanism or "PLAIN"}
 KAFKA_KEY {kafka_key}
 KAFKA_SECRET {kafka_secret}
 """
+    if schema_registry_url:
+        content += f"""KAFKA_SCHEMA_REGISTRY_URL {schema_registry_url}\n"""
+    if kafka_ssl_ca_pem:
+        content += f"""KAFKA_SSL_CA_PEM >\n    {kafka_ssl_ca_pem}\n"""
+    content += """# Learn more at https://www.tinybird.co/docs/forward/get-data-in/connectors/kafka#kafka-connection-settings
+"""
+
     return generate_connection_file(name, content, folder, skip_feedback=True)
 
 

tinybird/tb/modules/datafile/build.py

@@ -128,41 +128,33 @@ def folder_build(
                     filename = to_run[name]["filename"]
                     filename = filename.replace(f"{folder}/", "")
                     click.echo(FeedbackManager.info(message=f"✓ {filename}"))
-                else:
-                    if raise_on_exists:
-                        raise AlreadyExistsException(
-                            FeedbackManager.warning_name_already_exists(
-                                name=name if to_run[name]["version"] is None else f"{name}__v{to_run[name]['version']}"
-                            )
+                elif raise_on_exists:
+                    raise AlreadyExistsException(
+                        FeedbackManager.warning_name_already_exists(
+                            name=name if to_run[name]["version"] is None else f"{name}__v{to_run[name]['version']}"
                         )
+                    )
+                elif name_matches_existing_resource(resource, name, tb_client):
+                    if resource == "pipes":
+                        click.echo(FeedbackManager.error_pipe_cannot_be_pushed(name=name))
                     else:
-                        if name_matches_existing_resource(resource, name, tb_client):
-                            if resource == "pipes":
-                                click.echo(FeedbackManager.error_pipe_cannot_be_pushed(name=name))
-                            else:
-                                click.echo(FeedbackManager.error_datasource_cannot_be_pushed(name=name))
-                        else:
-                            click.echo(
-                                FeedbackManager.warning_name_already_exists(
-                                    name=(
-                                        name
-                                        if to_run[name]["version"] is None
-                                        else f"{name}__v{to_run[name]['version']}"
-                                    )
-                                )
-                            )
-            else:
-                if should_push_file(name, remote_resource_names, force, run_tests):
-                    extension = "pipe" if resource == "pipes" else "datasource"
-                    click.echo(FeedbackManager.info_building_resource(name=f"{name}.{extension}", version=""))
+                        click.echo(FeedbackManager.error_datasource_cannot_be_pushed(name=name))
                 else:
-                    if name_matches_existing_resource(resource, name, tb_client):
-                        if resource == "pipes":
-                            click.echo(FeedbackManager.warning_pipe_cannot_be_pushed(name=name))
-                        else:
-                            click.echo(FeedbackManager.warning_datasource_cannot_be_pushed(name=name))
-                    else:
-                        click.echo(FeedbackManager.warning_dry_name_already_exists(name=name))
+                    click.echo(
+                        FeedbackManager.warning_name_already_exists(
+                            name=(name if to_run[name]["version"] is None else f"{name}__v{to_run[name]['version']}")
+                        )
+                    )
+            elif should_push_file(name, remote_resource_names, force, run_tests):
+                extension = "pipe" if resource == "pipes" else "datasource"
+                click.echo(FeedbackManager.info_building_resource(name=f"{name}.{extension}", version=""))
+            elif name_matches_existing_resource(resource, name, tb_client):
+                if resource == "pipes":
+                    click.echo(FeedbackManager.warning_pipe_cannot_be_pushed(name=name))
+                else:
+                    click.echo(FeedbackManager.warning_datasource_cannot_be_pushed(name=name))
+            else:
+                click.echo(FeedbackManager.warning_dry_name_already_exists(name=name))
 
     def push_files(
         dependency_graph: GraphDependencies,
@@ -932,11 +924,8 @@ def process_file(
                 raise click.ClickException(FeedbackManager.error_missing_table_arn(datasource=datasource["name"]))
             if not params.get("import_export_bucket", None):
                 raise click.ClickException(FeedbackManager.error_missing_export_bucket(datasource=datasource["name"]))
-        else:
-            if not params.get("import_external_datasource", None):
-                raise click.ClickException(
-                    FeedbackManager.error_missing_external_datasource(datasource=datasource["name"])
-                )
+        elif not params.get("import_external_datasource", None):
+            raise click.ClickException(FeedbackManager.error_missing_external_datasource(datasource=datasource["name"]))
 
         return params
 
@@ -1059,7 +1048,7 @@ def process_file(
             #
             # Note: any unknown import_ parameter is leaved as is.
             for key in ImportReplacements.get_datafile_parameter_keys():
-                replacement, default_value = ImportReplacements.get_api_param_for_datafile_param(service, key)
+                replacement, default_value = ImportReplacements.get_api_param_for_datafile_param(key)
                 if not replacement:
                     continue  # We should not reach this never, but just in case...
 

tinybird/tb/modules/datafile/build_datasource.py

@@ -301,32 +301,28 @@ def new_ds(
     if alter_response and make_changes:
         # alter operation finished
         pass
+    elif (
+        os.getenv("TB_I_KNOW_WHAT_I_AM_DOING")
+        and click.prompt(FeedbackManager.info_ask_for_datasource_confirmation()) == ds_name
+    ):  # TODO move to CLI
+        try:
+            client.datasource_delete(ds_name)
+            click.echo(FeedbackManager.success_delete_datasource(datasource=ds_name))
+        except Exception:
+            raise click.ClickException(FeedbackManager.error_removing_datasource(datasource=ds_name))
+        return
+    elif alter_error_message:
+        raise click.ClickException(
+            FeedbackManager.error_datasource_already_exists_and_alter_failed(
+                datasource=ds_name, alter_error_message=alter_error_message
+            )
+        )
+    elif promote_error_message:
+        raise click.ClickException(
+            FeedbackManager.error_promoting_datasource(datasource=ds_name, error=promote_error_message)
+        )
     else:
-        # removed replacing by default. When a datasource is removed data is
-        # removed and all the references needs to be updated
-        if (
-            os.getenv("TB_I_KNOW_WHAT_I_AM_DOING")
-            and click.prompt(FeedbackManager.info_ask_for_datasource_confirmation()) == ds_name
-        ):  # TODO move to CLI
-            try:
-                client.datasource_delete(ds_name)
-                click.echo(FeedbackManager.success_delete_datasource(datasource=ds_name))
-            except Exception:
-                raise click.ClickException(FeedbackManager.error_removing_datasource(datasource=ds_name))
-            return
-        else:
-            if alter_error_message:
-                raise click.ClickException(
-                    FeedbackManager.error_datasource_already_exists_and_alter_failed(
-                        datasource=ds_name, alter_error_message=alter_error_message
-                    )
-                )
-            if promote_error_message:
-                raise click.ClickException(
-                    FeedbackManager.error_promoting_datasource(datasource=ds_name, error=promote_error_message)
-                )
-            else:
-                click.echo(FeedbackManager.warning_datasource_already_exists(datasource=ds_name))
+        click.echo(FeedbackManager.warning_datasource_already_exists(datasource=ds_name))
 
 
 def share_and_unshare_datasource(

tinybird/tb/modules/datafile/diff.py

@@ -168,7 +168,7 @@ def diff_command(
                 sys.stdout.writelines(diff_lines)
                 click.echo("")
 
-    for rfilename, _ in local_resources.items():
+    for rfilename in local_resources.keys():
         if rfilename not in changed:
             for resource in remote_datasources + remote_pipes:
                 properties = get_name_version(resource["name"])

tinybird/tb/modules/datafile/format_pipe.py

@@ -25,11 +25,25 @@ from tinybird.tb.modules.datafile.format_datasource import format_engine
 
 
 def format_node_sql(
-    file_parts: List[str], node: Dict[str, Any], line_length: Optional[int] = None, lower_keywords: bool = False
+    file_parts: List[str],
+    node: Dict[str, Any],
+    line_length: Optional[int] = None,
+    lower_keywords: bool = False,
+    resource_name: Optional[str] = None,
+    resource_source: Optional[str] = None,
 ) -> List[str]:
     file_parts.append("SQL >")
     file_parts.append(DATAFILE_NEW_LINE)
-    file_parts.append(format_sql(node["sql"], DATAFILE_INDENT, line_length=line_length, lower_keywords=lower_keywords))
+    file_parts.append(
+        format_sql(
+            node["sql"],
+            DATAFILE_INDENT,
+            line_length=line_length,
+            lower_keywords=lower_keywords,
+            resource_name=resource_name,
+            resource_source=resource_source,
+        )
+    )
     file_parts.append(DATAFILE_NEW_LINE)
     file_parts.append(DATAFILE_NEW_LINE)
     return file_parts
@@ -93,8 +107,21 @@ def format_pipe_include(file_parts: List[str], node: Dict[str, Any], includes: D
     return file_parts
 
 
-def format_sql(sql: str, DATAFILE_INDENT: str, line_length: Optional[int] = None, lower_keywords: bool = False) -> str:
-    sql = format_sql_template(sql.strip(), line_length=line_length, lower_keywords=lower_keywords)
+def format_sql(
+    sql: str,
+    DATAFILE_INDENT: str,
+    line_length: Optional[int] = None,
+    lower_keywords: bool = False,
+    resource_name: Optional[str] = None,
+    resource_source: Optional[str] = None,
+) -> str:
+    sql = format_sql_template(
+        sql.strip(),
+        line_length=line_length,
+        lower_keywords=lower_keywords,
+        resource_name=resource_name,
+        resource_source=resource_source,
+    )
     return "\n".join([f"{DATAFILE_INDENT}{part}" for part in sql.split("\n") if len(part.strip())])
 
 
@@ -105,6 +132,8 @@ def format_node(
     line_length: Optional[int] = None,
     unroll_includes: bool = False,
     lower_keywords: bool = False,
+    resource_name: Optional[str] = None,
+    resource_source: Optional[str] = None,
 ) -> None:
     if not unroll_includes:
         format_pipe_include(file_parts, node, includes)
@@ -119,7 +148,14 @@ def format_node(
 
     Doc = namedtuple("Doc", ["description"])
     format_description(file_parts, Doc(node.get("description", "")))
-    format_node_sql(file_parts, node, line_length=line_length, lower_keywords=lower_keywords)
+    format_node_sql(
+        file_parts,
+        node,
+        line_length=line_length,
+        lower_keywords=lower_keywords,
+        resource_name=resource_name,
+        resource_source=resource_source,
+    )
     format_node_type(file_parts, node)
 
 
@@ -132,6 +168,7 @@ def format_pipe(
     for_deploy_diff: bool = False,
     skip_eval: bool = False,
     content: Optional[str] = None,
+    resource_source: Optional[str] = None,
 ) -> str:
     if datafile:
         doc = datafile
@@ -162,7 +199,7 @@ def format_pipe(
                 if "." in include_file
                 else eval_var(include_file)
             )
-            included_pipe = parse_pipe(include_file, skip_eval=skip_eval).datafile
+            included_pipe = parse_pipe(str(include_file), skip_eval=skip_eval).datafile
             pipe_nodes = doc.nodes.copy()
             for included_node in included_pipe.nodes.copy():
                 unrolled_included_node = next(
@@ -178,10 +215,12 @@ def format_pipe(
             line_length=line_length,
             unroll_includes=unroll_includes,
             lower_keywords=bool(for_deploy_diff),
+            resource_name=filename,
+            resource_source=resource_source,
         )
 
     if not unroll_includes:
-        for k, _ in doc.includes.items():
+        for k in doc.includes.keys():
             if ".incl" not in k:
                 continue
             file_parts.append(f"INCLUDE {k}")