tinybird 0.0.1.dev287__tar.gz → 0.0.1.dev288__tar.gz

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tinybird
-Version: 0.0.1.dev287
+Version: 0.0.1.dev288
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/forward/commands
 Author: Tinybird

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/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.dev287'
-__revision__ = 'f35d4e1'
+__version__ = '0.0.1.dev288'
+__revision__ = '6cd5b6c'

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/agent/agent.py

@@ -1,4 +1,5 @@
 import asyncio
+import hashlib
 import shlex
 import subprocess
 import sys
@@ -47,7 +48,7 @@ from tinybird.tb.modules.agent.tools.deploy_check import deploy_check
 from tinybird.tb.modules.agent.tools.diff_resource import diff_resource
 from tinybird.tb.modules.agent.tools.get_endpoint_stats import get_endpoint_stats
 from tinybird.tb.modules.agent.tools.get_openapi_definition import get_openapi_definition
-from tinybird.tb.modules.agent.tools.plan import plan
+from tinybird.tb.modules.agent.tools.plan import complete_plan, plan
 from tinybird.tb.modules.agent.tools.secret import create_or_update_secrets
 from tinybird.tb.modules.agent.utils import AgentRunCancelled, TinybirdAgentContext, show_confirmation, show_input
 from tinybird.tb.modules.build_common import process as build_process
@@ -89,6 +90,7 @@ class TinybirdAgent:
         self.dangerously_skip_permissions = dangerously_skip_permissions or prompt_mode
         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:
