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

tinybird/tb/modules/cli.py

@@ -9,7 +9,7 @@ import os
 import pprint
 from os import getcwd
 from pathlib import Path
-from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 import click
 import humanfriendly
@@ -20,7 +20,8 @@ from tinybird.client import (
     AuthNoTokenException,
     TinyB,
 )
-from tinybird.config import SUPPORTED_CONNECTORS, get_config
+from tinybird.config import get_config
+from tinybird.prompts import ask_prompt
 from tinybird.tb import __cli__
 from tinybird.tb.modules.common import (
     CatchAuthExceptions,
@@ -30,7 +31,6 @@ from tinybird.tb.modules.common import (
     echo_safe_format_table,
     get_current_main_workspace,
     getenv_bool,
-    load_connector_config,
     try_update_config_with_remote,
 )
 from tinybird.tb.modules.config import CLIConfig
@@ -94,7 +94,7 @@ VERSION = f"{__cli__.__version__} (rev {__cli__.__revision__})"
     "--with-headers", help="Flag to enable connector to export with headers", is_flag=True, default=False, hidden=True
 )
 @click.option("--show-tokens", is_flag=True, default=False, help="Enable the output of tokens")
-@click.option("--local", is_flag=True, default=False, help="Run in local mode")
+@click.option("--prod", is_flag=True, default=False, help="Run against production")
 @click.version_option(version=VERSION)
 @click.pass_context
 @coro
@@ -117,14 +117,13 @@ async def cli(
     sf_stage,
     with_headers: bool,
     show_tokens: bool,
-    local: bool,
+    prod: bool,
 ) -> None:
     """
     Use `OBFUSCATE_REGEX_PATTERN` and `OBFUSCATE_PATTERN_SEPARATOR` environment variables to define a regex pattern and a separator (in case of a single string with multiple regex) to obfuscate secrets in the CLI output.
     """
-
     # We need to unpatch for our tests not to break
-    if show_tokens or local or ctx.invoked_subcommand == "build":
+    if show_tokens or not prod or ctx.invoked_subcommand == "build":
         __unpatch_click_output()
     else:
         __patch_click_output()
@@ -211,12 +210,11 @@ async def cli(
 
     logging.debug("debug enabled")
 
-    ctx.ensure_object(dict)["client"] = (
-        await get_tinybird_local_client() if local else _get_tb_client(config.get("token", None), config["host"])
-    )
+    skip_client = ctx.invoked_subcommand in ["login", "workspace"]
+    client = await create_ctx_client(config, prod, skip_client)
 
-    for connector in SUPPORTED_CONNECTORS:
-        load_connector_config(ctx, connector, debug, check_uninstalled=True)
+    if client:
+        ctx.ensure_object(dict)["client"] = client
 
 
 @cli.command(hidden=True)
@@ -489,11 +487,10 @@ async def sql(
 
 @cli.command(hidden=True)
 @click.argument("prompt")
-@click.option("--model", default="gpt-4o", help="The model to use for the LLM")
 @click.option("--folder", default=".", help="The folder to use for the project")
 @coro
-async def llm(prompt: str, model: str, folder: str) -> None:
-    """Run a prompt against the LLM api."""
+async def ask(prompt: str, folder: str) -> None:
+    """Ask things about your data project."""
     try:
         config = CLIConfig.get_project_config(folder)
         user_token = config.get_user_token()
@@ -501,9 +498,39 @@ async def llm(prompt: str, model: str, folder: str) -> None:
         if not user_token:
             raise CLIException("This action requires authentication. Run 'tb login' first.")
 
+        datasource_paths = [
+            Path(folder) / "datasources" / f
+            for f in os.listdir(Path(folder) / "datasources")
+            if f.endswith(".datasource")
+        ]
+        pipe_folders = ["endpoints", "pipes", "copies", "materializations", "sinks"]
+        pipe_paths = [
+            Path(folder) / pipe_folder / f
+            for pipe_folder in pipe_folders
+            if (Path(folder) / pipe_folder).exists()
+            for f in os.listdir(Path(folder) / pipe_folder)
+            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 pipe_paths
+                ]
+            ]
+        )
+
         client = config.get_client()
         llm = LLM(user_token=user_token, client=client)
