tinybird 0.0.1.dev250__tar.gz → 0.0.1.dev251__tar.gz

{tinybird-0.0.1.dev250 → tinybird-0.0.1.dev251}/PKG-INFO

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

{tinybird-0.0.1.dev250 → tinybird-0.0.1.dev251}/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.dev250'
-__revision__ = '1c47b89'
+__version__ = '0.0.1.dev251'
+__revision__ = 'c11760c'

{tinybird-0.0.1.dev250 → tinybird-0.0.1.dev251}/tinybird/tb/modules/agent/agent.py

@@ -41,6 +41,7 @@ from tinybird.tb.modules.agent.tools.deploy import deploy
 from tinybird.tb.modules.agent.tools.deploy_check import deploy_check
 from tinybird.tb.modules.agent.tools.explore import explore_data
 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.mock import mock
 from tinybird.tb.modules.agent.tools.plan import plan
 from tinybird.tb.modules.agent.tools.preview_datafile import preview_datafile
@@ -114,6 +115,7 @@ You have access to the following tools:
 9. `read_fixture_data` - Read a fixture data file present in the project folder.
 10. `append` - Append existing fixture to a datasource.
 11. `get_endpoint_stats` - Get metrics of the requests to an endpoint.
+12. `get_openapi_definition` - Get the OpenAPI definition for all endpoints that are built/deployed to Tinybird Cloud or Local.
 
 # When creating or updating datafiles:
 1. Use `plan` tool to plan the creation or update of resources.
@@ -195,6 +197,12 @@ Today is {datetime.now().strftime("%Y-%m-%d")}
                 Tool(
                     get_endpoint_stats, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True
                 ),
+                Tool(
+                    get_openapi_definition,
+                    docstring_format="google",
+                    require_parameter_descriptions=True,
+                    takes_ctx=True,
+                ),
             ],
         )
 
@@ -227,6 +235,9 @@ Today is {datetime.now().strftime("%Y-%m-%d")}
                 thinking_animation=thinking_animation,
                 workspace_name=self.project.workspace_name,
                 dangerously_skip_permissions=self.dangerously_skip_permissions,
+                token=self.token,
+                user_token=self.user_token,
+                host=self.host,
             ),
             message_history=self.messages,
         )
@@ -243,11 +254,13 @@ Today is {datetime.now().strftime("%Y-%m-%d")}
         )
         click.echo(result.output)
         click.echo("\n")
-        click.echo(f"Input tokens: {request_tokens}")
-        click.echo(f"Output tokens: {response_tokens}")
-        click.echo(f"Total tokens: {total_tokens}")
-        click.echo(f"Cost: ${cost:.6f}")
-        click.echo("\n")
+
+        if "@tinybird.co" in config.get("user_email", ""):
+            click.echo(f"Input tokens: {request_tokens}")
+            click.echo(f"Output tokens: {response_tokens}")
+            click.echo(f"Total tokens: {total_tokens}")
+            click.echo(f"Cost: ${cost:.6f}")
+            click.echo("\n")
 
 
 def run_agent(config: dict[str, Any], project: Project, dangerously_skip_permissions: bool):
