tinybird 0.0.1.dev20__tar.gz → 0.0.1.dev22__tar.gz

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tinybird
-Version: 0.0.1.dev20
+Version: 0.0.1.dev22
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/cli/introduction.html
 Author: Tinybird

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/prompts.py

@@ -122,7 +122,7 @@ FROM numbers({rows})
 
 """
 
-create_test_calls_prompt = """
+create_test_prompt = """
 You are a Tinybird expert. You will be given a pipe endpoint containing different nodes with SQL and Tinybird templating syntax. You will generate URLs to test it with different parameters combinations.
 
 <test>
@@ -141,6 +141,6 @@ You are a Tinybird expert. You will be given a pipe endpoint containing differen
     - The test command can have as many parameters as are needed to test the pipe.
     - The parameter within Tinybird templating syntax looks like this one {{String(my_param_name, default_value)}}.
     - If there are no parameters in the , you can omit parametrs and generate a single test command.
-    - Extra context: {context}
+    - Extra context: {prompt}
 </instructions>
 """

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/tb/__cli__.py

@@ -4,5 +4,5 @@ __description__ = 'Tinybird Command Line Tool'
 __url__ = 'https://www.tinybird.co/docs/cli/introduction.html'
 __author__ = 'Tinybird'
 __author_email__ = 'support@tinybird.co'
-__version__ = '0.0.1.dev20'
-__revision__ = '6125327'
+__version__ = '0.0.1.dev22'
+__revision__ = 'fad51ec'

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/tb/modules/create.py

