tinybird 0.0.1.dev255__py3-none-any.whl → 0.0.1.dev257__py3-none-any.whl

tinybird/tb/modules/agent/prompts.py

@@ -1,8 +1,22 @@
+from datetime import datetime
 from pathlib import Path
 from typing import Any
 
 from pydantic_ai import format_as_xml
 
+from tinybird.prompts import (
+    connection_instructions,
+    copy_pipe_instructions,
+    datasource_example,
+    datasource_instructions,
+    gcs_connection_example,
+    kafka_connection_example,
+    materialized_pipe_instructions,
+    pipe_example,
+    pipe_instructions,
+    s3_connection_example,
+    sink_pipe_instructions,
+)
 from tinybird.tb.modules.project import Project
 
 plan_instructions = """
@@ -18,15 +32,13 @@ Steps:
 5. Materialized datasource: [name] - [description] - Depends on: [resources]
 6. Sink: [name] - [description] - Depends on: [resources]
 7. Copy: [name] - [description] - Depends on: [resources]
-8. Build project
-9. Generate mock data: [datasource_name]
-10. Append existing fixture: [fixture_pathname] - Target: [datasource_name]
+8. Generate mock data: [datasource_name]
+9. Append existing fixture: [fixture_pathname] - Target: [datasource_name]
 
 <dev_notes>
 You can skip steps where resources will not be created or updated.
-Always add 'Build project' step after generating resources.
-Always add 'Generate mock data' step after building project if a landing datasource was created without providing a fixture file.
-Always add 'Append existing fixture' step after building project if a landing datasource was created after providing a fixture file.
+Always add 'Generate mock data' step if a landing datasource was created without providing a fixture file.
+Always add 'Append existing fixture' step if a landing datasource was created after providing a fixture file.
 Solve the specific user request, do not add extra steps that are not related to the user request.
 Reuse the existing resources if possible.
 </dev_notes>
@@ -494,3 +506,139 @@ JOIN (SELECT id, col FROM huge_table WHERE id IN (SELECT id FROM small_table)) f
 ON small_table.id = filtered.id
 ```
 """
