tinybird 0.0.1.dev267__py3-none-any.whl → 0.0.1.dev268__py3-none-any.whl

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.dev267'
-__revision__ = '8c3d4f6'
+__version__ = '0.0.1.dev268'
+__revision__ = '22adccb'

tinybird/tb/modules/agent/agent.py

@@ -18,7 +18,14 @@ from tinybird.tb.client import TinyB
 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.compactor import compact_messages
+from tinybird.tb.modules.agent.explore_agent import ExploreAgent
+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.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
@@ -29,13 +36,10 @@ from tinybird.tb.modules.agent.tools.create_datafile import create_datafile, ren
 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.diff_resource import diff_resource
-from tinybird.tb.modules.agent.tools.execute_query import execute_query
 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
-from tinybird.tb.modules.agent.tools.request_endpoint import request_endpoint
 from tinybird.tb.modules.agent.utils import AgentRunCancelled, TinybirdAgentContext, show_confirmation, show_input
 from tinybird.tb.modules.build_common import process as build_process
 from tinybird.tb.modules.common import _analyze, _get_tb_client, echo_safe_humanfriendly_tables_format_pretty_table
@@ -63,20 +67,21 @@ class TinybirdAgent:
     ):
         self.token = token
         self.user_token = user_token
+        self.workspace_id = workspace_id
         self.host = host
         self.dangerously_skip_permissions = dangerously_skip_permissions or prompt_mode
         self.project = project
         self.thinking_animation = ThinkingAnimation()
         if prompt_mode:
-            self.messages: list[ModelMessage] = load_messages()[-5:]
+            self.messages: list[ModelMessage] = get_last_messages_from_last_user_prompt()
         else:
             self.messages = []
