tinybird 0.0.1.dev29__py3-none-any.whl → 0.0.1.dev30__py3-none-any.whl

tinybird/tb/modules/common.py

@@ -175,7 +175,7 @@ def generate_datafile(
     if not f.exists() or force:
         with open(f"{f}", "w") as ds_file:
             ds_file.write(datafile)
-        click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder)))
+        click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder or ".")))
 
         if data and (base / "fixtures").exists():
             # Generating a fixture for Parquet files is not so trivial, since Parquet format
@@ -1177,31 +1177,6 @@ async def print_current_workspace(config: CLIConfig) -> None:
     echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
 
 
-async def print_current_branch(config: CLIConfig) -> None:
-    _ = await try_update_config_with_remote(config, only_if_needed=True)
-
-    response = await config.get_client().user_workspaces_and_branches()
-
-    columns = ["name", "id", "workspace"]
-    table = []
-
-    for workspace in response["workspaces"]:
-        if config["id"] == workspace["id"]:
-            click.echo(FeedbackManager.info_current_branch())
-            if workspace.get("is_branch"):
-                name = workspace["name"]
-                main_workspace = await get_current_main_workspace(config)
-                assert isinstance(main_workspace, dict)
-                main_name = main_workspace["name"]
-            else:
-                name = MAIN_BRANCH
-                main_name = workspace["name"]
-            table.append([name, workspace["id"], main_name])
-            break
-
-    echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
-
-
 class ConnectionReplacements:
     _PARAMS_REPLACEMENTS: Dict[str, Dict[str, str]] = {
         "s3": {

tinybird/tb/modules/config.py

@@ -323,14 +323,6 @@ class CLIConfig:
         path: str = os.path.join(working_dir, ".tinyb")
         return CLIConfig(path, parent=CLIConfig.get_global_config())
 
-    @staticmethod
-    def get_llm_config(working_dir: Optional[str] = None) -> Dict[str, Any]:
-        return (
-            CLIConfig.get_project_config(working_dir)
-            .get("llms", {})
-            .get("openai", {"model": "gpt-4o-mini", "api_key": None})
-        )
-
     @staticmethod
     def reset() -> None:
         CLIConfig._global = None

tinybird/tb/modules/create.py

@@ -1,4 +1,5 @@
 import os
+import re
 from os import getcwd
 from pathlib import Path
 from typing import Optional
@@ -260,7 +261,25 @@ def init_git(folder: str):
 
 
 def generate_pipe_file(name: str, content: str, folder: str):
-    base = Path(folder) / "endpoints"
+    def is_copy(content: str) -> bool:
+        return re.search(r"TYPE copy", content, re.IGNORECASE) is not None
+
+    def is_materialization(content: str) -> bool:
+        return re.search(r"TYPE materialized", content, re.IGNORECASE) is not None
+
+    def is_sink(content: str) -> bool:
+        return re.search(r"TYPE sink", content, re.IGNORECASE) is not None
+
+    if is_copy(content):
+        pathname = "copies"
+    elif is_materialization(content):
+        pathname = "materializations"
+    elif is_sink(content):
+        pathname = "sinks"
+    else:
+        pathname = "endpoints"
+
+    base = Path(folder) / pathname
     if not base.exists():
         base = Path()
     f = base / (f"{name}.pipe")

tinybird/tb/modules/datafile/build.py

@@ -768,13 +768,13 @@ async def process(
             raise click.ClickException(FeedbackManager.error_forkdownstream_pipes_with_engine(pipe=resource_name))
 
         to_run[resource_name] = r
-        file_deps = r.get("deps", [])
+        file_deps: List[str] = r.get("deps", [])
         deps += file_deps
         # calculate and look for deps
         dep_list = []
         for x in file_deps:
             if x not in INTERNAL_TABLES or is_internal:
-                f, ds = find_file_by_name(dir_path, x, verbose, vendor_paths=vendor_paths, resource=r)
+                f, ds = find_file_by_name(dir_path or ".", x, verbose, vendor_paths=vendor_paths, resource=r)
                 if f:
                     dep_list.append(f.rsplit(".", 1)[0])
                 if ds:

tinybird/tb/modules/datafile/build_pipe.py

@@ -266,7 +266,19 @@ async def new_pipe(
 
     if data.get("type") == "endpoint":
         token = tb_client.token
-        click.echo(f"""** => Test endpoint with:\n** $ curl {host}/v0/pipes/{p["name"]}.json?token={token}""")
+        try:
+            example_params = {
+                "format": "json",
+                "pipe": p["name"],
+                "q": "",
+                "token": token,
+            }
+            endpoint_url = await tb_client._req(f"/examples/query.http?{urlencode(example_params)}")
+            if endpoint_url:
+                endpoint_url = endpoint_url.replace("http://localhost:8001", host)
+                click.echo(f"""** => Test endpoint with:\n** $ curl {endpoint_url}""")
+        except Exception:
+            pass
 
 
 async def get_token_from_main_branch(branch_tb_client: TinyB) -> Optional[str]:

tinybird/tb/modules/datasource.py

@@ -542,7 +542,7 @@ async def datasource_share(ctx: Context, datasource_name: str, workspace_name_or
     """Share a datasource"""
 
     config = CLIConfig.get_project_config()
-    client = config.get_client()
+    client: TinyB = ctx.ensure_object(dict)["client"]
     host = config.get_host() or CLIConfig.DEFAULTS["host"]
     ui_host = get_display_host(host)
 

tinybird/tb/modules/llm.py

@@ -1,8 +1,10 @@
 import json
+import os
 import urllib.parse
 from copy import deepcopy
-from typing import List
+from typing import Any, List
 
+from anthropic import AnthropicVertex
 from pydantic import BaseModel
 
 from tinybird.client import TinyB
@@ -37,7 +39,7 @@ class LLM:
         self.user_client = deepcopy(client)
         self.user_client.token = user_token
 
-    async def ask(self, prompt: str, system_prompt: str = "", model: str = "o1-mini") -> str:
+    async def ask(self, prompt: str, system_prompt: str = "") -> str:
         """
         Calls the model with the given prompt and returns the response.
 
@@ -47,7 +49,7 @@ class LLM:
         Returns:
             str: The response from the language model.
         """
-        messages = []
+        messages: List[Any] = []
 
         if system_prompt:
             messages.append({"role": "user", "content": system_prompt})
@@ -55,8 +57,21 @@ class LLM:
         if prompt:
             messages.append({"role": "user", "content": prompt})
 
+        if gcloud_access_token := os.getenv("GCLOUD_ACCESS_TOKEN"):
+            client = AnthropicVertex(
+                region="europe-west1",
+                project_id="gen-lang-client-0705305160",
+                access_token=gcloud_access_token,
+            )
+            message = client.messages.create(
+                max_tokens=8000,
+                messages=messages,
+                model="claude-3-5-sonnet-v2@20241022",
+            )
+            return message.content[0].text or ""  # type: ignore
+
         data = {
-            "model": model,
+            "model": "o1-mini",
             "messages": messages,
         }
         response = await self.user_client._req(

tinybird/tb/modules/local.py

@@ -82,7 +82,7 @@ def get_docker_client():
         client.ping()
         return client
     except Exception:
-        raise CLIException("Docker is not running or installed. Make sure Docker is installed and running.")
+        raise CLIException("No container runtime is running. Make sure a Docker-compatible runtime is installed and running.")
 
 
 def stop_tinybird_local(docker_client):

tinybird/tb/modules/login.py

@@ -135,7 +135,12 @@ async def login(host: str, workspace: str):
     server_thread.start()
 
     # Open the browser to the auth page
-    client_id = "T6excMo8IKguvUw4vFNYfqlt9pe6msCU"
+    if "wadus" in host:
+        client_id = "Rpl7Uy9aSjqoPCSvHgGl3zNQuZcSOXBe"
+        base_auth_url = "https://auth.wadus1.tinybird.co"
+    else:
+        client_id = "T6excMo8IKguvUw4vFNYfqlt9pe6msCU"
+        base_auth_url = "https://auth.tinybird.co"
     callback_url = f"http://localhost:{AUTH_SERVER_PORT}"
     params = {
         "client_id": client_id,
@@ -143,7 +148,7 @@ async def login(host: str, workspace: str):
         "response_type": "token",
         "scope": "openid profile email",
     }
-    auth_url = f"https://auth.tinybird.co/authorize?{urlencode(params)}"
+    auth_url = f"{base_auth_url}/authorize?{urlencode(params)}"
     webbrowser.open(auth_url)
 
     # Wait for the authentication to complete or timeout

tinybird/tb/modules/shell.py

@@ -274,7 +274,7 @@ class Shell:
         elif arg.startswith("mock"):
             self.handle_mock(arg)
         else:
-            subprocess.run(f"tb --local {arg}", shell=True, text=True)
+            subprocess.run(f"tb {arg}", shell=True, text=True)
 
     def default(self, argline):
         click.echo("")
@@ -286,7 +286,7 @@ class Shell:
         elif len(arg.split()) == 1 and arg in self.endpoints + self.pipes + self.datasources + self.shared_datasources:
             self.run_sql(f"select * from {arg}")
         else:
-            subprocess.run(f"tb --local {arg}", shell=True, text=True)
+            subprocess.run(f"tb {arg}", shell=True, text=True)
 
     def run_sql(self, query, rows_limit=20):
         try:

tinybird/tb/modules/test.py

@@ -6,11 +6,13 @@
 import difflib
 import glob
 import os
+import urllib.parse
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import click
 import yaml
+from requests import Response
 
 from tinybird.prompts import test_create_prompt
 from tinybird.tb.modules.cli import cli
@@ -66,9 +68,9 @@ def test(ctx: click.Context) -> None:
     type=click.Path(exists=True, file_okay=False),
     help="Folder where datafiles will be placed",
 )
-@click.option("--prompt", type=str, default=None, help="Prompt to be used to create the test")
+@click.option("--prompt", type=str, default="", help="Prompt to be used to create the test")
 @coro
-async def test_create(name_or_filename: str, prompt: Optional[str], folder: str) -> None:
+async def test_create(name_or_filename: str, prompt: str, folder: str) -> None:
     """
     Create a test for an existing pipe
     """
@@ -80,15 +82,13 @@ async def test_create(name_or_filename: str, prompt: Optional[str], folder: str)
                 raise CLIException(FeedbackManager.error(message=f"Pipe {name_or_filename} not found"))
         else:
             pipe_folders = ("endpoints", "copies", "materializations", "sinks", "pipes")
-            pipe_path = next(
-                (
+            try:
+                pipe_path = next(
                     root_path / folder / f"{name_or_filename}.pipe"
                     for folder in pipe_folders
                     if (root_path / folder / f"{name_or_filename}.pipe").exists()
-                ),
-                None,
-            )
-            if not pipe_path:
+                )
+            except Exception:
                 raise CLIException(FeedbackManager.error(message=f"Pipe {name_or_filename} not found"))
 
         pipe_name = pipe_path.stem
@@ -97,8 +97,8 @@ async def test_create(name_or_filename: str, prompt: Optional[str], folder: str)
         pipe_content = pipe_path.read_text()
 
         client = await get_tinybird_local_client(os.path.abspath(folder))
-        pipe_nodes = await client._req(f"/v0/pipes/{pipe_name}")
-        parameters = set([param["name"] for node in pipe_nodes["nodes"] for param in node["params"]])
+        pipe = await 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,
@@ -107,24 +107,26 @@ async def test_create(name_or_filename: str, prompt: Optional[str], folder: str)
         )
         config = CLIConfig.get_project_config(folder)
         user_token = config.get_user_token()
+        if not user_token:
+            raise CLIException(FeedbackManager.error(message="No user token found"))
         llm = LLM(user_token=user_token, client=config.get_client())
 
-        response = await llm.ask(prompt, system_prompt=system_prompt)
-        response = extract_xml(response, "response")
-        tests_content = parse_xml(response, "test")
+        response_llm = await llm.ask(prompt, system_prompt=system_prompt)
+        response_xml = extract_xml(response_llm, "response")
+        tests_content = parse_xml(response_xml, "test")
 
         tests: List[Dict[str, Any]] = []
         for test_content in tests_content:
-            test = {}
+            test: Dict[str, Any] = {}
             test["name"] = extract_xml(test_content, "name")
             test["description"] = extract_xml(test_content, "description")
-            parameters = extract_xml(test_content, "parameters")
-            test["parameters"] = parameters.split("?")[1] if "?" in parameters else parameters
+            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 = await client._req_raw(f"/v0/pipes/{pipe_name}.ndjson?{parameters}")
+                response = await get_pipe_data(client, pipe_name=pipe_name, test_params=test["parameters"])
             except Exception:
                 pass
 
@@ -181,7 +183,7 @@ async def test_update(pipe: str, folder: str) -> None:
             test_params = test["parameters"].split("?")[1] if "?" in test["parameters"] else test["parameters"]
             response = None
             try:
-                response = await client._req_raw(f"/v0/pipes/{pipe_name}.ndjson?{test_params}")
+                response = await get_pipe_data(client, pipe_name=pipe_name, test_params=test_params)
             except Exception:
                 continue
 
@@ -221,20 +223,21 @@ async def run_tests(name: Tuple[str, ...], folder: str) -> None:
     client = await get_tinybird_local_client(os.path.abspath(folder))
     paths = [Path(n) for n in name]
     endpoints = [f"./tests/{p.stem}.yaml" for p in paths]
-    test_files: Iterable[str] = endpoints if len(endpoints) > 0 else glob.glob("./tests/**/*.y*ml", recursive=True)
+    test_files: List[str] = endpoints if len(endpoints) > 0 else glob.glob("./tests/**/*.y*ml", recursive=True)
 
     async def run_test(test_file):
         test_file_path = Path(test_file)
         click.echo(FeedbackManager.info(message=f"\n* {test_file_path.stem}{test_file_path.suffix}"))
         test_file_content = yaml.safe_load(test_file_path.read_text())
+
         for test in test_file_content:
             try:
                 test_params = test["parameters"].split("?")[1] if "?" in test["parameters"] else test["parameters"]
                 response = None
                 try:
-                    response = await client._req_raw(f"/v0/pipes/{test_file_path.stem}.ndjson?{test_params}")
+                    response = await get_pipe_data(client, pipe_name=test_file_path.stem, test_params=test_params)
                 except Exception:
-                    raise Exception("Expected to not fail but got an error")
+                    continue
 
                 expected_result = response.text
                 if response.status_code >= 400:
@@ -271,3 +274,19 @@ async def run_tests(name: Tuple[str, ...], folder: str) -> None:
         exit(1)
     else:
         click.echo(FeedbackManager.success(message=f"\n✓ {test_count}/{test_count} passed"))
+
+
+async def get_pipe_data(client, pipe_name: str, test_params: str) -> Response:
+    pipe = await client._req(f"/v0/pipes/{pipe_name}")
+    output_node = next(
+        (node for node in pipe["nodes"] if node["node_type"] != "default" and node["node_type"] != "standard"),
+        {"name": "not_found"},
+    )
+    if output_node["node_type"] == "endpoint":
+        return await client._req_raw(f"/v0/pipes/{pipe_name}.ndjson?{test_params}")
+
+    params = {
+        "q": output_node["sql"],
+        "pipeline": pipe_name,
+    }
+    return await client._req_raw(f"""/v0/sql?{urllib.parse.urlencode(params)}&{test_params}""")

tinybird/tb/modules/update.py

@@ -0,0 +1,182 @@
+import os
+import re
+from os import getcwd
+from pathlib import Path
+from typing import Optional
+
+import click
+
+from tinybird.client import TinyB
+from tinybird.prompts import mock_prompt, update_prompt
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import check_user_token_with_client, coro, generate_datafile
+from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.datafile.fixture import build_fixture_name, persist_fixture
+from tinybird.tb.modules.exceptions import CLIException
+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
+
+
+@cli.command()
+@click.argument("prompt")
+@click.option(
+    "--folder",
+    default=".",
+    type=click.Path(exists=False, file_okay=False),
+    help="Folder where project files will be placed",
+)
+@coro
+async def update(
+    prompt: str,
+    folder: str,
+) -> None:
+    """Update resources in the project."""
+    folder = folder or getcwd()
+    folder_path = Path(folder)
+    if not folder_path.exists():
+        folder_path.mkdir()
+
+    try:
+        config = CLIConfig.get_project_config(folder)
+        tb_client = config.get_client()
+        user_token: Optional[str] = None
+        try:
+            user_token = config.get_user_token()
+            if not user_token:
+                raise CLIException("No user token found")
+            await check_user_token_with_client(tb_client, token=user_token)
+        except Exception as e:
+            click.echo(
+                FeedbackManager.error(message=f"This action requires authentication. Run 'tb login' first. Error: {e}")
+            )
+            return
+
+        local_client = await get_tinybird_local_client(folder)
+
+        click.echo(FeedbackManager.highlight(message="\n» Updating resources..."))
+        datasources_updated = await update_resources(tb_client, user_token, prompt, folder)
+        click.echo(FeedbackManager.success(message="✓ Done!\n"))
+
+        if datasources_updated and user_token:
+            click.echo(FeedbackManager.highlight(message="\n» Generating fixtures..."))
+
+            datasource_files = [f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")]
+            for datasource_file in datasource_files:
+                datasource_path = Path(folder) / "datasources" / datasource_file
+                llm = LLM(user_token=user_token, client=tb_client)
+                datasource_name = datasource_path.stem
+                datasource_content = datasource_path.read_text()
+                has_json_path = "`json:" in datasource_content
+                if has_json_path:
+                    prompt = f"<datasource_schema>{datasource_content}</datasource_schema>\n<user_input>{prompt}</user_input>"
+                    response = await llm.ask(prompt, system_prompt=mock_prompt(rows=20))
+                    sql = extract_xml(response, "sql")
+                    sql = sql.split("FORMAT")[0]
+                    result = await local_client.query(f"{sql} FORMAT JSON")
+                    data = result.get("data", [])
+                    fixture_name = build_fixture_name(
+                        datasource_path.absolute().as_posix(), datasource_name, datasource_content
+                    )
+                    if data:
+                        persist_fixture(fixture_name, data, folder)
+                        click.echo(FeedbackManager.info(message=f"✓ /fixtures/{datasource_name}"))
+
+            click.echo(FeedbackManager.success(message="✓ Done!\n"))
+    except Exception as e:
+        click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
+
+
+async def update_resources(
+    tb_client: TinyB,
+    user_token: str,
+    prompt: str,
+    folder: str,
+):
+    datasource_paths = [
+        Path(folder) / "datasources" / f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")
+    ]
+    pipes_paths = [
+        Path(folder) / "endpoints" / f for f in os.listdir(Path(folder) / "endpoints") if f.endswith(".pipe")
+    ]
+    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
+            ]
+        ]
+    )
+    llm = LLM(user_token=user_token, client=tb_client)
+    result = await llm.ask(prompt, system_prompt=update_prompt(resources_xml))
+    result = extract_xml(result, "response")
+    resources = parse_xml(result, "resource")
+    datasources = []
+    pipes = []
+    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)
+
+    for ds in datasources:
+        content = ds["content"].replace("```", "")
+        filename = f"{ds['name']}.datasource"
+        generate_datafile(
+            content,
+            filename=filename,
+            data=None,
+            _format="ndjson",
+            force=True,
+            folder=folder,
+        )
+
+    for pipe in pipes:
+        content = pipe["content"].replace("```", "")
+        generate_pipe_file(pipe["name"], content, folder)
+
+    return len(datasources) > 0
+
+
+def generate_pipe_file(name: str, content: str, folder: str):
+    def is_copy(content: str) -> bool:
+        return re.search(r"TYPE copy", content, re.IGNORECASE) is not None
+
+    def is_materialization(content: str) -> bool:
+        return re.search(r"TYPE materialized", content, re.IGNORECASE) is not None
+
+    def is_sink(content: str) -> bool:
+        return re.search(r"TYPE sink", content, re.IGNORECASE) is not None
+
+    if is_copy(content):
+        pathname = "copies"
+    elif is_materialization(content):
+        pathname = "materializations"
+    elif is_sink(content):
+        pathname = "sinks"
+    else:
+        pathname = "endpoints"
+
+    base = Path(folder) / pathname
+    if not base.exists():
+        base = Path()
+    f = base / (f"{name}.pipe")
+    with open(f"{f}", "w") as file:
+        file.write(content)
+    click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder)))

{tinybird-0.0.1.dev29.dist-info → tinybird-0.0.1.dev30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tinybird
-Version: 0.0.1.dev29
+Version: 0.0.1.dev30
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/cli/introduction.html
 Author: Tinybird
@@ -8,6 +8,7 @@ Author-email: support@tinybird.co
 Requires-Python: >=3.9, <3.13
 Description-Content-Type: text/x-rst
 Requires-Dist: aiofiles (==24.1.0)
+Requires-Dist: anthropic (==0.42.0)
 Requires-Dist: click (<8.2,>=8.1.6)
 Requires-Dist: clickhouse-toolset (==0.33.dev0)
 Requires-Dist: colorama (==0.4.6)