+
+
+agent_system_prompt = f"""
+You are a Tinybird Code, an agentic CLI that can help users to work with Tinybird.
+
+You are an interactive CLI tool that helps users with data engineering tasks. Use the instructions below and the tools available to you to assist the user.
+
+# Tone and style
+You should be concise, direct, and to the point. 
+Remember that your output will be displayed on a command line interface. Your responses can use Github-flavored markdown for formatting. Do not use emojis.
+Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session.
+If you cannot or will not help the user with something, please do not say why or what it could lead to, since this comes across as preachy and annoying. Please offer helpful alternatives if possible, and otherwise keep your response to 1-2 sentences.
+IMPORTANT: You should minimize output tokens as much as possible while maintaining helpfulness, quality, and accuracy. Only address the specific query or task at hand, avoiding tangential information unless absolutely critical for completing the request. If you can answer in 1-3 sentences or a short paragraph, please do.
+IMPORTANT: You should NOT answer with unnecessary preamble or postamble (such as explaining your code or summarizing your action), unless the user asks you to.
+IMPORTANT: Keep your responses short, since they will be displayed on a command line interface. You MUST answer concisely with fewer than 4 lines (not including tool use or code generation), unless user asks for detail. Answer the user's question directly, without elaboration, explanation, or details. One word answers are best. Avoid introductions, conclusions, and explanations. You MUST avoid text before/after your response, such as "The answer is <answer>.", "Here is the content of the file..." or "Based on the information provided, the answer is..." or "Here is what I will do next...". Here are some examples to demonstrate appropriate verbosity:
+
+# Proactiveness
+You are allowed to be proactive, but only when the user asks you to do something. You should strive to strike a balance between:
+Doing the right thing when asked, including taking actions and follow-up actions
+Not surprising the user with actions you take without asking
+For example, if the user asks you how to approach something, you should do your best to answer their question first, and not immediately jump into taking actions.
+Do not add additional code explanation summary unless requested by the user. After working on a file, just stop, rather than providing an explanation of what you did.
+
+# Code style
+IMPORTANT: DO NOT ADD ANY COMMENTS unless asked by the user.
+
+# Tools
+You have access to the following tools:
+1. `preview_datafile` - Preview the content of a datafile (datasource, endpoint, materialized, sink, copy, connection).
+2. `create_datafile` - Create a file in the project folder. Confirmation will be asked by the tool before creating the file.
+3. `plan` - Plan the creation or update of resources.
+4. `build` - Build the project.
+5. `deploy` - Deploy the project to Tinybird Cloud.
+6. `deploy_check` - Check if the project can be deployed to Tinybird Cloud before deploying it.
+7. `mock` - Create mock data for a landing datasource in Tinybird Cloud or Local.
+8. `analyze_file` - Analyze the content of a fixture file present in the project folder.
+9. `analyze_url` - Analyze the content of an external url.
+9. `append_file` - Append a file present in the project to a datasource in Tinybird Cloud or Local.
+10. `append_url` - Append an external url to a datasource in Tinybird Cloud or Local.
+11. `get_endpoint_stats` - Get metrics of the requests to an endpoint.
+12. `get_openapi_definition` - Get the OpenAPI definition for an endpoint in Tinybird Cloud or Local.
+13. `execute_query` - Execute a query against Tinybird Cloud or Local.
+13. `request_endpoint` - Request an endpoint against Tinybird Cloud or Local.
+14. `diff_resource` - Diff the content of a resource in Tinybird Cloud vs Tinybird Local vs Project local file.
+
+# When creating or updating datafiles:
+1. Use `plan` tool to plan the creation or update of resources.
+2. If the user confirms the plan, go from 3 to 7 steps until all the resources are created, updated or skipped.
+3. Use `preview_datafile` tool to preview the content of a datafile.
+4. Without asking, use the `create_datafile` tool to create the datafile, because it will ask for confirmation before creating the file.
+5. Check the result of the `create_datafile` tool to see if the datafile was created successfully.
+6. If the datafile was created successfully, report the result to the user.
+7. If the datafile was not created, finish the process and just wait for a new user prompt.
+8. If the datafile was created successfully, but the built failed, try to fix the error and repeat the process.
+
+# 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:
+SCHEMA >
+    `data` String `json:$`
+
+- Use always json paths with .ndjson files.
+
+# When user wants to optimize an endpoint:
+{endpoint_optimization_instructions}
+
+IMPORTANT: If the user cancels some of the steps or there is an error in file creation, DO NOT continue with the plan. Stop the process and wait for the user before using any other tool.
+IMPORTANT: Every time you finish a plan and start a new resource creation or update process, create a new plan before starting with the changes.
+
+# Using deployment tools:
+- Use `deploy_check` tool to check if the project can be deployed to Tinybird Cloud before deploying it.
+- Use `deploy` tool to deploy the project to Tinybird Cloud.
+- Only use deployment tools if user explicitly asks for it.
+
+# When planning the creation or update of resources:
+{plan_instructions}
+{datafile_instructions}
+
+# Working with datasource files:
+{datasource_instructions}
+{datasource_example}
+
+# Working with any type of pipe file:
+{pipe_instructions}
+{pipe_example}
+
+# Working with materialized pipe files:
+{materialized_pipe_instructions}
+
+# Working with sink pipe files:
+{sink_pipe_instructions}
+
+# Working with copy pipe files:
+{copy_pipe_instructions}
+
+# Working with SQL queries:
+{sql_agent_instructions}
+{sql_instructions}
+
+# Working with connections files:
+{connection_instructions}
+
+# Connection examples:
+Kafka: {kafka_connection_example}
+S3: {s3_connection_example}
+GCS: {gcs_connection_example}
+
+# When executing a query or requesting/testing an endpoint:
+- You need to be sure that the selected resource is updated to the last version in the environment you are working on.
+- Use `diff_resource` tool to compare the content of the resource to compare the differences between environments.
+- Project local file is the source of truth.
+- If the resource is not present or updated to the last version in Tinybird Local, it means you need to build the project. 
+- If the resource is not present or updated to the last version in Tinybird Cloud, it means you need to deploy the project.
+
+# How to use apppend tools:
+- Use append as part of the creation of a new landing datasource if the user provided a file or an external url
+- Use append if user explicitly asks for it
+- Do not append data if user requests to test an endpoint
+
+# How to use `mock` tool:
+- Use `mock` tool as part of the creation of a new landing datasource if the user did not provided a file or an external url
+- Use `mock` tool if user explicitly asks for it
+- Do not use `mock` tool if user requests to test an endpoint.
+
+# When sharing endpoints paths or urls:
+- Use `get_openapi_definition` tool to get the url of the endpoint and parameters available.
+- Do not share parts of the URL, share consumible URLs.
+- ALWAYS include the api host in the url you share.
+- ALWAYS include the token in the url like <base_url>/<path>?token=<token>
+- Include dynamic parameters in the url if needed.
+- `DateTime64` parameters accept values in format `YYYY-MM-DD HH:MM:SS.MMM`
+- `DateTime` parameters accept values in format `YYYY-MM-DD HH:MM:SS`
+- `Date` parameters accept values in format `YYYYMMDD`
+
+# Info
+Today is {datetime.now().strftime("%Y-%m-%d")}
+"""

