tinybird 0.0.1.dev265__tar.gz → 0.0.1.dev266__tar.gz

{tinybird-0.0.1.dev265 → tinybird-0.0.1.dev266}/PKG-INFO

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

{tinybird-0.0.1.dev265 → tinybird-0.0.1.dev266}/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.dev265'
-__revision__ = 'b723917'
+__version__ = '0.0.1.dev266'
+__revision__ = 'b4193e9'

{tinybird-0.0.1.dev265 → tinybird-0.0.1.dev266}/tinybird/tb/client.py

@@ -870,9 +870,15 @@ class TinyB:
             data={"operation": "change_role", "users": users, "new_role": role},
         )
 
-    def workspace(self, workspace_id: str, with_token: bool = False):
-        params = {"with_token": "true" if with_token else "false"}
-        return self._req(f"/v0/workspaces/{workspace_id}?{urlencode(params)}")
+    def workspace(
+        self, workspace_id: str, with_token: bool = False, with_organization: bool = False, version: str = "v0"
+    ):
+        params = {}
+        if with_token:
+            params["with_token"] = "true"
+        if with_organization:
+            params["with_organization"] = "true"
+        return self._req(f"/{version}/workspaces/{workspace_id}?{urlencode(params)}")
 
     def workspace_info(self, version: str = "v0") -> Dict[str, Any]:
         return self._req(f"/{version}/workspace")
@@ -880,6 +886,12 @@ class TinyB:
     def organization(self, organization_id: str):
         return self._req(f"/v0/organizations/{organization_id}")
 