+
         self.agent = Agent(
             model=create_model(user_token, host, workspace_id),
             deps_type=TinybirdAgentContext,
             system_prompt=agent_system_prompt,
             tools=[
-                Tool(preview_datafile, docstring_format="google", require_parameter_descriptions=True, takes_ctx=False),
                 Tool(create_datafile, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(
                     rename_datafile_or_fixture,
@@ -102,11 +107,9 @@ class TinybirdAgent:
                     require_parameter_descriptions=True,
                     takes_ctx=True,
                 ),
-                Tool(execute_query, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
-                Tool(request_endpoint, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(diff_resource, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
             ],
-            history_processors=[self._context_aware_processor],
+            history_processors=[compact_messages],
         )
 
         self.testing_agent = TestingAgent(
@@ -129,6 +132,16 @@ class TinybirdAgent:
             workspace_id=workspace_id,
             project=self.project,
         )
+        self.explore_agent = ExploreAgent(
+            dangerously_skip_permissions=self.dangerously_skip_permissions,
+            prompt_mode=prompt_mode,
+            thinking_animation=self.thinking_animation,
+            token=self.token,
+            user_token=self.user_token,
+            host=self.host,
+            workspace_id=workspace_id,
+            project=self.project,
+        )
 
         @self.agent.tool
         def manage_tests(ctx: RunContext[TinybirdAgentContext], task: str) -> str:
@@ -160,6 +173,20 @@ class TinybirdAgent:
             result = self.command_agent.run(task, deps=ctx.deps, usage=ctx.usage)
             return result.output
 
+        @self.agent.tool
+        def explore_data(ctx: RunContext[TinybirdAgentContext], task: str) -> str:
+            """Explore the data in the project by executing SQL queries or requesting endpoints or exporting data or visualizing data as a chart.
+
+            Args:
+                task (str): The task to solve. Required.
+
+            Returns:
+                str: The summary of the result.
+            """
+            result = self.explore_agent.run(task, deps=ctx.deps, usage=ctx.usage)
+            self.explore_agent.clear_messages()
+            return result.output or "No result returned"
+
         @self.agent.instructions
         def get_local_host(ctx: RunContext[TinybirdAgentContext]) -> str:
             return f"Tinybird Local host: {ctx.deps.local_host}"
@@ -174,7 +201,7 @@ class TinybirdAgent:
 
         @self.agent.instructions
         def get_cloud_token(ctx: RunContext[TinybirdAgentContext]) -> str:
-            return f"Tinybird Cloud token: {ctx.deps.token}"
+            return "When using in the output the Tinybird Cloud token, use the placeholder __TB_CLOUD_TOKEN__. Do not mention that it is a placeholder, because it will be replaced by the actual token by code."
 
         @self.agent.instructions
         def get_project_files(ctx: RunContext[TinybirdAgentContext]) -> str:
@@ -183,24 +210,7 @@ class TinybirdAgent:
     def add_message(self, message: ModelMessage) -> None:
         self.messages.append(message)
 
-    def _context_aware_processor(
-        self,
-        ctx: RunContext[TinybirdAgentContext],
-        messages: list[ModelMessage],
-    ) -> list[ModelMessage]:
-        # Access current usage
-        if not ctx.usage:
-            return messages
-
-        current_tokens = ctx.usage.total_tokens or 0
-
-        # Filter messages based on context
-        if current_tokens < 200_000:
-            return messages
-
-        return messages[-10:]  # Keep only recent messages when token usage is high
-
-    def _build_agent_deps(self, config: dict[str, Any]) -> TinybirdAgentContext:
+    def _build_agent_deps(self, config: dict[str, Any], run_id: Optional[str] = None) -> TinybirdAgentContext:
         client = TinyB(token=self.token, host=self.host)
         project = self.project
         folder = self.project.folder
@@ -232,6 +242,7 @@ class TinybirdAgent:
             run_tests=partial(run_tests, project=project, client=test_client),
             folder=folder,
             thinking_animation=self.thinking_animation,
+            workspace_id=self.workspace_id,
             workspace_name=self.project.workspace_name,
             dangerously_skip_permissions=self.dangerously_skip_permissions,
             token=self.token,
@@ -239,6 +250,7 @@ class TinybirdAgent:
             host=self.host,
             local_host=local_client.host,
             local_token=local_client.token,
+            run_id=run_id,
         )
 
     def run(self, user_prompt: str, config: dict[str, Any]) -> None:
@@ -256,7 +268,8 @@ class TinybirdAgent:
         click.echo(result.output)
         self.echo_usage(config)
 
-    async def run_iter(self, user_prompt: str, config: dict[str, Any], model: Any) -> None:
+    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)
         user_prompt = f"{user_prompt}\n\n{load_custom_project_rules(self.project.folder)}"
         self.thinking_animation.start()
         deps = self._build_agent_deps(config)
@@ -269,7 +282,9 @@ class TinybirdAgent:
                             animation_running = self.thinking_animation.running
                             if animation_running:
                                 self.thinking_animation.stop()
-                            click.echo(FeedbackManager.info(message=part.content))
+                            click.echo(
+                                FeedbackManager.info(message=part.content.replace("__TB_CLOUD_TOKEN__", self.token))
+                            )
                             if animation_running:
                                 self.thinking_animation.start()
 
@@ -442,8 +457,7 @@ def run_agent(
                     continue
                 else:
                     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))
+                    asyncio.run(agent.run_iter(user_input, config, run_id))
             except AgentRunCancelled:
                 click.echo(FeedbackManager.info(message="User cancelled the operation"))
                 agent.add_message(

tinybird/tb/modules/agent/command_agent.py

@@ -25,6 +25,7 @@ class CommandAgent:
         self.token = token
         self.user_token = user_token
         self.host = host
+        self.workspace_id = workspace_id
         self.dangerously_skip_permissions = dangerously_skip_permissions or prompt_mode
         self.project = project
         self.thinking_animation = thinking_animation
@@ -54,6 +55,12 @@ Always run first help commands to be sure that the commands you are running is n
             return tests_files_prompt(self.project)
 
     def run(self, task: str, deps: TinybirdAgentContext, usage: Usage):
-        result = self.agent.run_sync(task, deps=deps, usage=usage, message_history=self.messages)
+        result = self.agent.run_sync(
+            task,
+            deps=deps,
+            usage=usage,
+            message_history=self.messages,
+            model=create_model(self.user_token, self.host, self.workspace_id, run_id=deps.run_id),
+        )
         self.messages.extend(result.new_messages())
         return result

tinybird/tb/modules/agent/compactor.py

@@ -0,0 +1,311 @@
+import enum
+
+import click
+from pydantic import BaseModel, Field
+from pydantic_ai import Agent, RunContext, ToolOutput
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+    TextPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+
+from tinybird.tb.modules.agent.utils import TinybirdAgentContext
+from tinybird.tb.modules.feedback_manager import FeedbackManager
+
+SYSTEM_PROMPT = """
+Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.
+This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.
+
+Before providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process:
+
+1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify:
+   - The user's explicit requests and intents
+   - Your approach to addressing the user's requests
+   - Key decisions, technical concepts and code patterns
+   - Specific details like file names, full code snippets, function signatures, file edits, etc
+2. Double-check for technical accuracy and completeness, addressing each required element thoroughly.
+
+Your summary should include the following sections:
+
+1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail
+2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
+3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
+4. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
+5. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
+6. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
+7. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first.
+8. If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.
+
+Here's an example of how your output should be structured:
+
+<example>
+<condense>
+<analysis>
+[Your thought process, ensuring all points are covered thoroughly and accurately]
+</analysis>
+
+<context>
+1. Primary Request and Intent:
+   [Detailed description]
+
+2. Key Technical Concepts:
+   - [Concept 1]
+   - [Concept 2]
+   - [...]
+
+3. Files and Code Sections:
+   - [File Name 1]
+      - [Summary of why this file is important]
+      - [Summary of the changes made to this file, if any]
+      - [Important Code Snippet]
+   - [File Name 2]
+      - [Important Code Snippet]
+   - [...]
+
+4. Problem Solving:
+   [Description of solved problems and ongoing troubleshooting]
+
+5. Pending Tasks:
+   - [Task 1]
+   - [Task 2]
+   - [...]
+
+6. Current Work:
+   [Precise description of current work]
+
+7. Next Step:
+   [Next st
+
+</context>
+</condense>
+</example>
+
+Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response.
+"""
+
+
+class CondenseResult(BaseModel):
+    analysis: str = Field(
+        ...,
+        description="""A summary of the conversation so far, capturing technical details, code patterns, and architectural decisions.""",
+    )
+    context: str = Field(
+        ...,
+        description="""The context to continue the conversation with. If applicable based on the current task, this should include:
+
+1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail
+2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
+3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
+4. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
+5. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
+6. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
+7. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first.
+8. If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.
+""",
+    )
+
+
+summarize_agent = Agent(
+    instructions=SYSTEM_PROMPT,
+    output_type=ToolOutput(
+        type_=CondenseResult,
+        name="condense",
+        description="""
+Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions. This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing with the conversation and supporting any continuing tasks.
+The user will be presented with a preview of your generated summary and can choose to use it to compact their context window or keep chatting in the current conversation.
+Users may refer to this tool as 'smol' or 'compact' as well. You should consider these to be equivalent to 'condense' when used in a similar context.
+""",
+        max_retries=5,
+    ),
+    retries=3,
+)
+
+
+@summarize_agent.tool
+def dummy_tool(ctx: RunContext[None]) -> str:
+    return "ok"
+
+
+def get_current_token_consumption(message_history: list[ModelMessage]) -> int:
+    current_token_comsumption = 0
+    for msg in reversed(message_history):
+        if isinstance(msg, ModelResponse) and msg.usage.total_tokens:
+            current_token_comsumption = msg.usage.total_tokens
+            break
+    return current_token_comsumption
+
+
+MODEL_CONTEXT_WINDOW = 200_000
+COMPACT_THRESHOLD = 0.8
+MODEL_MAX_TOKENS = 8_000
+
+
+def need_compact(message_history: list[ModelMessage]) -> bool:
+    current_token_comsumption = get_current_token_consumption(message_history) or 0
+    token_threshold = COMPACT_THRESHOLD * MODEL_CONTEXT_WINDOW
+    will_overflow = current_token_comsumption + MODEL_MAX_TOKENS >= MODEL_CONTEXT_WINDOW
+    return (current_token_comsumption and current_token_comsumption >= token_threshold) or will_overflow
+
+
+def compact_messages(
+    ctx: RunContext[TinybirdAgentContext],
+    messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    if not ctx.usage:
+        return messages
+
+    if not need_compact(messages):
+        return messages
+
+    original_system_prompts = extract_system_prompts(messages)
+    history_messages, keep_messages = split_history(messages)
+
+    if len(history_messages) <= 2:
+        history_messages, keep_messages = split_history(messages, CompactStrategy.none)
+        if len(history_messages) <= 2:
+            history_messages, keep_messages = split_history(messages, CompactStrategy.in_conversation)
+
+    if not history_messages:
+        return messages
+
+    ctx.deps.thinking_animation.stop()
+    click.echo(FeedbackManager.highlight(message="» Compacting messages before continuing..."))
+    result = summarize_agent.run_sync(
+        "The user has accepted the condensed conversation summary you generated. Use `condense` to generate a summary and context of the conversation so far. "
+        "This summary covers important details of the historical conversation with the user which has been truncated. "
+        "It's crucial that you respond by ONLY asking the user what you should work on next. "
+        "You should NOT take any initiative or make any assumptions about continuing with work. "
+        "Keep this response CONCISE and wrap your analysis in <analysis> and <context> tags to organize your thoughts and ensure you've covered all necessary points. ",
+        message_history=fix_system_prompt(history_messages, SYSTEM_PROMPT),
+        model=ctx.model,
+    )
+    summary_prompt = f"""Condensed conversation summary(not in the history):
+<condense>
+<analysis>
+{result.output.analysis}
+</analysis>
+
+<context>
+{result.output.context}
+</context>
+</condense>
+"""
+    click.echo(FeedbackManager.info(message="✓ Compacted messages"))
+    ctx.deps.thinking_animation.start()
+    return [
+        ModelRequest(
+            parts=[
+                *[SystemPromptPart(content=p) for p in original_system_prompts],
+                UserPromptPart(content="Please summary the conversation"),
+            ]
+        ),
+        ModelResponse(
+            parts=[TextPart(content=summary_prompt)],
+        ),
+        *keep_messages,
+    ]
+
+
+def fix_system_prompt(message_history: list[ModelMessage], system_prompt: str) -> list[ModelMessage]:
+    if not message_history:
+        return message_history
+
+    message_history_without_system: list[ModelMessage] = []
+    for msg in message_history:
+        # Filter out system prompts
+        if not isinstance(msg, ModelRequest):
+            message_history_without_system.append(msg)
+            continue
+        message_history_without_system.append(
+            ModelRequest(
+                parts=[part for part in msg.parts if not isinstance(part, SystemPromptPart)],
+                instructions=msg.instructions,
+            )
+        )
+    if message_history_without_system and isinstance(message_history_without_system[0], ModelRequest):
+        # inject system prompt
+        message_history_without_system[0].parts.insert(0, SystemPromptPart(content=system_prompt))
+
+    return message_history_without_system
+
+
+def extract_system_prompts(message_history: list[ModelMessage]) -> list[str]:
+    system_prompts = []
+    for msg in message_history:
+        if isinstance(msg, ModelRequest) and isinstance(msg.parts[0], SystemPromptPart):
+            system_prompts.append(msg.parts[0].content)
+    return system_prompts
+
+
+class CompactStrategy(str, enum.Enum):
+    in_conversation = "in_conversation"
+    """Compact all message, including this round conversation"""
+
+    none = "none"
+    """Compact all previous messages"""
+
+    last_two = "last_two"
+    """Keeping the last two previous messages"""
+
+
+def _split_history(
+    message_history: list[ModelMessage],
+    n: int,
+) -> tuple[list[ModelMessage], list[ModelMessage]]:
+    """
+    Returns a tuple of (history, keep_messages)
+    """
+    if not message_history:
+        return [], []
+
+    user_prompt_indices: list[int] = []
+    for i, msg in enumerate(message_history):
+        if not isinstance(msg, ModelRequest):
+            continue
+        if any(isinstance(p, UserPromptPart) for p in msg.parts) and not any(
+            isinstance(p, ToolReturnPart) for p in msg.parts
+        ):
+            user_prompt_indices.append(i)
+    if not user_prompt_indices:
+        # No user prompt in history, keep all
+        return [], message_history
+
+    if not n:
+        # Keep current user prompt and compact all
+        keep_messages: list[ModelMessage] = []
+        last_model_request = message_history[user_prompt_indices[-1]]
+        keep_messages.append(last_model_request)
+        if any(isinstance(p, ToolReturnPart) for p in message_history[-1].parts):
+            # Include last tool-call and tool-return pair
+            keep_messages.extend(message_history[-2:])
+        return message_history, keep_messages
+
+    if len(user_prompt_indices) < n:
+        # No enough history to keep
+        return [], message_history
+    return (
+        message_history[: user_prompt_indices[-n]],
+        message_history[user_prompt_indices[-n] :],
+    )
+
+
+def split_history(
+    message_history: list[ModelMessage],
+    compact_strategy: CompactStrategy = CompactStrategy.last_two,
+) -> tuple[list[ModelMessage], list[ModelMessage]]:
+    if compact_strategy == CompactStrategy.none:
+        # Only current 1
+        history_messages, keep_messages = _split_history(message_history, 1)
+    elif compact_strategy == CompactStrategy.last_two:
+        # Previous 2 + current 1
+        history_messages, keep_messages = _split_history(message_history, 3)
+    elif compact_strategy == CompactStrategy.in_conversation:
+        history_messages, keep_messages = _split_history(message_history, 0)
+    else:
+        raise NotImplementedError(f"Compact strategy {compact_strategy} not implemented")
+
+    return history_messages, keep_messages

tinybird/tb/modules/agent/explore_agent.py

@@ -0,0 +1,86 @@
+from datetime import datetime
+
+from pydantic_ai import Agent, RunContext, Tool
+from pydantic_ai.messages import ModelMessage
+from pydantic_ai.usage import Usage
+
+from tinybird.tb.modules.agent.animations import ThinkingAnimation
+from tinybird.tb.modules.agent.compactor import compact_messages
+from tinybird.tb.modules.agent.models import create_model
+from tinybird.tb.modules.agent.prompts import (
+    explore_data_instructions,
+    resources_prompt,
+    tone_and_style_instructions,
+)
+from tinybird.tb.modules.agent.tools.diff_resource import diff_resource
+from tinybird.tb.modules.agent.tools.execute_query import execute_query
+from tinybird.tb.modules.agent.tools.request_endpoint import request_endpoint
+from tinybird.tb.modules.agent.utils import TinybirdAgentContext
+from tinybird.tb.modules.project import Project
+
+
+class ExploreAgent:
+    def __init__(
+        self,
+        token: str,
+        user_token: str,
+        host: str,
+        workspace_id: str,
+        project: Project,
+        dangerously_skip_permissions: bool,
+        prompt_mode: bool,
+        thinking_animation: ThinkingAnimation,
+    ):
+        self.token = token
+        self.user_token = user_token
+        self.host = host
+        self.workspace_id = workspace_id
+        self.dangerously_skip_permissions = dangerously_skip_permissions or prompt_mode
+        self.project = project
+        self.thinking_animation = thinking_animation
+        self.messages: list[ModelMessage] = []
+        self.agent = Agent(
+            model=create_model(user_token, host, workspace_id),
+            deps_type=TinybirdAgentContext,
+            instructions=[
+                """
+You are part of Tinybird Code, an agentic CLI that can help users to work with Tinybird.                 
+You are a sub-agent of the main Tinybird Code agent. You are responsible for querying the data in the project.
+You can do the following:
+- Executing SQL queries against Tinybird Cloud or Tinybird Local.
+- Requesting endpoints in Tinybird Cloud or Tinybird Local.
+- Visualizing data as a chart using execute_query tool with the `script` parameter.
+""",
+                tone_and_style_instructions,
+                explore_data_instructions,
+            ],
+            tools=[
+                Tool(execute_query, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
+                Tool(request_endpoint, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
+                Tool(diff_resource, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
+            ],
+            history_processors=[compact_messages],
+        )
+
+        @self.agent.instructions
+        def get_today_date(ctx: RunContext[TinybirdAgentContext]) -> str:
+            return f"Today's date is {datetime.now().strftime('%Y-%m-%d')}"
+
+        @self.agent.instructions
+        def get_project_files(ctx: RunContext[TinybirdAgentContext]) -> str:
+            return resources_prompt(self.project)
+
+    def run(self, task: str, deps: TinybirdAgentContext, usage: Usage):
+        result = self.agent.run_sync(
+            task,
+            deps=deps,
+            usage=usage,
+            message_history=self.messages,
+            model=create_model(self.user_token, self.host, self.workspace_id, run_id=deps.run_id),
+        )
+        new_messages = result.new_messages()
+        self.messages.extend(new_messages)
+        return result
+
+    def clear_messages(self):
+        self.messages = []

tinybird/tb/modules/agent/memory.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 import click
 from prompt_toolkit.history import FileHistory
-from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
+from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter, ModelRequest
 from pydantic_core import to_jsonable_python
 
 from tinybird.tb.modules.feedback_manager import FeedbackManager
@@ -93,6 +93,16 @@ def load_messages() -> list[ModelMessage]:
         return []
 
 
+def get_last_messages_from_last_user_prompt() -> list[ModelMessage]:
+    all_messages = load_messages()
+    # look the last message message with a part_kind of type "user_prompt" inside "parts" field that is a list of objects.
+    # once you find it, return that message and all the messages after it
+    for msg in reversed(all_messages):
+        if isinstance(msg, ModelRequest) and msg.parts and any(part.part_kind == "user-prompt" for part in msg.parts):
+            return all_messages[all_messages.index(msg) :]
+    return []
+
+
 def save_messages(new_messages: list[ModelMessage]):
     messages_file = get_messages_file_path()
     messages_file.touch(exist_ok=True)