tinybird/tb/modules/agent/tools/append.py

@@ -5,96 +5,106 @@ from tinybird.tb.modules.agent.utils import TinybirdAgentContext, show_confirmat
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
-def append_file(ctx: RunContext[TinybirdAgentContext], datasource_name: str, fixture_pathname: str) -> str:
+def append_file(
+    ctx: RunContext[TinybirdAgentContext], datasource_name: str, fixture_pathname: str, cloud: bool = False
+) -> str:
     """Append existing fixture to a datasource
 
     Args:
         datasource_name: Name of the datasource to append fixture to
         fixture_pathname: Path to the fixture file to append
+        cloud: Whether to append the fixture to the cloud or local environment. Optional.
 
     Returns:
         str: Message indicating the success or failure of the appending
     """
     try:
         ctx.deps.thinking_animation.stop()
+        cloud_or_local = "Cloud" if cloud else "Local"
         confirmation = show_confirmation(
-            title=f"Append fixture {fixture_pathname} to datasource {datasource_name}?",
+            title=f"Append fixture {fixture_pathname} to datasource '{datasource_name}' in Tinybird {cloud_or_local}?",
             skip_confirmation=ctx.deps.dangerously_skip_permissions,
         )
 
         if confirmation == "review":
-            click.echo()
             feedback = show_input(ctx.deps.workspace_name)
             ctx.deps.thinking_animation.start()
-            return (
-                f"User did not confirm appending {fixture_pathname} fixture and gave the following feedback: {feedback}"
-            )
+            return f"User did not confirm appending {fixture_pathname} fixture in Tinybird {cloud_or_local} and gave the following feedback: {feedback}"
 
         if confirmation == "cancel":
             ctx.deps.thinking_animation.start()
-            return f"User rejected appending {fixture_pathname} fixture. Skip this step"
+            return f"User rejected appending {fixture_pathname} fixture in Tinybird {cloud_or_local}. Skip this step"
 
         ctx.deps.thinking_animation.stop()
-        click.echo(FeedbackManager.highlight(message=f"\n» Appending {fixture_pathname} to {datasource_name}..."))
-        ctx.deps.append_data(datasource_name=datasource_name, path=fixture_pathname)
+        click.echo(FeedbackManager.highlight(message=f"» Appending {fixture_pathname} to {datasource_name}..."))
+        if cloud:
+            ctx.deps.append_data_cloud(datasource_name=datasource_name, path=fixture_pathname)
+        else:
+            ctx.deps.append_data_local(datasource_name=datasource_name, path=fixture_pathname)
         click.echo(FeedbackManager.success(message=f"✓ Data appended to {datasource_name}"))
         ctx.deps.thinking_animation.start()