@@ -109,6 +111,7 @@ class TinybirdAgent:
                     takes_ctx=True,
                 ),
                 Tool(plan, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
+                Tool(complete_plan, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(build, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(deploy, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(deploy_check, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
@@ -297,6 +300,18 @@ class TinybirdAgent:
     def add_message(self, message: ModelMessage) -> None:
         self.messages.append(message)
 
+    def start_plan(self, plan) -> str:
+        self.confirmed_plan_id = hashlib.sha256(plan.encode()).hexdigest()[:16]
+        return self.confirmed_plan_id
+
+    def cancel_plan(self) -> Optional[str]:
+        plan_id = self.confirmed_plan_id
+        self.confirmed_plan_id = None
+        return plan_id
+
+    def get_plan(self) -> Optional[str]:
+        return self.confirmed_plan_id
+
     def _build_agent_deps(self, config: dict[str, Any], run_id: Optional[str] = None) -> TinybirdAgentContext:
         project = self.project
         folder = self.project.folder
@@ -336,6 +351,9 @@ class TinybirdAgent:
             local_host=local_client.host,
             local_token=local_client.token,
             run_id=run_id,
+            get_plan=self.get_plan,
+            start_plan=self.start_plan,
+            cancel_plan=self.cancel_plan,
         )
 
     def run(self, user_prompt: str, config: dict[str, Any]) -> None:
@@ -389,8 +407,8 @@ class TinybirdAgent:
             ai_credits_limits = limits_data.get("limits", {}).get("ai_credits", {})
             current_ai_credits = ai_credits_limits.get("quantity") or 0
             ai_credits = ai_credits_limits.get("max") or 0
-            remaining_credits = max(ai_credits - current_ai_credits, 0)
-            current_ai_credits = min(ai_credits, current_ai_credits)
+            remaining_credits = round(max(ai_credits - current_ai_credits, 0), 2)
+            current_ai_credits = round(min(ai_credits, current_ai_credits), 2)
             if not ai_credits:
                 return
             warning_threshold = ai_credits * 0.8
@@ -569,6 +587,7 @@ def run_agent(
                         ]
                     )
                 )
+                agent.cancel_plan()
                 continue
             except KeyboardInterrupt:
                 click.echo(FeedbackManager.info(message="Goodbye!"))

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/agent/prompts.py

@@ -876,6 +876,7 @@ You have access to the following tools:
 15. `run_command` - Run a command using the Tinybird CLI.
 16. `diff_resource` - Diff the content of a resource in Tinybird Cloud vs Tinybird Local vs Project local file.
 17. `rename_datafile_or_fixture` - Rename a datafile or fixture.
+18. `complete_plan` - Complete a plan.
 
 # When creating, updating, or deleting files:
 1. Use `plan` tool to plan the creation, update, rename, or deletion of resources.
@@ -885,6 +886,7 @@ You have access to the following tools:
 5. If the file was created or removed successfully, report the result to the user.
 6. If the file was not created or removed, finish the process and just wait for a new user prompt.
 7. If the file was created or removed successfully, but the build failed, try to fix the error and repeat the process.
+8. If the plan is completed or cancelled, use the `complete_plan` tool to complete the plan.
 
 # When creating a landing datasource given a .ndjson file:
 - If the user does not specify anything about the desired schema, create a schema like this (sorting key not needed in this case)

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/agent/tools/append.py

@@ -53,9 +53,10 @@ def append_file(
                     return "Append operation cancelled by user."
 
         cloud_or_local = "Cloud" if cloud else "Local"
+        active_plan = ctx.deps.get_plan() is not None and not cloud
         confirmation = show_confirmation(
             title=f"Append fixture {fixture_pathname} to datasource '{datasource_name}' in Tinybird {cloud_or_local}?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
 
         if confirmation == "review":
@@ -112,9 +113,10 @@ def append_url(
                     return "Append operation cancelled by user."
 
         cloud_or_local = "Cloud" if cloud else "Local"
+        active_plan = ctx.deps.get_plan() is not None and not cloud
         confirmation = show_confirmation(
             title=f"Append URL {fixture_url} to datasource '{datasource_name}' in Tinybird {cloud_or_local}?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
 
         if confirmation == "review":

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/agent/tools/datafile.py

@@ -50,9 +50,10 @@ def create_datafile(
             content = create_terminal_box(resource.content, title=resource.pathname)
         click.echo(content)
         action = "Create" if not exists else "Update"
+        active_plan = ctx.deps.get_plan() is not None
         confirmation = show_confirmation(
             title=f"{action} '{resource.pathname}'?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
 
         if confirmation == "review":
@@ -146,9 +147,10 @@ def rename_datafile_or_fixture(ctx: RunContext[TinybirdAgentContext], path: str,
     """
     try:
         ctx.deps.thinking_animation.stop()
+        active_plan = ctx.deps.get_plan() is not None
         confirmation = show_confirmation(
             title=f"Rename '{path}' to '{new_path}'?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
 
         if confirmation == "review":
@@ -211,10 +213,10 @@ def remove_file(ctx: RunContext[TinybirdAgentContext], path: str) -> str:
             click.echo(FeedbackManager.error(message=f"Error: File {path} not found"))
             ctx.deps.thinking_animation.start()
             return f"Error: File {path} not found (double check the file path)"
-
+        active_plan = ctx.deps.get_plan() is not None
         confirmation = show_confirmation(
             title=f"Delete '{path}'?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
 
         if confirmation == "review":

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/agent/tools/mock.py

@@ -70,9 +70,10 @@ def generate_mock_fixture(
         content = create_terminal_box(preview_content, title=f"fixtures/{datasource_name}.{data_format}")
         click.echo(content)
         click.echo("Showing a preview of the file.\n")
+        active_plan = ctx.deps.get_plan() is not None
         confirmation = show_confirmation(
             title=f"Create fixture file for datasource '{datasource_name}'?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
 
         if confirmation == "review":
@@ -85,7 +86,7 @@ def generate_mock_fixture(
         click.echo(FeedbackManager.success(message=f"✓ {fixture_path_name} created"))
         confirmation = show_confirmation(
             title=f"Append {fixture_path_name} to datasource '{datasource_name}'?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
         if confirmation == "review":
             feedback = show_input(ctx.deps.workspace_name)

tinybird-0.0.1.dev288/tinybird/tb/modules/agent/tools/plan.py

@@ -0,0 +1,86 @@
+from typing import Literal
+
+import click
+from pydantic_ai import RunContext
+
+from tinybird.tb.modules.agent.utils import (
+    AgentRunCancelled,
+    TinybirdAgentContext,
+    show_input,
+    show_options,
+)
+
+PlanConfirmationResult = Literal["yes", "review", "yes_and_auto_implement"]
+
+
+def show_plan_confirmation(skip_confirmation: bool = False) -> PlanConfirmationResult:
+    if skip_confirmation:
+        return "yes"
+
+    title = "Do you want to continue with the plan?"
+    while True:
+        result = show_options(
+            options=[
+                "Yes, continue",
+                "Yes, continue and implement all",
+                "No, tell Tinybird Code what to do",
+                "Cancel",
+            ],
+            title=title,
+        )
+
+        if result is None:  # Cancelled
+            raise AgentRunCancelled(f"User cancelled the operation: {title}")
+
+        if result.startswith("Yes, continue and implement all"):
+            return "yes_and_auto_implement"
+        if result.startswith("Yes"):
+            return "yes"
+        elif result.startswith("No"):
+            return "review"
+
+        raise AgentRunCancelled(f"User cancelled the operation: {title}")
+
+
+def plan(ctx: RunContext[TinybirdAgentContext], plan: str) -> str:
+    """Given a plan, ask the user for confirmation to implement it
+
+    Args:
+        plan (str): The plan to implement. Required.
+
+    Returns:
+        str: If the plan was implemented or not.
+    """
+    ctx.deps.thinking_animation.stop()
+    plan = plan.strip()
+
+    click.echo(plan)
+    confirmation = show_plan_confirmation(skip_confirmation=ctx.deps.dangerously_skip_permissions)
+
+    if confirmation == "review":
+        feedback = show_input(ctx.deps.workspace_name)
+        ctx.deps.thinking_animation.start()
+        ctx.deps.cancel_plan()
+        return f"User did not confirm the proposed plan and gave the following feedback: {feedback}"
+
+    ctx.deps.thinking_animation.start()
+
+    if confirmation == "yes_and_auto_implement":
+        plan_id = ctx.deps.start_plan(plan=plan)
+        return f"User confirmed the plan {plan_id}. Implementing..."
+    else:
+        return "User confirmed the plan. Implementing..."
+
+
+def complete_plan(ctx: RunContext[TinybirdAgentContext]) -> str:
+    """Given an ongoing plan, complete it
+
+    Args:
+        ctx (RunContext[TinybirdAgentContext]): The context of the agent.
+
+    Returns:
+        str: The result of the plan.
+    """
+
+    plan_id = ctx.deps.cancel_plan()
+    return f"Plan {plan_id} completed"

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/agent/tools/secret.py

@@ -70,9 +70,10 @@ def create_or_update_secrets(ctx: RunContext[TinybirdAgentContext], secrets: dic
         preview_content = create_terminal_box(new_content, title=".env.local")
         click.echo(preview_content)
 
+        active_plan = ctx.deps.get_plan() is not None
         confirmation = show_confirmation(
             title=f"{action} {len(secrets)} secret(s) in .env.local?",
-            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+            skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
         )
 
         if confirmation == "review":

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/agent/utils.py

@@ -67,6 +67,9 @@ class TinybirdAgentContext(BaseModel):
     local_host: str
     local_token: str
     run_id: Optional[str] = None
+    get_plan: Callable[..., Optional[str]]
+    start_plan: Callable[..., str]
+    cancel_plan: Callable[..., Optional[str]]
 
 
 default_style = PromptStyle.from_dict(
@@ -778,10 +781,10 @@ def copy_fixture_to_project_folder_if_needed(
         if input_path.exists() and not _is_path_inside_project(input_path, project_folder):
             # Ask for confirmation to copy the file
             click.echo(FeedbackManager.highlight(message=f"» File {fixture_pathname} is outside the project folder."))
-
+            active_plan = ctx.deps.get_plan() is not None
             confirmation = show_confirmation(
                 title=f"Copy {input_path.name} to project folder for analysis?",
-                skip_confirmation=ctx.deps.dangerously_skip_permissions,
+                skip_confirmation=ctx.deps.dangerously_skip_permissions or active_plan,
             )
 
             if confirmation == "review":

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird/tb/modules/deployment_common.py

@@ -134,7 +134,11 @@ def promote_deployment(host: Optional[str], headers: dict, wait: bool, ingest_hi
             result = api_fetch(TINYBIRD_API_URL, headers=headers)
 
             last_deployment = result.get("deployment")
-            if last_deployment.get("status") == "deleted":
+            if not last_deployment:
+                click.echo(FeedbackManager.error(message="Error parsing deployment from response"))
+                sys_exit("deployment_error", "Error parsing deployment from response")
+
+            if last_deployment and last_deployment.get("status") == "deleted":
                 click.echo(FeedbackManager.success(message=f"✓ Deployment #{candidate_deployment.get('id')} is live!"))
                 break
 

{tinybird-0.0.1.dev287 → tinybird-0.0.1.dev288}/tinybird.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tinybird
-Version: 0.0.1.dev287
+Version: 0.0.1.dev288
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/forward/commands
 Author: Tinybird

tinybird-0.0.1.dev287/tinybird/tb/modules/agent/tools/plan.py

@@ -1,29 +0,0 @@
-import click
-from pydantic_ai import RunContext
-
-from tinybird.tb.modules.agent.utils import TinybirdAgentContext, show_confirmation, show_input
-
-
-def plan(ctx: RunContext[TinybirdAgentContext], plan: str) -> str:
-    """Given a plan, ask the user for confirmation to implement it
-
-    Args:
-        plan (str): The plan to implement. Required.
-
-    Returns:
-        str: If the plan was implemented or not.
-    """
-    ctx.deps.thinking_animation.stop()
-    plan = plan.strip()
-    click.echo(plan)
-    confirmation = show_confirmation(
-        title="Do you want to continue with the plan?", skip_confirmation=ctx.deps.dangerously_skip_permissions
-    )
-
-    if confirmation == "review":
-        feedback = show_input(ctx.deps.workspace_name)
-        ctx.deps.thinking_animation.start()
-        return f"User did not confirm the proposed plan and gave the following feedback: {feedback}"
-
-    ctx.deps.thinking_animation.start()
-    return "User confirmed the plan. Implementing..."