-        click.echo(await llm.ask(prompt, model=model))
+        click.echo(await llm.ask(system_prompt=ask_prompt(resources_xml), prompt=prompt))
     except Exception as e:
         raise CLIException(FeedbackManager.error_exception(error=e))
 
@@ -543,3 +570,13 @@ def __patch_click_output():
 def __unpatch_click_output():
     click.echo = __old_click_echo
     click.secho = __old_click_secho
+
+
+async def create_ctx_client(config: Dict[str, Any], prod: bool, skip_client: bool):
+    if skip_client:
+        return None
+
+    if prod:
+        return _get_tb_client(config.get("token", None), config["host"])
+
+    return await get_tinybird_local_client()

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
@@ -108,7 +109,7 @@ async def create(
                     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))
+                        response = await llm.ask(system_prompt=mock_prompt(rows), prompt=prompt)
                         sql = extract_xml(response, "sql")
                         sql = sql.split("FORMAT")[0]
                         result = await local_client.query(f"{sql} FORMAT JSON")
@@ -205,7 +206,7 @@ TYPE ENDPOINT
             ]
         )
         llm = LLM(user_token=user_token, client=tb_client)
-        result = await llm.ask(prompt, system_prompt=create_prompt(resources_xml))
+        result = await llm.ask(system_prompt=create_prompt(resources_xml), prompt=prompt)
         result = extract_xml(result, "response")
         resources = parse_xml(result, "resource")
         datasources = []
@@ -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,33 +1,9 @@
-import json
-import urllib.parse
 from copy import deepcopy
-from typing import List
-
-from pydantic import BaseModel
+from typing import Optional
 
 from tinybird.client import TinyB
 
 
-class DataFile(BaseModel):
-    name: str
-    content: str
-
-
-class DataProject(BaseModel):
-    datasources: List[DataFile]
-    pipes: List[DataFile]
-
-
-class TestExpectation(BaseModel):
-    name: str
-    description: str
-    parameters: str
-
-
-class TestExpectations(BaseModel):
-    tests: List[TestExpectation]
-
-
 class LLM:
     def __init__(
         self,
@@ -37,69 +13,26 @@ 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, system_prompt: str, prompt: Optional[str] = None) -> str:
         """
         Calls the model with the given prompt and returns the response.
 
         Args:
+            system_prompt (str): The system prompt to send to the model.
             prompt (str): The user prompt to send to the model.
 
         Returns:
             str: The response from the language model.
         """
-        messages = []
 
-        if system_prompt:
-            messages.append({"role": "user", "content": system_prompt})
+        data = {"system": system_prompt}
 
         if prompt:
-            messages.append({"role": "user", "content": prompt})
+            data["prompt"] = prompt
 
-        data = {
-            "model": model,
-            "messages": messages,
-        }
         response = await self.user_client._req(
             "/v0/llm",
             method="POST",
-            data=json.dumps(data),
-            headers={"Content-Type": "application/json"},
+            data=data,
         )
         return response.get("result", "")