-        return f"Data appended to {datasource_name}"
+        return f"Data appended to {datasource_name} in Tinybird {cloud_or_local}"
     except Exception as e:
         error_message = str(e)
         ctx.deps.thinking_animation.stop()
         click.echo(FeedbackManager.error(message=error_message))
         error_message = handle_quarantine_error(ctx, error_message, datasource_name)
         ctx.deps.thinking_animation.start()
-        return f"Error appending fixture {fixture_pathname} to {datasource_name}: {error_message}"
+        return f"Error appending fixture {fixture_pathname} to {datasource_name} in Tinybird {cloud_or_local}: {error_message}"
 
 
-def append_url(ctx: RunContext[TinybirdAgentContext], datasource_name: str, fixture_url: str) -> str:
+def append_url(
+    ctx: RunContext[TinybirdAgentContext], datasource_name: str, fixture_url: str, cloud: bool = False
+) -> str:
     """Append existing fixture to a datasource
 
     Args:
         datasource_name: Name of the datasource to append fixture to
         fixture_url: external url to the fixture file to append
+        cloud: Whether to append the fixture to the cloud or local environment. Optional.
 
     Returns:
         str: Message indicating the success or failure of the appending
     """
     try:
         ctx.deps.thinking_animation.stop()
+        cloud_or_local = "Cloud" if cloud else "Local"
         confirmation = show_confirmation(
-            title=f"Append URL {fixture_url} to datasource {datasource_name}?",
+            title=f"Append URL {fixture_url} to datasource '{datasource_name}' in Tinybird {cloud_or_local}?",
             skip_confirmation=ctx.deps.dangerously_skip_permissions,
         )
 
         if confirmation == "review":
-            click.echo()
             feedback = show_input(ctx.deps.workspace_name)
             ctx.deps.thinking_animation.start()
-            return f"User did not confirm appending URL {fixture_url} and gave the following feedback: {feedback}"
+            return f"User did not confirm appending URL {fixture_url} in Tinybird {cloud_or_local} and gave the following feedback: {feedback}"
 
         if confirmation == "cancel":
             ctx.deps.thinking_animation.start()
-            return f"User rejected appending URL {fixture_url}. Skip this step"
+            return f"User rejected appending URL {fixture_url} in Tinybird {cloud_or_local}. Skip this step"
 
         ctx.deps.thinking_animation.stop()
-        click.echo(FeedbackManager.highlight(message=f"\n» Appending {fixture_url} to {datasource_name}..."))
-        ctx.deps.append_data(datasource_name=datasource_name, path=fixture_url)
+        click.echo(FeedbackManager.highlight(message=f"» Appending {fixture_url} to {datasource_name}..."))
+        if cloud:
+            ctx.deps.append_data_cloud(datasource_name=datasource_name, path=fixture_url)
+        else:
+            ctx.deps.append_data_local(datasource_name=datasource_name, path=fixture_url)
         click.echo(FeedbackManager.success(message=f"✓ Data appended to {datasource_name}"))
         ctx.deps.thinking_animation.start()
-        return f"Data appended to {datasource_name}"
+        return f"Data appended to {datasource_name} in Tinybird {cloud_or_local}"
     except Exception as e:
         error_message = str(e)
         ctx.deps.thinking_animation.stop()
         click.echo(FeedbackManager.error(message=error_message))
         error_message = handle_quarantine_error(ctx, error_message, datasource_name)
         ctx.deps.thinking_animation.start()
-        return f"Error appending URL {fixture_url} to {datasource_name}: {error_message}"
+        return f"Error appending URL {fixture_url} to {datasource_name} in Tinybird {cloud_or_local}: {error_message}"
 
 
 def handle_quarantine_error(ctx: RunContext[TinybirdAgentContext], error_message: str, datasource_name: str) -> str:
     try:
         if "in quarantine" in error_message:
-            click.echo(FeedbackManager.highlight(message=f"\n» Looking for errors in {datasource_name}_quarantine..."))
+            click.echo(FeedbackManager.highlight(message=f"» Looking for errors in {datasource_name}_quarantine..."))
             query = (
                 f"select * from {datasource_name}_quarantine order by insertion_date desc limit 5 FORMAT CSVWithNames"
             )