+    def organization_limits(self, organization_id: str):
+        return self._req(f"/v1/billing/{organization_id}/limits")
+
+    def organization_subscription(self, organization_id: str):
+        return self._req(f"/v1/billing/{organization_id}/subscription")
+
     def create_organization(
         self,
         name: str,

{tinybird-0.0.1.dev265 → tinybird-0.0.1.dev266}/tinybird/tb/modules/agent/agent.py

@@ -3,6 +3,7 @@ import shlex
 import subprocess
 import sys
 import urllib.parse
+import uuid
 from functools import partial
 from pathlib import Path
 from typing import Any, Optional
@@ -19,7 +20,7 @@ from tinybird.tb.modules.agent.animations import ThinkingAnimation
 from tinybird.tb.modules.agent.banner import display_banner
 from tinybird.tb.modules.agent.command_agent import CommandAgent
 from tinybird.tb.modules.agent.memory import clear_history, clear_messages, load_messages, save_messages
-from tinybird.tb.modules.agent.models import create_model, model_costs
+from tinybird.tb.modules.agent.models import create_model
 from tinybird.tb.modules.agent.prompts import agent_system_prompt, load_custom_project_rules, resources_prompt
 from tinybird.tb.modules.agent.testing_agent import TestingAgent
 from tinybird.tb.modules.agent.tools.analyze import analyze_file, analyze_url
@@ -256,12 +257,12 @@ class TinybirdAgent:
         click.echo(result.output)
         self._echo_usage(config, result)
 
-    async def run_iter(self, user_prompt: str, config: dict[str, Any]) -> None:
+    async def run_iter(self, user_prompt: str, config: dict[str, Any], model: Any) -> None:
         user_prompt = f"{user_prompt}\n\n{load_custom_project_rules(self.project.folder)}"
         self.thinking_animation.start()
         deps = self._build_agent_deps(config)
 
-        async with self.agent.iter(user_prompt, deps=deps, message_history=self.messages) as agent_run:
+        async with self.agent.iter(user_prompt, deps=deps, message_history=self.messages, model=model) as agent_run:
             async for node in agent_run:
                 if hasattr(node, "model_response"):
                     for _i, part in enumerate(node.model_response.parts):
@@ -281,19 +282,32 @@ class TinybirdAgent:
             self._echo_usage(config, agent_run.result)
 
     def _echo_usage(self, config: dict[str, Any], result: AgentRunResult) -> None:
-        if "@tinybird.co" in config.get("user_email", ""):
-            usage = result.usage()
-            request_tokens = usage.request_tokens or 0
-            response_tokens = usage.response_tokens or 0
-            total_tokens = usage.total_tokens or 0
-            cost = (
-                request_tokens * model_costs["input_cost_per_token"]
-                + response_tokens * model_costs["output_cost_per_token"]
+        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")
+            limits_data = client.organization_limits(workspace["organization"]["id"])
+            ai_requests_limits = limits_data.get("limits", {}).get("ai_requests", {})
+            current_ai_requests = ai_requests_limits.get("quantity", 0)
+            max_ai_requests = ai_requests_limits.get("max", 0)
+            if not max_ai_requests:
+                return
+            remaining_requests = max(max_ai_requests - current_ai_requests, 0)
+            warning_threshold = max_ai_requests * 0.8
+            if current_ai_requests >= warning_threshold:
+                message_color = FeedbackManager.warning
+            else:
+                message_color = FeedbackManager.gray
+
+            current_ai_requests = min(max_ai_requests, current_ai_requests)
+
+            click.echo(
+                message_color(
+                    message=f"{remaining_requests} agent requests left ({current_ai_requests}/{max_ai_requests}). This message is informative. Limits will be enforced soon."
+                )
             )
-            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}")
+        except Exception:
+            pass
 
 
 def run_agent(
@@ -426,7 +440,9 @@ def run_agent(
                 elif user_input.strip() == "":
                     continue
                 else:
-                    asyncio.run(agent.run_iter(user_input, config))
+                    run_id = str(uuid.uuid4())
+                    model = create_model(user_token, host, workspace_id, run_id=run_id)
+                    asyncio.run(agent.run_iter(user_input, config, model))
             except AgentRunCancelled:
                 click.echo(FeedbackManager.info(message="User cancelled the operation"))
                 agent.add_message(

{tinybird-0.0.1.dev265 → tinybird-0.0.1.dev266}/tinybird/tb/modules/agent/models.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from anthropic import AsyncAnthropic
 from httpx import AsyncClient
 from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelName
@@ -9,11 +11,17 @@ def create_model(
     base_url: str,
     workspace_id: str,
     model: AnthropicModelName = "claude-4-sonnet-20250514",
+    run_id: Optional[str] = None,
 ):
+    default_headers = {}
+    if run_id:
+        default_headers["X-Run-Id"] = run_id
+
     client = AsyncAnthropic(
         base_url=base_url,
         http_client=AsyncClient(params={"token": token, "workspace_id": workspace_id}),
         auth_token=token,
+        default_headers=default_headers,
     )
     return AnthropicModel(
         model_name=model,

{tinybird-0.0.1.dev265 → tinybird-0.0.1.dev266}/tinybird/tb/modules/token.py

@@ -163,7 +163,7 @@ def create(ctx: Context) -> None:
 
     * Static Tokens do not have a TTL and can have any valid scope (ADMIN, TOKENS, or ORG_DATASOURCES:READ).
 
-    * JWT tokens have a TTL and can only have the PIPES:READ scope.Their main use case is allow your users to call your endpoints without exposing your API key.
+    * JWT tokens have a TTL and can have PIPES:READ and DATASOURCES:READ scopes. Their main use case is allow your users to call your endpoints or read datasources without exposing your API key.
 
 
     Examples:
@@ -173,6 +173,8 @@ def create(ctx: Context) -> None:
     tb token create static my_static_token --scope TOKENS
 
     tb token create jwt my_jwt_token --ttl 1h --scope PIPES:READ --resource my_pipe
+
+    tb token create jwt my_jwt_token --ttl 1h --scope DATASOURCES:READ --resource my_datasource --filter "column = 'value'"
     """
 
 
@@ -182,16 +184,21 @@ def create(ctx: Context) -> None:
 @click.option(
     "--scope",
     multiple=True,
-    type=click.Choice(["PIPES:READ"]),
+    type=click.Choice(["PIPES:READ", "DATASOURCES:READ"]),
     required=True,
-    help="Scope of the token (only PIPES:READ is allowed for JWT tokens)",
+    help="Scope of the token (only PIPES:READ and DATASOURCES:READ are allowed for JWT tokens)",
 )
 @click.option("--resource", multiple=True, required=True, help="Resource associated with the scope")
 @click.option(
     "--fixed-params", multiple=True, help="Fixed parameters in key=value format, multiple values separated by commas"
 )
+@click.option(
+    "--filter",
+    multiple=True,
+    help="SQL filter to apply when reading a datasource (only applicable for DATASOURCES:READ scope)",
+)
 @click.pass_context
-def create_jwt_token(ctx: Context, name: str, ttl: timedelta, scope, resource, fixed_params) -> None:
+def create_jwt_token(ctx: Context, name: str, ttl: timedelta, scope, resource, fixed_params, filter) -> None:
     """Create a JWT token with a TTL specify."""
 
     obj: Dict[str, Any] = ctx.ensure_object(dict)
@@ -205,6 +212,10 @@ def create_jwt_token(ctx: Context, name: str, ttl: timedelta, scope, resource, f
     if fixed_params and len(fixed_params) > len(scope):
         raise CLITokenException(FeedbackManager.error_number_of_fixed_params_and_resources_mismatch())
 
+    # Ensure the number of filters does not exceed the number of scope/resource pairs
+    if filter and len(filter) > len(scope):
+        raise CLITokenException("The number of SQL filters must match the number of scopes")
+
     # Parse fixed params
     parsed_fixed_params = parse_fixed_params(fixed_params) if fixed_params else []
 
@@ -213,15 +224,24 @@ def create_jwt_token(ctx: Context, name: str, ttl: timedelta, scope, resource, f
     for i, params in enumerate(parsed_fixed_params):
         fixed_params_list[i] = params
 
+    # Create a list of filters for each scope/resource pair, defaulting to None if not provided
+    filters_list: List[Optional[str]] = [None] * len(scope)
+    for i, f in enumerate(filter):
+        filters_list[i] = f
+
     scopes = []
-    for sc, res, fparams in zip(scope, resource, fixed_params_list):
-        scopes.append(
-            {
-                "type": sc,
-                "resource": res,
-                "fixed_params": fparams,
-            }
-        )
+    for sc, res, fparams, filter in zip(scope, resource, fixed_params_list, filters_list):
+        scope_data = {
+            "type": sc,
+            "resource": res,
+        }
+
+        if sc == "PIPES:READ" and fparams:
+            scope_data["fixed_params"] = fparams
+        elif sc == "DATASOURCES:READ" and filter:
+            scope_data["filter"] = filter
+
+        scopes.append(scope_data)
 
     try:
         response = client.create_jwt_token(name, expiration_time, scopes)
@@ -258,6 +278,7 @@ class DynamicOptionsCommand(click.Command):
         # Options
         dynamic_options_help += "Options:\n"
         dynamic_options_help += f"  --scope [{','.join(valid_scopes)}]   Scope for the token [Required]\n"
+        dynamic_options_help += "  --resource TEXT    Resource you want to associate the scope with\n"
         dynamic_options_help += "  -h, --help            Show this message and exit.\n"
 
         return dynamic_options_help

{tinybird-0.0.1.dev265 → tinybird-0.0.1.dev266}/tinybird.egg-info/PKG-INFO

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