-
-    async def create_project(self, prompt: str) -> DataProject:
-        try:
-            prompt = (
-                prompt
-                + "\n#More extra context\n- If you add some array data type remember that the json path should be like this: `json:$.array_field[:]`"
-            )
-            response = await self.user_client._req(
-                "/v0/llm/create",
-                method="POST",
-                data=f'{{"prompt": {json.dumps(prompt)}}}',
-                headers={"Content-Type": "application/json"},
-            )
-
-            return DataProject.model_validate(response.get("result", {}))
-        except Exception:
-            return DataProject(datasources=[], pipes=[])
-
-    async def generate_sql_sample_data(self, schema: str, rows: int = 20, prompt: str = "") -> str:
-        response = await self.user_client._req(
-            "/v0/llm/mock",
-            method="POST",
-            data=f'{{"schema": "{urllib.parse.quote(schema)}", "rows": {rows}, "context": "{prompt}"}}',
-            headers={"Content-Type": "application/json"},
-        )
-        result = response.get("result", "")
-        return result.replace("elementAt", "arrayElement")
-
-    async def create_tests(self, pipe_content: str, pipe_params: set[str], prompt: str = "") -> TestExpectations:
-        response = await self.user_client._req(
-            "/v0/llm/create/tests",
-            method="POST",
-            data=json.dumps({"pipe_content": pipe_content, "pipe_params": list(pipe_params), "prompt": prompt}),
-            headers={"Content-Type": "application/json"},
-        )
-        result = response.get("result", "")
-        return TestExpectations.model_validate(result)

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

@@ -8,7 +8,7 @@ from tinybird.client import TinyB
 from tinybird.tb.modules.config import CLIConfig
 from tinybird.tb.modules.exceptions import CLIException
 
-TB_IMAGE_NAME = "tinybirdco/tinybird-local:latest"
+TB_IMAGE_NAME = "tinybirdco/tinybird-local:beta"
 TB_CONTAINER_NAME = "tinybird-local"
 TB_LOCAL_PORT = int(os.getenv("TB_LOCAL_PORT", 80))
 TB_LOCAL_HOST = f"http://localhost:{TB_LOCAL_PORT}"

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

@@ -66,7 +66,7 @@ async def mock(datasource: str, rows: int, prompt: str, folder: str) -> None:
         llm = LLM(user_token=user_token, client=user_client)
         tb_client = await get_tinybird_local_client(os.path.abspath(folder))
         prompt = f"<datasource_schema>{datasource_content}</datasource_schema>\n<user_input>{prompt}</user_input>"
-        response = await llm.ask(prompt, system_prompt=mock_prompt(rows))
+        response = await llm.ask(system_prompt=mock_prompt(rows), prompt=prompt)
         sql = extract_xml(response, "sql")
         if os.environ.get("TB_DEBUG", "") != "":
             logging.debug(sql)

tinybird/tb/modules/shell.py

@@ -1,8 +1,10 @@
 import asyncio
 import concurrent.futures
+import glob
 import os
 import subprocess
 import sys
+from pathlib import Path
 from typing import List
 
 import click
@@ -21,17 +23,22 @@ from tinybird.tb.modules.table import format_table
 
 
 class DynamicCompleter(Completer):
-    def __init__(self, datasources: List[str], shared_datasources: List[str], endpoints: List[str], pipes: List[str]):
-        self.datasources = datasources
-        self.shared_datasources = shared_datasources
-        self.endpoints = endpoints
-        self.pipes = pipes
+    def __init__(self, folder: str):
+        self.folder = folder
         self.static_commands = ["create", "mock", "test", "select"]
         self.test_commands = ["create", "run", "update"]
         self.mock_flags = ["--prompt", "--rows"]
         self.common_rows = ["10", "50", "100", "500", "1000"]
         self.sql_keywords = ["select", "from", "where", "group by", "order by", "limit"]
 
+    @property
+    def datasources(self):
+        return [Path(f).stem for f in glob.glob(f"{self.folder}/**/*.datasource", recursive=True)]
+
+    @property
+    def pipes(self):
+        return [Path(f).stem for f in glob.glob(f"{self.folder}/**/*.pipe", recursive=True)]
+
     def get_completions(self, document, complete_event):
         text = document.text_before_cursor.strip()
         words = text.split()
@@ -75,9 +82,9 @@ class DynamicCompleter(Completer):
         if words[-1] == "from" or (
             "from" in words and len(words) > words.index("from") + 1 and text_lower.endswith(" ")
         ):