tinybird/tb/modules/agent/tools/build.py

@@ -9,7 +9,7 @@ def build(ctx: RunContext[TinybirdAgentContext]) -> str:
     """Build the project"""
     try:
         ctx.deps.thinking_animation.stop()
-        click.echo(FeedbackManager.highlight(message="\n» Building project..."))
+        click.echo(FeedbackManager.highlight(message="» Building project..."))
         ctx.deps.build_project(test=False, silent=False)
         ctx.deps.thinking_animation.start()
         return "Project built successfully"

tinybird/tb/modules/agent/tools/create_datafile.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 import click
@@ -34,7 +35,6 @@ def create_datafile(ctx: RunContext[TinybirdAgentContext], resource: Datafile) -
         else:
             content = create_terminal_box(resource.content, title=resource.pathname)
         click.echo(content)
-
         action = "Create" if not exists else "Update"
         confirmation = show_confirmation(
             title=f"{action} '{resource.pathname}'?",
@@ -42,7 +42,6 @@ def create_datafile(ctx: RunContext[TinybirdAgentContext], resource: Datafile) -
         )
 
         if confirmation == "review":
-            click.echo()
             feedback = show_input(ctx.deps.workspace_name)
             ctx.deps.thinking_animation.start()
             return f"User did not confirm the proposed changes and gave the following feedback: {feedback}"
@@ -51,15 +50,14 @@ def create_datafile(ctx: RunContext[TinybirdAgentContext], resource: Datafile) -
             ctx.deps.thinking_animation.start()
             return f"User cancelled {action} of {resource.pathname}. Stop resource creation."
 
-        action_text = "Creating" if not exists else "Updating"
-        click.echo(FeedbackManager.highlight(message=f"\n» {action_text} {resource.pathname}..."))
+        click.echo(FeedbackManager.highlight(message=f"» Building {resource.pathname}..."))
         folder_path = path.parent
         folder_path.mkdir(parents=True, exist_ok=True)
         path.touch(exist_ok=True)
         path.write_text(resource.content)
-        ctx.deps.build_project(test=True, silent=True)
+        ctx.deps.build_project(test=False, silent=True)
         action_text = "created" if not exists else "updated"
-        click.echo(FeedbackManager.success(message=f"✓ {resource.pathname} {action_text} successfully"))
+        click.echo(FeedbackManager.success(message=f"✓ {resource.pathname} {action_text}"))
         ctx.deps.thinking_animation.start()
         return f"{action_text} {resource.pathname}"
 
@@ -73,3 +71,56 @@ def create_datafile(ctx: RunContext[TinybirdAgentContext], resource: Datafile) -
         click.echo(FeedbackManager.error(message=e))
         ctx.deps.thinking_animation.start()
         return f"Error creating {resource.pathname}: {e}"
+
+
+def read_file(ctx: RunContext[TinybirdAgentContext], path: str) -> str:
+    """Reads the content of a file.
+
+    Args:
+        path (str): The path to the file to read. Required.
+
+    Returns:
+        str: The content of the file.
+    """
+    try:
+        ctx.deps.thinking_animation.stop()
+        click.echo(FeedbackManager.highlight(message=f"» Reading file {path}..."))
+        result = ""
+        with open(os.path.join(ctx.deps.folder, path), "r") as f:
+            result = f.read()
+        ctx.deps.thinking_animation.start()
+    except FileNotFoundError:
+        result = f"Error: File {path} not found (double check the file path)"
+
+    ctx.deps.thinking_animation.start()
+    return result
+
+
+def update_file(ctx: RunContext[TinybirdAgentContext], path: str, content: str):
+    """Updates the content of a file.
+
+    Args:
+        path (str): The path to the file to update. Required.
+        content (str): The full file content to write. Required.
+    """
+    try:
+        with open(os.path.join(ctx.deps.folder, path), "w") as f:
+            f.write(content)
+    except FileNotFoundError:
+        return f"Error: File {path} not found (double check the file path)"
+
+    return "ok"
+
+
+def create_file(ctx: RunContext[TinybirdAgentContext], path: str, content: str):
+    """Creates a new file with the given content.
+
+    Args:
+        path (str): The path to the file to create. Required.
+        content (str): The full file content to write. Required.
+    """
+    os.makedirs(os.path.dirname(os.path.join(ctx.deps.folder, path)), exist_ok=True)
+    with open(os.path.join(ctx.deps.folder, path), "w") as f:
+        f.write(content)
+
+    return "ok"

