tinybird 0.0.1.dev300__py3-none-any.whl → 0.0.1.dev302__py3-none-any.whl

tinybird/connectors.py

@@ -369,13 +369,7 @@ class Snowflake(Connector):
             the_type = "String"
             if t.startswith("NUMBER"):
                 the_type = "Int32"
-            if (
-                t.startswith("FLOAT")
-                or t.startswith("DOUBLE")
-                or t.startswith("REAL")
-                or t.startswith("NUMERIC")
-                or t.startswith("DECIMAL")
-            ):
+            if t.startswith(("FLOAT", "DOUBLE", "REAL", "NUMERIC", "DECIMAL")):
                 the_type = "Float32"
             if t == "DATE":
                 the_type = "Date"

tinybird/datafile/common.py

@@ -2051,11 +2051,7 @@ def parse(
                 lexer = list(sa)
             if lexer:
                 cmd, args = lexer[0], lexer[1:]
-                if (
-                    parser_state.multiline
-                    and cmd.lower() in cmds
-                    and not (line.startswith(" ") or line.startswith("\t"))
-                ):
+                if parser_state.multiline and cmd.lower() in cmds and not line.startswith((" ", "\t")):
                     cmds[parser_state.command](
                         parser_state.multiline_string,
                         lineno=lineno,
@@ -2508,7 +2504,7 @@ def is_file_a_datasource(filename: str) -> bool:
 
         for line in lines:
             trimmed_line = line.strip().lower()
-            if trimmed_line.startswith("schema") or trimmed_line.startswith("engine"):
+            if trimmed_line.startswith(("schema", "engine")):
                 return True
 
     return False

tinybird/service_datasources.py

@@ -446,7 +446,6 @@ def get_tinybird_service_datasources() -> List[Dict[str, Any]]:
                 {"name": "workspace_name", "type": "String"},
                 {"name": "user_email", "type": "String"},
                 {"name": "request_id", "type": "String"},
-                {"name": "model", "type": "LowCardinality(String)"},
                 {"name": "prompt_tokens", "type": "UInt32"},
                 {"name": "completion_tokens", "type": "UInt32"},
                 {"name": "total_tokens", "type": "UInt32"},
@@ -927,7 +926,6 @@ def get_organization_service_datasources() -> List[Dict[str, Any]]:
                 {"name": "workspace_name", "type": "String"},
                 {"name": "user_email", "type": "String"},
                 {"name": "request_id", "type": "String"},
-                {"name": "model", "type": "LowCardinality(String)"},
                 {"name": "prompt_tokens", "type": "UInt32"},
                 {"name": "completion_tokens", "type": "UInt32"},
                 {"name": "total_tokens", "type": "UInt32"},

tinybird/tb/__cli__.py

@@ -4,5 +4,5 @@ __description__ = 'Tinybird Command Line Tool'
 __url__ = 'https://www.tinybird.co/docs/forward/commands'
 __author__ = 'Tinybird'
 __author_email__ = 'support@tinybird.co'
-__version__ = '0.0.1.dev300'
-__revision__ = '9791d55'
+__version__ = '0.0.1.dev302'
+__revision__ = 'cf4d28e'

tinybird/tb/client.py

@@ -1390,7 +1390,3 @@ class TinyB:
 
     def delete_tag(self, name: str):
         self._req(f"/v0/tags/{name}", method="DELETE")
-
-    def explore_data(self, prompt: str) -> str:
-        params = urlencode({"prompt": prompt, "host": self.host, "origin": "cli"})
-        return self._req(f"/v1/agents/explore?{params}")

tinybird/tb/modules/agent/agent.py

@@ -6,8 +6,7 @@ import sys
 import urllib.parse
 import uuid
 from functools import partial
-from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import click
 import humanfriendly
@@ -27,7 +26,6 @@ from tinybird.tb.modules.agent.file_agent import FileAgent
 from tinybird.tb.modules.agent.memory import (
     clear_history,
     clear_messages,
-    get_last_messages_from_last_user_prompt,
     save_messages,
 )
 from tinybird.tb.modules.agent.mock_agent import MockAgent
@@ -67,16 +65,17 @@ from tinybird.tb.modules.common import (
     echo_safe_humanfriendly_tables_format_pretty_table,
     get_region_from_host,
     get_regions,
+    sys_exit,
     update_cli,
 )
 from tinybird.tb.modules.config import CLIConfig
 from tinybird.tb.modules.deployment_common import create_deployment
-from tinybird.tb.modules.exceptions import CLIBuildException, CLIDeploymentException, CLIMockException
+from tinybird.tb.modules.exceptions import CLIBuildException, CLIDeploymentException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.llm import LLM
 from tinybird.tb.modules.local_common import get_tinybird_local_client
 from tinybird.tb.modules.login_common import login
-from tinybird.tb.modules.mock_common import append_mock_data, create_mock_data
+from tinybird.tb.modules.mock_common import append_mock_data
 from tinybird.tb.modules.project import Project
 from tinybird.tb.modules.test_common import run_tests as run_tests_common
 
@@ -91,6 +90,7 @@ class TinybirdAgent:
         project: Project,
         dangerously_skip_permissions: bool,
         prompt_mode: bool,
+        feature: Optional[str] = None,
     ):
         self.token = token
         self.user_token = user_token
@@ -100,14 +100,12 @@ class TinybirdAgent:
         self.project = project
         self.thinking_animation = ThinkingAnimation()
         self.confirmed_plan_id: Optional[str] = None
-        if prompt_mode:
-            self.messages: list[ModelMessage] = get_last_messages_from_last_user_prompt()
-        else:
-            self.messages = []
+        self.feature = feature
+        self.messages: list[ModelMessage] = []
         cli_config = CLIConfig.get_project_config()
         regions = get_regions(cli_config)
         self.agent = Agent(
-            model=create_model(user_token, host, workspace_id),
+            model=create_model(user_token, host, workspace_id, feature=feature),
             deps_type=TinybirdAgentContext,
             instructions=[agent_system_prompt],
             tools=[
@@ -381,7 +379,6 @@ class TinybirdAgent:
             build_project=partial(build_project, project=project, config=config),
             deploy_project=partial(deploy_project, project=project, config=config),
             deploy_check_project=partial(deploy_check_project, project=project, config=config),
-            mock_data=partial(mock_data, project=project, config=config),
             append_data_local=partial(append_data_local, config=config),
             append_data_cloud=partial(append_data_cloud, config=config),
             analyze_fixture=partial(analyze_fixture, config=config),
@@ -428,7 +425,6 @@ class TinybirdAgent:
         save_messages(new_messages)
         self.thinking_animation.stop()
         click.echo(result.output)
-        self.echo_usage(config)
 
     async def run_iter(self, user_prompt: str, config: dict[str, Any], run_id: Optional[str] = None) -> None:
         model = create_model(self.user_token, self.host, self.workspace_id, run_id=run_id)
@@ -455,34 +451,67 @@ class TinybirdAgent:
             self.messages.extend(new_messages)
             save_messages(new_messages)
             self.thinking_animation.stop()
-            self.echo_usage(config)
 
-    def echo_usage(self, config: dict[str, Any]) -> None:
+    def echo_usage(self, config: dict[str, Any], show_credits: bool = False) -> None:
         try:
             client = _get_tb_client(config["user_token"], config["host"])
             workspace_id = config.get("id", "")
             workspace = client.workspace(workspace_id, with_organization=True, version="v1")
+            is_free_plan = workspace["organization"]["plan"].get("billing") == "free_shared_infrastructure_usage"
+
+            if not is_free_plan and not show_credits:
+                return
+
             limits_data = client.organization_limits(workspace["organization"]["id"])
             llm_usage_limits = limits_data.get("limits", {}).get("llm_usage", {})
             current_llm_usage = llm_usage_limits.get("quantity") or 0
             llm_usage = llm_usage_limits.get("max") or 0
             remaining_credits = round(max(llm_usage - current_llm_usage, 0), 2)
             current_llm_usage = round(min(llm_usage, current_llm_usage), 2)
+
             if not llm_usage:
                 return
-            warning_threshold = llm_usage * 0.8
-            message_color = FeedbackManager.warning if current_llm_usage >= warning_threshold else FeedbackManager.gray
+
+            def get_message(current_llm_usage, llm_usage: int) -> tuple[Callable[..., str], str, bool]:
+                warning_threshold = llm_usage * 0.8
+                ui_host = get_display_cloud_host(config["host"])
+
+                if is_free_plan:
+                    upgrade_link = f"{ui_host}/organizations/{workspace['organization']['name']}/upgrade?from=agent"
+                    if current_llm_usage >= llm_usage:
+                        return (
+                            FeedbackManager.error,
+                            f" You have reached the maximum number of credits. Please upgrade to continue using Tinybird Code: {upgrade_link}",
+                            True,
+                        )
+                    if current_llm_usage >= warning_threshold:
+                        return (
+                            FeedbackManager.warning,
+                            f" You are reaching the maximum number of credits. Please upgrade to continue using Tinybird Code: {upgrade_link}",
+                            False,
+                        )
+                return FeedbackManager.gray, "", False
+
+            message_color, upgrade_message, should_exit = get_message(current_llm_usage, llm_usage)
             click.echo(
                 message_color(
-                    message=f"{remaining_credits} credits left ({current_llm_usage}/{llm_usage}). You can continue using Tinybird Code. Limits will be enforced soon."
+                    message=f"{remaining_credits} credits left ({current_llm_usage}/{llm_usage}).{upgrade_message}"
                 )
             )
+
+            if should_exit:
+                sys_exit("tinybird_code_error", "Maximum number of credits reached")
+
         except Exception:
             pass
 
 
 def run_agent(
-    config: dict[str, Any], project: Project, dangerously_skip_permissions: bool, prompt: Optional[str] = None
+    config: dict[str, Any],
+    project: Project,
+    dangerously_skip_permissions: bool,
+    prompt: Optional[str] = None,
+    feature: Optional[str] = None,
 ):
     if not prompt:
         latest_version = CheckPypi().get_latest_version()
@@ -497,7 +526,10 @@ def run_agent(
             )
             if yes:
                 update_cli()
-    click.echo(FeedbackManager.highlight(message="» Initializing Tinybird Code..."))
+
+    if not prompt:
+        click.echo(FeedbackManager.highlight(message="» Initializing Tinybird Code..."))
+
     token = config.get("token", "")
     host = config.get("host", "")
     user_token = config.get("user_token", "")
@@ -556,6 +588,7 @@ def run_agent(
             project,
             dangerously_skip_permissions,
             prompt_mode,
+            feature,
         )
 
         # Print mode: run once with the provided prompt and exit
@@ -577,7 +610,6 @@ def run_agent(
             )
         )
         agent.echo_usage(config)
-        click.echo()
 
     except Exception as e:
         click.echo(FeedbackManager.error(message=f"Failed to initialize agent: {e}"))
@@ -630,6 +662,9 @@ def run_agent(
                 elif user_input.lower() == "/help":
                     subprocess.run(["tb", "--help"], check=True)
                     continue
+                elif user_input.lower() == "/usage":
+                    agent.echo_usage(config, show_credits=True)
+                    continue
                 elif user_input.strip() == "":
                     continue
                 else:
@@ -732,36 +767,6 @@ def append_data_cloud(config: dict[str, Any], datasource_name: str, path: str) -
     append_mock_data(client, datasource_name, path)
 
 
-def mock_data(
-    config: dict[str, Any],
-    project: Project,
-    datasource_name: str,
-    data_format: str,
-    rows: int,
-    context: Optional[str] = None,
-) -> list[dict[str, Any]]:
-    client = get_tinybird_local_client(config, test=False, silent=False)
-    cli_config = CLIConfig.get_project_config()
-    datasource_path = project.get_resource_path(datasource_name, "datasource")
-
-    if not datasource_path:
-        raise CLIMockException(f"Datasource {datasource_name} not found")
-
-    datasource_content = Path(datasource_path).read_text()
-    return create_mock_data(
-        datasource_name,
-        datasource_content,
-        rows,
-        context or "",
-        cli_config,
-        config,
-        cli_config.get_user_token() or "",
-        client,
-        data_format,
-        project.folder,
-    )
-
-
 def analyze_fixture(config: dict[str, Any], fixture_path: str, format: str = "json") -> dict[str, Any]:
     local_client = get_tinybird_local_client(config, test=False, silent=True)
     meta, _data = _analyze(fixture_path, local_client, format)

tinybird/tb/modules/agent/models.py

@@ -8,7 +8,7 @@ from pydantic_ai.retries import AsyncTenacityTransport, wait_retry_after
 from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 
-def create_retrying_client(token: str, workspace_id: str):
+def create_retrying_client(token: str, workspace_id: str, feature: Optional[str] = None):
     """Create a client with smart retry handling for multiple error types."""
 
     def should_retry_status(response):
@@ -29,7 +29,10 @@ def create_retrying_client(token: str, workspace_id: str):
         ),
         validate_response=should_retry_status,
     )
-    return AsyncClient(transport=transport, params={"token": token, "workspace_id": workspace_id})
+    params = {"token": token, "workspace_id": workspace_id}
+    if feature:
+        params["feature"] = feature
+    return AsyncClient(transport=transport, params=params)
 
 
 def create_model(
@@ -38,6 +41,7 @@ def create_model(
     workspace_id: str,
     model: AnthropicModelName = "claude-4-sonnet-20250514",
     run_id: Optional[str] = None,
+    feature: Optional[str] = None,
 ):
     default_headers = {}
     if run_id:

tinybird/tb/modules/agent/utils.py

@@ -44,7 +44,6 @@ class TinybirdAgentContext(BaseModel):
     build_project_test: Callable[..., None]
     deploy_project: Callable[..., None]
     deploy_check_project: Callable[[], None]
-    mock_data: Callable[..., list[dict[str, Any]]]
     append_data_local: Callable[..., None]
     append_data_cloud: Callable[..., None]
     analyze_fixture: Callable[..., dict[str, Any]]
@@ -512,7 +511,7 @@ def create_terminal_box(content: str, new_content: Optional[str] = None, title:
                     new_line_num = int(match.group(2))
                     old_index = old_line_num - 1
                     new_index = new_line_num - 1
-            elif line.startswith("---") or line.startswith("+++"):
+            elif line.startswith(("---", "+++")):
                 # Skip file headers
                 pass
             elif line.startswith(" "):

tinybird/tb/modules/build.py

@@ -8,7 +8,7 @@ from urllib.parse import urlencode
 
 import click
 
-import tinybird.context as context
+from tinybird import context
 from tinybird.datafile.exceptions import ParseException
 from tinybird.datafile.parse_datasource import parse_datasource
 from tinybird.datafile.parse_pipe import parse_pipe
@@ -140,7 +140,7 @@ def dev_cloud(
     context.disable_template_security_validation.set(True)
 
     def process(filenames: List[str], watch: bool = False):
-        datafiles = [f for f in filenames if f.endswith(".datasource") or f.endswith(".pipe")]
+        datafiles = [f for f in filenames if f.endswith((".datasource", ".pipe"))]
         if len(datafiles) > 0:
             check_filenames(filenames=datafiles)
             folder_playground(

tinybird/tb/modules/build_common.py

@@ -51,7 +51,7 @@ def process(
             return build_status.error
         else:
             build_status.building = True
-    if file_changed and (file_changed.endswith(FixtureExtension.NDJSON) or file_changed.endswith(FixtureExtension.CSV)):
+    if file_changed and file_changed.endswith((FixtureExtension.NDJSON, FixtureExtension.CSV)):
         rebuild_fixture(project, tb_client, file_changed)
         if build_status:
             build_status.building = False
@@ -61,7 +61,7 @@ def process(
         if build_status:
             build_status.building = False
             build_status.error = None
-    elif file_changed and (file_changed.endswith(".env.local") or file_changed.endswith(".env")):
+    elif file_changed and file_changed.endswith((".env.local", ".env")):
         if build_status:
             build_status.building = False
             build_status.error = None

tinybird/tb/modules/cli.py

@@ -302,7 +302,7 @@ def sql(
             )
 
             query = ""
-            for _, elem in dependencies_graph.to_run.items():
+            for elem in dependencies_graph.to_run.values():
                 for _node in elem["nodes"]:
                     if _node["params"]["name"].lower() == node.lower():
                         query = "".join(_node["sql"])
@@ -488,7 +488,6 @@ def create_ctx_client(
         "logout",
         "update",
         "upgrade",
-        "create",
         "info",
         "tag",
         "push",
@@ -502,7 +501,7 @@ def create_ctx_client(
         return None
 
     commands_always_cloud = ["infra"]
-    commands_always_local = ["build", "dev"]
+    commands_always_local = ["build", "dev", "create"]
     command_always_test = ["test"]
 
     if (

tinybird/tb/modules/common.py

@@ -908,7 +908,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"
@@ -916,7 +916,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"