-            for x in self.datasources + self.shared_datasources:
+            for x in self.datasources:
                 yield Completion(x, start_position=0, display=x, style="class:completion.datasource")
-            for x in self.endpoints + self.pipes:
+            for x in self.pipes:
                 yield Completion(x, start_position=0, display=x, style="class:completion.pipe")
             return
 
@@ -116,7 +123,7 @@ class DynamicCompleter(Completer):
                 yield Completion(cmd, start_position=0, display=cmd, style="class:completion.cmd")
             return
         elif len(words) == 2:
-            for cmd in self.endpoints:
+            for cmd in self.pipes:
                 yield Completion(cmd, start_position=0, display=cmd, style="class:completion.pipe")
             return
 
@@ -130,7 +137,7 @@ class DynamicCompleter(Completer):
                     style="class:completion.cmd",
                 )
 
-        for cmd in self.datasources + self.shared_datasources:
+        for cmd in self.datasources:
             if cmd.startswith(current_word):
                 yield Completion(
                     cmd,
@@ -139,7 +146,7 @@ class DynamicCompleter(Completer):
                     style="class:completion.datasource",
                 )
 
-        for cmd in self.endpoints + self.pipes:
+        for cmd in self.pipes:
             if cmd.startswith(current_word):
                 yield Completion(
                     cmd,
@@ -176,26 +183,14 @@ def _(event):
 
 
 class Shell:
-    def __init__(
-        self,
-        folder: str,
-        client: TinyB,
-        datasources: List[str],
-        shared_datasources: List[str],
-        pipes: List[str],
-        endpoints: List[str],
-    ):
+    def __init__(self, folder: str, client: TinyB):
         self.history = self.get_history()
         self.folder = folder
         self.client = client
-        self.datasources = datasources
-        self.shared_datasources = shared_datasources
-        self.pipes = pipes
-        self.endpoints = endpoints
         self.prompt_message = "\ntb > "
         self.commands = ["create", "mock", "test", "tb", "select"]
         self.session: PromptSession = PromptSession(
-            completer=DynamicCompleter(self.datasources, self.shared_datasources, self.endpoints, self.pipes),
+            completer=DynamicCompleter(folder),
             complete_style=CompleteStyle.COLUMN,
             complete_while_typing=True,
             history=self.history,
@@ -274,7 +269,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("")
@@ -283,10 +278,10 @@ class Shell:
             return
         if arg.startswith("with") or arg.startswith("select"):
             self.run_sql(argline)
-        elif len(arg.split()) == 1 and arg in self.endpoints + self.pipes + self.datasources + self.shared_datasources:
+        elif len(arg.split()) == 1 and arg in self.pipes + self.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:
@@ -310,16 +305,16 @@ class Shell:
             with concurrent.futures.ThreadPoolExecutor() as executor:
                 res = executor.submit(run_query_in_thread).result()
 
-        except Exception as e:
-            click.echo(FeedbackManager.error_exception(error=str(e)))
+                if isinstance(res, dict) and "error" in res:
+                    click.echo(FeedbackManager.error_exception(error=res["error"]))
 
-        if isinstance(res, dict) and "error" in res:
-            click.echo(FeedbackManager.error_exception(error=res["error"]))
+                if isinstance(res, dict) and "data" in res and res["data"]:
+                    print_table_formatted(res, "QUERY")
+                else:
+                    click.echo(FeedbackManager.info_no_rows())
 
-        if isinstance(res, dict) and "data" in res and res["data"]:
-            print_table_formatted(res, "QUERY")
-        else:
-            click.echo(FeedbackManager.info_no_rows())
+        except Exception as e:
+            click.echo(FeedbackManager.error_exception(error=str(e)))
 
     def reprint_prompt(self):
         click.echo(f"{bcolors.OKGREEN}{self.prompt_message}{bcolors.ENDC}", nl=False)