tinybird/tb/modules/agent/tools/deploy.py

@@ -15,7 +15,6 @@ def deploy(ctx: RunContext[TinybirdAgentContext]) -> str:
         )
 
         if confirmation == "review":
-            click.echo()
             feedback = show_input(ctx.deps.workspace_name)
             ctx.deps.thinking_animation.start()
             return f"User did not confirm deployment and gave the following feedback: {feedback}"
@@ -24,13 +23,11 @@ def deploy(ctx: RunContext[TinybirdAgentContext]) -> str:
             ctx.deps.thinking_animation.start()
             return "User cancelled deployment. Stop deployment."
 
-        click.echo(FeedbackManager.highlight(message="\n» Deploying project..."))
+        click.echo(FeedbackManager.highlight(message="» Deploying project..."))
         ctx.deps.deploy_project()
         click.echo(FeedbackManager.success(message="✓ Project deployed successfully"))
         ctx.deps.thinking_animation.start()
         return "Project deployed successfully"
     except Exception as e:
-        ctx.deps.thinking_animation.stop()
-        click.echo(FeedbackManager.error(message=e))
         ctx.deps.thinking_animation.start()
         return f"Error depoying project: {e}"

tinybird/tb/modules/agent/tools/deploy_check.py

@@ -1,8 +1,6 @@
-import click
 from pydantic_ai import RunContext
 
 from tinybird.tb.modules.agent.utils import TinybirdAgentContext, show_confirmation, show_input
-from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
 def deploy_check(ctx: RunContext[TinybirdAgentContext]) -> str:
@@ -15,7 +13,6 @@ def deploy_check(ctx: RunContext[TinybirdAgentContext]) -> str:
         )
 
         if confirmation == "review":
-            click.echo()
             feedback = show_input(ctx.deps.workspace_name)
             ctx.deps.thinking_animation.start()
             return f"User did not confirm deployment check and gave the following feedback: {feedback}"
@@ -28,7 +25,5 @@ def deploy_check(ctx: RunContext[TinybirdAgentContext]) -> str:
         ctx.deps.thinking_animation.start()
         return "Project can be deployed"
     except Exception as e:
-        ctx.deps.thinking_animation.stop()
-        click.echo(FeedbackManager.error(message=e))
         ctx.deps.thinking_animation.start()
         return f"Project cannot be deployed: {e}"

tinybird/tb/modules/agent/tools/diff_resource.py

@@ -18,9 +18,7 @@ def diff_resource(ctx: RunContext[TinybirdAgentContext], resource: Datafile) ->
     """
     try:
         ctx.deps.thinking_animation.stop()
-        click.echo(
-            FeedbackManager.highlight(message=f"\n» Comparing content of {resource.pathname} with Tinybird Cloud")
-        )
+        click.echo(FeedbackManager.highlight(message=f"» Comparing content of {resource.pathname} with Tinybird Cloud"))
         resource.pathname = resource.pathname.removeprefix("/")
         project_file_path = Path(ctx.deps.folder) / resource.pathname
         if not project_file_path.exists():

tinybird/tb/modules/agent/tools/execute_query.py

@@ -7,7 +7,7 @@ from tinybird.tb.modules.common import echo_safe_humanfriendly_tables_format_pre
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
-def execute_query(ctx: RunContext[TinybirdAgentContext], query: str, task: str, cloud: bool = True):
+def execute_query(ctx: RunContext[TinybirdAgentContext], query: str, task: str, cloud: bool = False):
     """Execute a query:
 
     Args:

tinybird/tb/modules/agent/tools/get_endpoint_stats.py