@@ -378,7 +391,12 @@ def append_data(config: dict[str, Any], datasource_name: str, path: str) -> None
 
 
 def mock_data(
-    config: dict[str, Any], project: Project, datasource_name: str, data_format: str, rows: int
+    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()
@@ -388,12 +406,11 @@ def mock_data(
         raise CLIMockException(f"Datasource {datasource_name} not found")
 
     datasource_content = Path(datasource_path).read_text()
-    prompt = ""
     return create_mock_data(
         datasource_name,
         datasource_content,
         rows,
-        prompt,
+        context or "",
         cli_config,
         config,
         cli_config.get_user_token() or "",

tinybird-0.0.1.dev251/tinybird/tb/modules/agent/tools/get_openapi_definition.py

@@ -0,0 +1,39 @@
+import click
+import requests
+from pydantic_ai import RunContext
+
+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:
+
+    Args:
+        cloud (bool): Whether to get the OpenAPI spec from cloud or local. Optional.
+
+    Returns:
+        str: The OpenAPI spec for all endpoints in the project.
+    """
+
+    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"))
+    try:
+        url = f"{ctx.deps.host}/v0/pipes/openapi.json"
+
+        result = requests.get(url, headers={"Authorization": f"Bearer {ctx.deps.token}"})
+        if result.status_code != 200:
+            raise Exception(f"Failed to get OpenAPI spec: {result.status_code} {result.text}")
+
+        openapi_spec = result.text
+
+        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>"
+    except Exception as e:
+        click.echo(FeedbackManager.error(message=f"Error getting OpenAPI spec: {e}"))
+        ctx.deps.thinking_animation.start()
+        return f"Error getting OpenAPI spec: {e}"

{tinybird-0.0.1.dev250 → tinybird-0.0.1.dev251}/tinybird/tb/modules/agent/tools/mock.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import click
 from pydantic_ai import RunContext
 
@@ -6,13 +8,20 @@ from tinybird.tb.modules.datafile.fixture import persist_fixture
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
-def mock(ctx: RunContext[TinybirdAgentContext], datasource_name: str, data_format: str, rows: int) -> str:
+def mock(
+    ctx: RunContext[TinybirdAgentContext],
+    datasource_name: str,
+    data_format: str,
+    rows: int,
+    context: Optional[str] = None,
+) -> str:
     """Create mock data for a datasource
 
     Args:
         datasource_name: Name of the datasource to create mock data for
         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.
 
     Returns:
         str: Message indicating the success or failure of the mock data generation
@@ -35,7 +44,7 @@ def mock(ctx: RunContext[TinybirdAgentContext], datasource_name: str, data_forma
             return "User cancelled mock data generation. Stop mock data generation."
 
         click.echo(FeedbackManager.highlight(message=f"\n» Generating mock data for {datasource_name}..."))
-        data = ctx.deps.mock_data(datasource_name=datasource_name, data_format=data_format, rows=rows)
+        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))
         click.echo(FeedbackManager.success(message=f"✓ Data generated for {datasource_name}"))
@@ -43,6 +52,22 @@ def mock(ctx: RunContext[TinybirdAgentContext], datasource_name: str, data_forma
         return f"Mock data generated successfully for datasource {datasource_name}"
     except Exception as e:
         ctx.deps.thinking_animation.stop()
-        click.echo(FeedbackManager.error(message=e))
+        error_message = str(e)
+        click.echo(FeedbackManager.error(message=error_message))
+        try:
+            if "in quarantine" in error_message:
+                click.echo(
+                    FeedbackManager.highlight(message=f"\n» 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_local_query(query=query)
+                error_message = (
+                    error_message
+                    + f"\nThese are the first 5 rows of the quarantine table for datasource '{datasource_name}':\n{quarantine_data}. Use again `mock` tool but add this issue to the context."
+                )
+
+        except Exception as quarantine_error:
+            error_message = error_message + f"\nError accessing to {datasource_name}_quarantine: {quarantine_error}"
+
         ctx.deps.thinking_animation.start()
-        return f"Error generating mock data: {e}"
+        return f"Error generating mock data: {error_message}"

{tinybird-0.0.1.dev250 → tinybird-0.0.1.dev251}/tinybird/tb/modules/agent/utils.py

@@ -45,6 +45,9 @@ class TinybirdAgentContext(BaseModel):
     execute_cloud_query: Callable[..., str]
     execute_local_query: Callable[..., str]
     dangerously_skip_permissions: bool
+    token: str
+    user_token: str
+    host: str
 
 
 default_style = PromptStyle.from_dict(

{tinybird-0.0.1.dev250 → tinybird-0.0.1.dev251}/tinybird.egg-info/PKG-INFO

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

{tinybird-0.0.1.dev250 → tinybird-0.0.1.dev251}/tinybird.egg-info/SOURCES.txt

@@ -94,6 +94,7 @@ tinybird/tb/modules/agent/tools/deploy.py
 tinybird/tb/modules/agent/tools/deploy_check.py
 tinybird/tb/modules/agent/tools/explore.py
 tinybird/tb/modules/agent/tools/get_endpoint_stats.py
+tinybird/tb/modules/agent/tools/get_openapi_definition.py
 tinybird/tb/modules/agent/tools/mock.py
 tinybird/tb/modules/agent/tools/plan.py
 tinybird/tb/modules/agent/tools/preview_datafile.py