@@ -56,10 +56,11 @@ async def create(
     folder = folder or getcwd()
     try:
         config = CLIConfig.get_project_config(folder)
+        user_token: Optional[str] = None
 
         if prompt:
-            user_token = config.get_user_token()
             try:
+                user_token = config.get_user_token()
                 if not user_token:
                     raise CLIException("No user token found")
                 await check_user_token(ctx, token=user_token)
@@ -69,7 +70,7 @@ async def create(
 
         tb_client = config.get_client()
         click.echo(FeedbackManager.gray(message="Creating new project structure..."))
-        await project_create(tb_client, data, prompt, folder)
+        await project_create(tb_client, user_token, data, prompt, folder)
         click.echo(FeedbackManager.success(message="✓ Scaffolding completed!\n"))
 
         click.echo(FeedbackManager.gray(message="\nCreating CI/CD files for GitHub and GitLab..."))
@@ -145,11 +146,11 @@ async def create(
             )
             click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
             persist_fixture(fixture_name, data_content)
-        elif prompt:
+        elif prompt and user_token:
             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(client=tb_client)
+                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
@@ -171,6 +172,7 @@ async def create(
 
 async def project_create(
     client: TinyB,
+    user_token: Optional[str],
     data: Optional[str],
     prompt: Optional[str],
     folder: str,
@@ -203,9 +205,9 @@ TYPE ENDPOINT
             """,
             folder,
         )
-    elif prompt:
+    elif prompt and user_token:
         try:
-            llm = LLM(client=client)
+            llm = LLM(user_token=user_token, client=client)
             result = await llm.create_project(prompt)
             for ds in result.datasources:
                 content = ds.content.replace("```", "")

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/tb/modules/datafile/build.py

@@ -689,6 +689,148 @@ class GraphDependencies:
     all_resources: Dict[str, Dict[str, Any]]
 
 
+async def process(
+    filename: str,
+    tb_client: TinyB,
+    deps: List[str],
+    dep_map: Dict[str, Any],
+    to_run: Dict[str, Any],
+    vendor_paths: Optional[List[Tuple[str, str]]] = None,
+    skip_connectors: bool = False,
+    current_ws: Optional[Dict[str, Any]] = None,
+    changed: Optional[Dict[str, Any]] = None,
+    fork_downstream: Optional[bool] = False,
+    is_internal: Optional[bool] = False,
+    dir_path: Optional[str] = None,
+    verbose: bool = False,
+    embedded_datasources: Optional[Dict[str, Any]] = None,
+):
+    name, kind = filename.rsplit(".", 1)
+    warnings = []
+
+    try:
+        res = await process_file(
+            filename,
+            tb_client,
+            skip_connectors=skip_connectors,
+            current_ws=current_ws,
+        )
+    except click.ClickException as e:
+        raise e
+    except IncludeFileNotFoundException as e:
+        raise click.ClickException(FeedbackManager.error_deleted_include(include_file=str(e), filename=filename))
+    except Exception as e:
+        raise click.ClickException(str(e))
+
+    for r in res:
+        resource_name = r["resource_name"]
+        warnings = r.get("warnings", [])
+        if (
+            changed
+            and resource_name in changed
+            and (not changed[resource_name] or changed[resource_name] in ["shared", "remote"])
+        ):
+            continue
+
+        if (
+            fork_downstream
+            and r.get("resource", "") == "pipes"
+            and any(["engine" in x.get("params", {}) for x in r.get("nodes", [])])
+        ):
+            raise click.ClickException(FeedbackManager.error_forkdownstream_pipes_with_engine(pipe=resource_name))
+
+        to_run[resource_name] = r
+        file_deps = 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)
+                if f:
+                    dep_list.append(f.rsplit(".", 1)[0])
+                if ds:
+                    ds_fn = ds["resource_name"]
+                    prev = to_run.get(ds_fn, {})
+                    to_run[ds_fn] = deepcopy(r)
+                    try:
+                        to_run[ds_fn]["deps"] = list(
+                            set(to_run[ds_fn].get("deps", []) + prev.get("deps", []) + [resource_name])
+                        )
+                    except ValueError:
+                        pass
+                    embedded_datasources[x] = to_run[ds_fn]
+                else:
+                    e_ds = embedded_datasources.get(x, None)
+                    if e_ds:
+                        dep_list.append(e_ds["resource_name"])
+
+        dep_map[resource_name] = set(dep_list)
+    return os.path.basename(name), warnings
+
+
+async def get_processed(
+    filenames: Iterable[str],
+    changed: Optional[Dict[str, Any]] = None,
+    verbose: bool = False,
+    deps: List[str] = [],
+    dep_map: Dict[str, Any] = {},
+    to_run: Dict[str, Any] = {},
+    vendor_paths: Optional[List[Tuple[str, str]]] = None,
+    processed: Set[str] = set(),
+    tb_client: TinyB = None,
+    skip_connectors: bool = False,
+    current_ws: Optional[Dict[str, Any]] = None,
+    fork_downstream: Optional[bool] = False,
+    is_internal: Optional[bool] = False,
+    dir_path: Optional[str] = None,
+    embedded_datasources: Optional[Dict[str, Any]] = None,
+):
+    for filename in filenames:
+        # just process changed filenames (tb deploy and --only-changes)
+        if changed:
+            resource = Path(filename).resolve().stem
+            if resource in changed and (not changed[resource] or changed[resource] in ["shared", "remote"]):
+                continue
+        if os.path.isdir(filename):
+            await get_processed(filenames=get_project_filenames(filename))
+        else:
+            if verbose:
+                click.echo(FeedbackManager.info_processing_file(filename=filename))
+
+            if ".incl" in filename:
+                click.echo(FeedbackManager.warning_skipping_include_file(file=filename))
+
+            name, warnings = await process(
+                filename=filename,
+                tb_client=tb_client,
+                deps=deps,
+                dep_map=dep_map,
+                to_run=to_run,
+                vendor_paths=vendor_paths,
+                skip_connectors=skip_connectors,
+                current_ws=current_ws,
+                changed=changed,
+                fork_downstream=fork_downstream,
+                is_internal=is_internal,
+                dir_path=dir_path,
+                verbose=verbose,
+                embedded_datasources=embedded_datasources,
+            )
+            processed.add(name)
+
+            if verbose:
+                if len(warnings) == 1:
+                    click.echo(FeedbackManager.warning_pipe_restricted_param(word=warnings[0]))
+                elif len(warnings) > 1:
+                    click.echo(
+                        FeedbackManager.warning_pipe_restricted_params(
+                            words=", ".join(["'{}'".format(param) for param in warnings[:-1]]),
+                            last_word=warnings[-1],
+                        )
+                    )
+
+
 async def build_graph(
     filenames: Iterable[str],
     tb_client: TinyB,
@@ -740,109 +882,24 @@ async def build_graph(
         all_dep_map = all_dependencies_graph.dep_map
         all_resources = all_dependencies_graph.to_run
 
-    async def process(
-        filename: str,
-        deps: List[str],
-        dep_map: Dict[str, Any],
-        to_run: Dict[str, Any],
-        vendor_paths: Optional[List[Tuple[str, str]]],
-    ):
-        name, kind = filename.rsplit(".", 1)
-        warnings = []
-
-        try:
-            res = await process_file(
-                filename,
-                tb_client,
-                skip_connectors=skip_connectors,
-                current_ws=current_ws,
-            )
-        except click.ClickException as e:
-            raise e
-        except IncludeFileNotFoundException as e:
-            raise click.ClickException(FeedbackManager.error_deleted_include(include_file=str(e), filename=filename))
-        except Exception as e:
-            raise click.ClickException(str(e))
-
-        for r in res:
-            resource_name = r["resource_name"]
-            warnings = r.get("warnings", [])
-            if (
-                changed
-                and resource_name in changed
-                and (not changed[resource_name] or changed[resource_name] in ["shared", "remote"])
-            ):
-                continue
-
-            if (
-                fork_downstream
-                and r.get("resource", "") == "pipes"
-                and any(["engine" in x.get("params", {}) for x in r.get("nodes", [])])
-            ):
-                raise click.ClickException(FeedbackManager.error_forkdownstream_pipes_with_engine(pipe=resource_name))
-
-            to_run[resource_name] = r
-            file_deps = 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)
-                    if f:
-                        dep_list.append(f.rsplit(".", 1)[0])
-                    if ds:
-                        ds_fn = ds["resource_name"]
-                        prev = to_run.get(ds_fn, {})
-                        to_run[ds_fn] = deepcopy(r)
-                        try:
-                            to_run[ds_fn]["deps"] = list(
-                                set(to_run[ds_fn].get("deps", []) + prev.get("deps", []) + [resource_name])
-                            )
-                        except ValueError:
-                            pass
-                        embedded_datasources[x] = to_run[ds_fn]
-                    else:
-                        e_ds = embedded_datasources.get(x, None)
-                        if e_ds:
-                            dep_list.append(e_ds["resource_name"])
-
-            dep_map[resource_name] = set(dep_list)
-        return os.path.basename(name), warnings
-
     processed = set()
 
-    async def get_processed(filenames: Iterable[str]):
-        for filename in filenames:
-            # just process changed filenames (tb deploy and --only-changes)
-            if changed:
-                resource = Path(filename).resolve().stem
-                if resource in changed and (not changed[resource] or changed[resource] in ["shared", "remote"]):
-                    continue
-            if os.path.isdir(filename):
-                await get_processed(filenames=get_project_filenames(filename))
-            else:
-                if verbose:
-                    click.echo(FeedbackManager.info_processing_file(filename=filename))
-
-                if ".incl" in filename:
-                    click.echo(FeedbackManager.warning_skipping_include_file(file=filename))
-
-                name, warnings = await process(filename, deps, dep_map, to_run, vendor_paths)
-                processed.add(name)
-
-                if verbose:
-                    if len(warnings) == 1:
-                        click.echo(FeedbackManager.warning_pipe_restricted_param(word=warnings[0]))
-                    elif len(warnings) > 1:
-                        click.echo(
-                            FeedbackManager.warning_pipe_restricted_params(
-                                words=", ".join(["'{}'".format(param) for param in warnings[:-1]]),
-                                last_word=warnings[-1],
-                            )
-                        )
-
-    await get_processed(filenames=filenames)
+    await get_processed(
+        filenames=filenames,
+        tb_client=tb_client,
+        changed=changed,
+        deps=deps,
+        dep_map=dep_map,
+        to_run=to_run,
+        vendor_paths=vendor_paths,
+        processed=processed,
+        skip_connectors=skip_connectors,
+        current_ws=current_ws,
+        fork_downstream=fork_downstream,
+        is_internal=is_internal,
+        dir_path=dir_path,
+        embedded_datasources=embedded_datasources,
+    )
 
     if process_dependencies:
         if only_changes:
@@ -878,7 +935,21 @@ async def build_graph(
                                 click.echo(FeedbackManager.info_skipping_resource(resource=processed_filename))
                                 continue
                             click.echo(FeedbackManager.info_processing_file(filename=processed_filename))
-                        await process(str(f), deps, dep_map, to_run, vendor_paths)
+                        await process(
+                            filename=str(f),
+                            tb_client=tb_client,
+                            deps=deps,
+                            dep_map=dep_map,
+                            to_run=to_run,
+                            vendor_paths=vendor_paths,
+                            skip_connectors=skip_connectors,
+                            current_ws=current_ws,
+                            fork_downstream=fork_downstream,
+                            is_internal=is_internal,
+                            dir_path=dir_path,
+                            verbose=verbose,
+                            embedded_datasources=embedded_datasources,
+                        )
 
     return GraphDependencies(dep_map, to_run, all_dep_map, all_resources)
 

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/tb/modules/llm.py

@@ -8,8 +8,6 @@ from openai import OpenAI
 from pydantic import BaseModel
 
 from tinybird.client import TinyB
-from tinybird.prompts import create_test_calls_prompt
-from tinybird.tb.modules.config import CLIConfig
 
 
 class DataFile(BaseModel):
@@ -33,13 +31,10 @@ class TestExpectations(BaseModel):
 
 
 class LLM:
-    def __init__(self, client: TinyB, api_key: Optional[str] = None):
-        self.client = client
-        user_token = CLIConfig.get_project_config().get_user_token()
-        user_client = deepcopy(client)
-        if user_token:
-            user_client.token = user_token
-        self.user_client = user_client
+    def __init__(self, user_token: str, client: TinyB, api_key: Optional[str] = None):
+        self.user_client = deepcopy(client)
+        self.user_client.token = user_token
+
         self.openai = OpenAI(api_key=api_key) if api_key else None
 
     async def _execute(self, action_fn: Callable[[], Awaitable[str]], checker_fn: Callable[[str], bool]):
@@ -79,20 +74,12 @@ class LLM:
         result = response.get("result", "")
         return result.replace("elementAt", "arrayElement")
 
-    async def create_test_commands(
-        self, pipe_content: str, pipe_params: set[str], context: Optional[str] = None
-    ) -> TestExpectations:
-        if not self.openai:
-            raise ValueError("OpenAI API key is not set")
-
-        completion = self.openai.beta.chat.completions.parse(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": create_test_calls_prompt.format(context=context or "")},
-                {"role": "user", "content": f"Pipe content: {pipe_content}\nPipe params: {pipe_params}"},
-            ],
-            temperature=0.2,
-            seed=42,
-            response_format=TestExpectations,
+    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"},
         )
-        return completion.choices[0].message.parsed or TestExpectations(tests=[])
+        result = response.get("result", "")
+        return TestExpectations.model_validate(result)

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/tb/modules/local.py

@@ -59,18 +59,15 @@ def start_tinybird_local(
         )
 
     click.echo(FeedbackManager.info(message="* Waiting for Tinybird Local to be ready..."))
-    for attempt in range(10):
-        try:
-            run = container.exec_run("tb --no-version-warning sql 'SELECT 1 AS healthcheck' --format json").output
-            # dont parse the json as docker sometimes returns warning messages
-            # todo: rafa, make this rigth
-            if b'"healthcheck": 1' in run:
-                break
-            raise RuntimeError("Unexpected response from Tinybird")
-        except Exception:
-            if attempt == 9:  # Last attempt
-                raise CLIException("Tinybird Local not ready yet. Please try again in a few seconds.")
-            time.sleep(5)  # Wait 5 seconds before retrying
+    while True:
+        container.reload()  # Refresh container attributes
+        health = container.attrs.get("State", {}).get("Health", {}).get("Status")
+        if health == "healthy":
+            break
+        if health == "unhealthy":
+            raise CLIException("Tinybird Local is unhealthy. Please try running `tb local restart` in a few seconds.")
+
+        time.sleep(5)
 
     # Remove tinybird-local dangling images to avoid running out of disk space
     images = docker_client.images.list(name=re.sub(r":.*$", "", TB_IMAGE_NAME), all=True, filters={"dangling": True})

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/tb/modules/mock.py

@@ -41,7 +41,6 @@ async def mock(ctx: click.Context, datasource: str, rows: int, prompt: str, fold
         if not prompt:
             # load the prompt from the fixture.prompt file if it exists
             if prompt_path.exists():
-                click.echo(FeedbackManager.gray(message=f"Using prompt for {prompt_path}..."))
                 prompt = prompt_path.read_text()
         else:
             click.echo(FeedbackManager.gray(message=f"Overriding prompt for {datasource_name}..."))
@@ -61,7 +60,7 @@ async def mock(ctx: click.Context, datasource: str, rows: int, prompt: str, fold
             click.echo(FeedbackManager.error(message="This action requires authentication. Run 'tb login' first."))
             return
 
-        llm = LLM(client=user_client)
+        llm = LLM(user_token=user_token, client=user_client)
         tb_client = await get_tinybird_local_client(os.path.abspath(folder))
         sql = await llm.generate_sql_sample_data(datasource_content, rows=rows, prompt=prompt)
         if os.environ.get("TB_DEBUG", "") != "":

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird/tb/modules/test.py

@@ -89,11 +89,12 @@ async def test_create(pipe: str, prompt: Optional[str], folder: str) -> None:
         pipe_nodes = await client._req(f"/v0/pipes/{pipe_name}")
         pipe_params = set([param["name"] for node in pipe_nodes["nodes"] for param in node["params"]])
 
-        llm_config = CLIConfig.get_llm_config()
-        llm = LLM(client=client, api_key=llm_config["api_key"])
+        config = CLIConfig.get_project_config(folder)
+        user_token = config.get_user_token()
+        llm = LLM(user_token=user_token, client=config.get_client())
 
-        test_expectations = await llm.create_test_commands(
-            pipe_content=pipe_content, pipe_params=pipe_params, context=prompt
+        test_expectations = await llm.create_tests(
+            pipe_content=pipe_content, pipe_params=pipe_params, prompt=prompt or ""
         )
         valid_test_expectations = []
         for test in test_expectations.tests:

{tinybird-0.0.1.dev20 → tinybird-0.0.1.dev22}/tinybird.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tinybird
-Version: 0.0.1.dev20
+Version: 0.0.1.dev22
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/cli/introduction.html
 Author: Tinybird