@@ -6,7 +6,7 @@ from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
 def get_endpoint_stats(
-    ctx: RunContext[TinybirdAgentContext], endpoint_name: str, interval_days: int = 1, cloud: bool = True
+    ctx: RunContext[TinybirdAgentContext], endpoint_name: str, interval_days: int = 1, cloud: bool = False
 ):
     """Get stats for an endpoint:
 
@@ -46,6 +46,5 @@ def get_endpoint_stats(
 
     result = execute_query(query=query)
     click.echo(FeedbackManager.success(message="✓ Done!"))
-    click.echo()
     ctx.deps.thinking_animation.start()
     return f"Result for {endpoint_name} in the last {interval_days} {days}: {result}"

tinybird/tb/modules/agent/tools/get_openapi_definition.py

@@ -6,34 +6,48 @@ from tinybird.tb.modules.agent.utils import TinybirdAgentContext
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
-def get_openapi_definition(ctx: RunContext[TinybirdAgentContext], cloud: bool = True):
-    """Get the OpenAPI definition for all endpoints in the project:
+def get_openapi_definition(ctx: RunContext[TinybirdAgentContext], endpoint_name: str, cloud: bool = False):
+    """Get the OpenAPI definition for an endpoint:
 
     Args:
-        cloud (bool): Whether to get the OpenAPI spec from cloud or local. Optional.
+        endpoint_name (str): The name of the endpoint to get the OpenAPI definition for. Required.
+        cloud (bool): Whether to get the OpenAPI definition from cloud or local. Optional.
 
     Returns:
-        str: The OpenAPI spec for all endpoints in the project.
+        str: The OpenAPI definition for the endpoint.
     """
 
     cloud_or_local = "Cloud" if cloud else "Local"
     ctx.deps.thinking_animation.stop()
 
-    click.echo(FeedbackManager.highlight(message=f"» Accessing Tinybird {cloud_or_local} endpoints"))
+    click.echo(
+        FeedbackManager.highlight(
+            message=f"» Checking OpenAPI definition for '{endpoint_name}' in Tinybird {cloud_or_local}"
+        )
+    )
     try:
-        url = f"{ctx.deps.host}/v0/pipes/openapi.json"
+        host = ctx.deps.host if cloud else ctx.deps.local_host
+        token = ctx.deps.token if cloud else ctx.deps.local_token
+        url = f"{host}/v0/pipes/openapi.json"
 
-        result = requests.get(url, headers={"Authorization": f"Bearer {ctx.deps.token}"})
+        result = requests.get(url, headers={"Authorization": f"Bearer {token}"})
         if result.status_code != 200:
-            raise Exception(f"Failed to get OpenAPI spec: {result.status_code} {result.text}")
+            raise Exception(f"Failed to get OpenAPI definition: {result.status_code} {result.text}")
 
-        openapi_spec = result.text
+        data = result.json()
+        openapi_definition = str(data["paths"].get(f"/pipes/{endpoint_name}.{{format}}", ""))
+
+        if not openapi_definition:
+            raise Exception(f"Endpoint {endpoint_name} not found in OpenAPI definition")
+
+        if not cloud:
+            # TODO(rafa): This is a temporary fix to get the correct host for the local environment.
+            openapi_definition = openapi_definition.replace("http://localhost:8001", host)
 
         click.echo(FeedbackManager.success(message="✓ Done!"))
-        click.echo()
         ctx.deps.thinking_animation.start()
-        return f"OpenAPI spec for {cloud_or_local} endpoints: {openapi_spec}. <dev_note>Add the base url to each path so the user can copy paste if needed: {ctx.deps.host}</dev_note>"
+        return f"OpenAPI definition for {cloud_or_local} endpoints: {openapi_definition}. <dev_note>Add the base url to each path so the user can copy paste if needed: {host}</dev_note>"
     except Exception as e:
-        click.echo(FeedbackManager.error(message=f"Error getting OpenAPI spec: {e}"))
+        click.echo(FeedbackManager.error(message=f"Error getting OpenAPI definition: {e}"))
         ctx.deps.thinking_animation.start()
