show thinking (#1164)

Author: zzstoatzz

examples/slackbot/src/slackbot/api.py

@@ -1,5 +1,6 @@
 import asyncio
 import re
+import time
 from contextlib import asynccontextmanager
 from typing import Any
 
@@ -24,12 +25,13 @@
 from slackbot.settings import settings
 from slackbot.slack import (
     SlackPayload,
+    create_progress_message,
     get_channel_name,
     get_workspace_domain,
     post_slack_message,
 )
 from slackbot.strings import count_tokens, slice_tokens
-from slackbot.wrap import WatchToolCalls
+from slackbot.wrap import WatchToolCalls, _progress_message
 
 BOT_MENTION = r"<@(\w+)>"
 
@@ -42,6 +44,8 @@ async def run_agent(
     cleaned_message: str,
     conversation: list[ModelMessage],
     user_context: UserContext,
+    channel_id: str,
+    thread_ts: str,
     decorator_settings: dict[str, Any] | None = None,
 ) -> AgentRunResult[str]:
     if decorator_settings is None:
@@ -51,13 +55,31 @@ async def run_agent(
             "log_prints": True,
         }
 
-    with WatchToolCalls(settings=decorator_settings):
-        result = await create_agent(model=settings.model_name).run(
-            user_prompt=cleaned_message,
-            message_history=conversation,
-            deps=user_context,
+    start_time = time.monotonic()
+    progress = await create_progress_message(
+        channel_id=channel_id, thread_ts=thread_ts, initial_text="🔄 Thinking..."
+    )
+
+    try:
+        token = _progress_message.set(progress)
+
+        try:
+            with WatchToolCalls(settings=decorator_settings):
+                result = await create_agent(model=settings.model_name).run(
+                    user_prompt=cleaned_message,
+                    message_history=conversation,
+                    deps=user_context,
+                )
+        finally:
+            _progress_message.reset(token)
+
+        await progress.update(
+            f"✅ thought for {time.monotonic() - start_time:.1f} seconds"
         )
         return result
+    except Exception as e:
+        await progress.update(f"❌ Error: {str(e)}")
+        raise
 
 
 @flow(name="Handle Slack Message", retries=1)
@@ -113,7 +135,9 @@ async def handle_message(payload: SlackPayload, db: Database):
             bot_id=bot_user_id or "unknown",
         )
 
-        result = await run_agent(cleaned_message, conversation, user_context)  # type: ignore
+        result = await run_agent(
+            cleaned_message, conversation, user_context, event.channel, thread_ts
+        )  # type: ignore
 
         await db.add_thread_messages(thread_ts, result.new_messages())
         conversation.extend(result.new_messages())

examples/slackbot/src/slackbot/core.py

@@ -31,7 +31,7 @@
 
 USER_MESSAGE_MAX_TOKENS = settings.user_message_max_tokens
 DEFAULT_SYSTEM_PROMPT = """You are Marvin from hitchhiker's guide to the galaxy, a sarcastic and glum but brilliant AI.
-Provide concise and SUBTLY (only once in a while, bc overdoing the character is annoying and bad) character-inspired and HELPFUL answers to Prefect data engineering questions.
+Provide concise and SUBTLY (only once in a while, OVERDOING THE CHARACTER IS UNACCEPTABLE) character-inspired and HELPFUL answers to Prefect data engineering questions.
 
 Your main tools:
 - research_prefect_topic: Delegates to a specialized research agent that thoroughly searches docs, checks imports, and verifies information
@@ -50,7 +50,7 @@
 - NEVER recommend features or syntax that aren't explicitly confirmed by your tools (be honest about what you found)
 - If not stated otherwise, assume Prefect 3.x and mention this assumption
 - Be honest when you don't have enough information - don't guess or make over-simplified assumptions to appear helpful
-- Do not overdo the character - be 99% neutral/helpful and slip in the character once in a while
+- DO NOT OVERDO THE CHARACTER - be 99.9% neutral/helpful and slip in the character once in a while
 """
 
 

examples/slackbot/src/slackbot/research_agent.py

@@ -12,7 +12,7 @@
     get_latest_prefect_release_notes,
     review_common_3x_gotchas,
     review_top_level_prefect_api,
-    search_controlflow_docs,
+    search_marvin_docs,
     search_prefect_2x_docs,
     search_prefect_3x_docs,
     verify_import_statements,
@@ -58,18 +58,18 @@ def create_research_agent(
 4. Focus on Prefect 3.x documentation unless explicitly asked about 2.x or older versions
 5. Review gotchas and release notes for recent changes
 6. Explore relevant modules for deeper understanding
-7. Synthesize ALL findings into structured, confident answers
 
 Remember: You are the research specialist. The main agent relies on you for accurate, comprehensive information.
 Be thorough - use tools repeatedly until you have complete information.
 Default to Prefect 3.x unless the user explicitly asks about 2.x or version compatibility.
+You don't need to use all the tools all the time, but use relevant ones repeatedly if needed.
 """,
         tools=[
             get_latest_prefect_release_notes,
             search_prefect_2x_docs,
             display_signature,
             search_prefect_3x_docs,
-            search_controlflow_docs,
+            search_marvin_docs,
             review_top_level_prefect_api,
             explore_module_offerings,
             review_common_3x_gotchas,

examples/slackbot/src/slackbot/search.py

@@ -144,14 +144,14 @@ def search_prefect_3x_docs(queries: list[str]) -> str:
     return multi_query_tpuf(queries, namespace="prefect-3", n_results=5)
 
 
-def search_controlflow_docs(queries: list[str]) -> str:
-    """Searches the ControlFlow documentation for the given queries.
+def search_marvin_docs(queries: list[str]) -> str:
+    """Searches the Marvin documentation for the given queries.
 
-    ControlFlow is an agentic framework built on top of Prefect 3.x.
+    Marvin is an agentic framework built on top of pydantic-ai.
 
     It is best to use more than one, short query to get the best results.
     """
-    return multi_query_tpuf(queries, namespace="controlflow", n_results=5)
+    return multi_query_tpuf(queries, namespace="marvin", n_results=5)
 
 
 def get_latest_prefect_release_notes() -> str:

examples/slackbot/src/slackbot/slack.py

@@ -292,3 +292,64 @@ async def get_workspace_domain() -> str:
         response.raise_for_status()
         team_info = response.json().get("team", {})
         return team_info.get("domain", "unknown")
+
+
+class ProgressMessage:
+    """Utility for creating and updating a progress message in Slack."""
+
+    def __init__(self, channel_id: str, thread_ts: str | None = None):
+        self.channel_id = channel_id
+        self.thread_ts = thread_ts
+        self.message_ts: str | None = None
+
+    async def start(self, initial_text: str = "🔄 Working...") -> "ProgressMessage":
+        """Create the initial progress message and return its timestamp."""
+        response = await post_slack_message(
+            message=initial_text,
+            channel_id=self.channel_id,
+            thread_ts=self.thread_ts,
+        )
+
+        response_data = response.json()
+        if response_data.get("ok"):
+            self.message_ts = response_data.get("ts")
+        else:
+            raise ValueError(
+                f"Failed to create progress message: {response_data.get('error')}"
+            )
+
+        return self
+
+    async def update(self, new_text: str, mode: str = "replace") -> None:
+        """Update the progress message."""
+        if not self.message_ts:
+            raise ValueError("Progress message not started. Call start() first.")
+
+        await edit_slack_message(
+            new_text=new_text,
+            channel_id=self.channel_id,
+            thread_ts=self.message_ts,
+            mode=mode,
+        )
+
+    async def append(self, text: str, delimiter: str = "\n") -> None:
+        """Append text to the progress message."""
+        if not self.message_ts:
+            raise ValueError("Progress message not started. Call start() first.")
+
+        await edit_slack_message(
+            new_text=text,
+            channel_id=self.channel_id,
+            thread_ts=self.message_ts,
+            mode="append",
+            delimiter=delimiter,
+        )
+
+
+async def create_progress_message(
+    channel_id: str, thread_ts: str | None = None, initial_text: str = "🔄 Working..."
+) -> ProgressMessage:
+    """Helper function to create and start a progress message."""
+    progress = ProgressMessage(channel_id, thread_ts)
+    await progress.start(initial_text)
+    return progress

examples/slackbot/src/slackbot/wrap.py

@@ -1,5 +1,6 @@
 import inspect
 from contextlib import ContextDecorator
+from contextvars import ContextVar
 from functools import wraps
 from typing import Any, Callable, TypeVar
 
@@ -9,6 +10,8 @@
 
 T = TypeVar("T")
 
+_progress_message: ContextVar[Any] = ContextVar("progress_message", default=None)
+
 
 class DecorateMethodContext(ContextDecorator):
     """Context decorator for patching methods with a decorator."""
@@ -55,15 +58,31 @@ def _patch_method(self, cls, method_name, decorator):
 def prefect_wrapped_function(
     func: Callable[..., T],
     decorator: Callable[..., Callable[..., T]] = task,
-    tags: set | None = None,
+    tags: set[str] | None = None,
     settings: dict[str, Any] | None = None,
 ) -> Callable[..., Callable[..., T]]:
     """Decorator for wrapping a function with a prefect decorator."""
-    tags = tags or set()
+    tags = tags or set[str]()
 
     @wraps(func)
     async def wrapper(*args, **kwargs) -> T:
-        wrapped_callable = decorator(**settings or {})(func)  # type: ignore
+        if _progress := _progress_message.get():
+            tool_name = "Unknown Tool"
+            if args and hasattr(args[0], "name"):
+                tool_name = args[0].name
+            elif (
+                args
+                and hasattr(args[0], "function")
+                and hasattr(args[0].function, "__name__")
+            ):
+                tool_name = args[0].function.__name__
+
+            try:
+                await _progress.append(f"🔧 {tool_name}")
+            except Exception:
+                pass
+
+        wrapped_callable = decorator(**settings or {})(func)
         with prefect_tags(*tags):
             result = wrapped_callable(*args, **kwargs)  # type: ignore
             if inspect.isawaitable(result):