tinybird 0.0.1.dev300__py3-none-any.whl → 0.0.1.dev302__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
 
-    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/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

@@ -181,7 +181,7 @@ def format_pipe(
         )
 
     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}")

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/exceptions.py

@@ -127,6 +127,13 @@ class CLIMockException(CLIException):
         super().__init__(message, "mock_error", **kw_telemetry_event_data)
 
 
+class CLIAgentException(CLIException):
+    """Exceptions generated by the agent commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "agent_error", **kw_telemetry_event_data)
+
+
 class CLILoginException(CLIException):
     """Exceptions generated by the login commands"""
 

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/shell.py

@@ -294,7 +294,7 @@ class Shell:
         arg = argline.strip().lower()
         if not arg:
             return
-        if arg.startswith("with") or arg.startswith("select"):
+        if arg.startswith(("with", "select")):
             self.run_sql(argline)
         elif len(arg.split()) == 1 and arg in self.project.pipes + self.project.datasources:
             self.run_sql(f"select * from {argline}")

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/tb_cli_modules/common.py

@@ -1042,7 +1042,7 @@ def get_format_from_filename_or_url(filename_or_url: str) -> str:
     'csv'
     """
     filename_or_url = filename_or_url.lower()
-    if filename_or_url.endswith("json") or filename_or_url.endswith("ndjson"):
+    if filename_or_url.endswith(("json", "ndjson")):
         return "ndjson"
     if filename_or_url.endswith("parquet"):
         return "parquet"
@@ -1050,7 +1050,7 @@ def get_format_from_filename_or_url(filename_or_url: str) -> str:
         return "csv"
     try:
         parsed = urlparse(filename_or_url)
-        if parsed.path.endswith("json") or parsed.path.endswith("ndjson"):
+        if parsed.path.endswith(("json", "ndjson")):
             return "ndjson"
         if parsed.path.endswith("parquet"):
             return "parquet"