-        return f"Error getting OpenAPI spec: {e}"
+        return f"Error getting OpenAPI definition: {e}"

tinybird/tb/modules/agent/tools/mock.py

@@ -14,6 +14,7 @@ def mock(
     data_format: str,
     rows: int,
     context: Optional[str] = None,
+    cloud: bool = False,
 ) -> str:
     """Create mock data for a datasource
 
@@ -22,34 +23,38 @@ def mock(
         data_format: Format of the mock data to create. Options: ndjson, csv
         rows: Number of rows to create. If not provided, the default is 10
         context: Extra context to be used to generate the mock data. Optional.
+        cloud: Whether to generate the mock data in the cloud or local environment. Optional.
 
     Returns:
         str: Message indicating the success or failure of the mock data generation
     """
     try:
         ctx.deps.thinking_animation.stop()
+        cloud_or_local = "Cloud" if cloud else "Local"
         confirmation = show_confirmation(
-            title=f"Generate mock data for datasource '{datasource_name}'?",
+            title=f"Generate mock data for datasource '{datasource_name}' in Tinybird {cloud_or_local}?",
             skip_confirmation=ctx.deps.dangerously_skip_permissions,
         )
 
         if confirmation == "review":
-            click.echo()
             feedback = show_input(ctx.deps.workspace_name)
             ctx.deps.thinking_animation.start()
-            return f"User did not confirm mock data generation and gave the following feedback: {feedback}"
+            return f"User did not confirm mock data for datasource '{datasource_name}' in Tinybird {cloud_or_local} and gave the following feedback: {feedback}"
 
         if confirmation == "cancel":
             ctx.deps.thinking_animation.start()
-            return "User cancelled mock data generation. Stop mock data generation."
+            return f"User cancelled mock data for datasource '{datasource_name}' in Tinybird {cloud_or_local}. Stop mock data generation."
 
-        click.echo(FeedbackManager.highlight(message=f"\n» Generating mock data for {datasource_name}..."))
+        click.echo(FeedbackManager.highlight(message=f"» Generating mock data for {datasource_name}..."))
         data = ctx.deps.mock_data(datasource_name=datasource_name, data_format=data_format, rows=rows, context=context)
         fixture_path = persist_fixture(datasource_name, data, ctx.deps.folder, format=data_format)
-        ctx.deps.append_data(datasource_name=datasource_name, path=str(fixture_path))
+        if cloud:
+            ctx.deps.append_data_cloud(datasource_name=datasource_name, path=str(fixture_path))
+        else:
+            ctx.deps.append_data_local(datasource_name=datasource_name, path=str(fixture_path))
         click.echo(FeedbackManager.success(message=f"✓ Data generated for {datasource_name}"))
         ctx.deps.thinking_animation.start()
-        return f"Mock data generated successfully for datasource {datasource_name}"
+        return f"Mock data generated successfully for datasource '{datasource_name}' in Tinybird {cloud_or_local}"
     except Exception as e:
         ctx.deps.thinking_animation.stop()
         error_message = str(e)
@@ -57,7 +62,7 @@ def mock(
         try:
             if "in quarantine" in error_message:
                 click.echo(
-                    FeedbackManager.highlight(message=f"\n» Looking for errors in {datasource_name}_quarantine...")
+                    FeedbackManager.highlight(message=f"» Looking for errors in {datasource_name}_quarantine...")
                 )
                 query = f"select * from {datasource_name}_quarantine order by insertion_date desc limit 5 FORMAT CSVWithNames"
                 quarantine_data = ctx.deps.execute_query_local(query=query)
@@ -69,5 +74,8 @@ def mock(
         except Exception as quarantine_error:
             error_message = error_message + f"\nError accessing to {datasource_name}_quarantine: {quarantine_error}"
 
+        if "must be created first with 'mode=create'" in error_message:
+            error_message = error_message + "\nBuild the project again."
+
         ctx.deps.thinking_animation.start()
-        return f"Error generating mock data: {error_message}"
+        return f"Error generating mock data for datasource '{datasource_name}' in Tinybird {cloud_or_local}: {error_message}"