Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
Commit ·
44be1d7
1
Parent(s): da0136f
feat: emit tool_log events from research sub-agent for UI liveness
Browse filesWithout this, the UI shows "research tool called" then goes silent
for 30-120s while the sub-agent runs internally. Now emits progress
events: "Starting research sub-agent...", each tool call name, and
"Research complete." — so the user sees what's happening.
- agent/tools/research_tool.py +29 -0
agent/tools/research_tool.py
CHANGED
|
@@ -14,6 +14,8 @@ from typing import Any
|
|
| 14 |
|
| 15 |
from litellm import Message, acompletion
|
| 16 |
|
|
|
|
|
|
|
| 17 |
logger = logging.getLogger(__name__)
|
| 18 |
|
| 19 |
# Tools the research agent can use (read-only subset)
|
|
@@ -210,6 +212,17 @@ async def research_handler(
|
|
| 210 |
if spec["function"]["name"] in RESEARCH_TOOL_NAMES
|
| 211 |
]
|
| 212 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 213 |
# Run the research loop (max 20 iterations — research should be focused)
|
| 214 |
max_iterations = 20
|
| 215 |
for _iteration in range(max_iterations):
|
|
@@ -231,6 +244,7 @@ async def research_handler(
|
|
| 231 |
|
| 232 |
# If no tool calls, we have our final answer
|
| 233 |
if not msg.tool_calls:
|
|
|
|
| 234 |
content = msg.content or "Research completed but no summary generated."
|
| 235 |
return content, True
|
| 236 |
|
|
@@ -263,6 +277,21 @@ async def research_handler(
|
|
| 263 |
continue
|
| 264 |
|
| 265 |
try:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 266 |
output, _success = await session.tool_router.call_tool(
|
| 267 |
tool_name, tool_args, session=session
|
| 268 |
)
|
|
|
|
| 14 |
|
| 15 |
from litellm import Message, acompletion
|
| 16 |
|
| 17 |
+
from agent.core.session import Event
|
| 18 |
+
|
| 19 |
logger = logging.getLogger(__name__)
|
| 20 |
|
| 21 |
# Tools the research agent can use (read-only subset)
|
|
|
|
| 212 |
if spec["function"]["name"] in RESEARCH_TOOL_NAMES
|
| 213 |
]
|
| 214 |
|
| 215 |
+
async def _log(text: str) -> None:
|
| 216 |
+
"""Send a progress event to the UI so it doesn't look frozen."""
|
| 217 |
+
try:
|
| 218 |
+
await session.send_event(
|
| 219 |
+
Event(event_type="tool_log", data={"tool": "research", "log": text})
|
| 220 |
+
)
|
| 221 |
+
except Exception:
|
| 222 |
+
pass
|
| 223 |
+
|
| 224 |
+
await _log("Starting research sub-agent...")
|
| 225 |
+
|
| 226 |
# Run the research loop (max 20 iterations — research should be focused)
|
| 227 |
max_iterations = 20
|
| 228 |
for _iteration in range(max_iterations):
|
|
|
|
| 244 |
|
| 245 |
# If no tool calls, we have our final answer
|
| 246 |
if not msg.tool_calls:
|
| 247 |
+
await _log("Research complete.")
|
| 248 |
content = msg.content or "Research completed but no summary generated."
|
| 249 |
return content, True
|
| 250 |
|
|
|
|
| 277 |
continue
|
| 278 |
|
| 279 |
try:
|
| 280 |
+
# Brief description for the UI
|
| 281 |
+
brief = tool_name
|
| 282 |
+
if tool_name == "github_find_examples":
|
| 283 |
+
brief = f"github_find_examples({tool_args.get('repo', '')}/{tool_args.get('keyword', '')})"
|
| 284 |
+
elif tool_name == "github_read_file":
|
| 285 |
+
brief = f"github_read_file({tool_args.get('path', '')})"
|
| 286 |
+
elif tool_name == "explore_hf_docs":
|
| 287 |
+
brief = f"explore_hf_docs({tool_args.get('endpoint', '')})"
|
| 288 |
+
elif tool_name == "fetch_hf_docs":
|
| 289 |
+
url = tool_args.get("url", "")
|
| 290 |
+
brief = f"fetch_hf_docs(...{url[-50:]})" if len(url) > 50 else f"fetch_hf_docs({url})"
|
| 291 |
+
elif tool_name == "hf_inspect_dataset":
|
| 292 |
+
brief = f"hf_inspect_dataset({tool_args.get('dataset', '')})"
|
| 293 |
+
await _log(brief)
|
| 294 |
+
|
| 295 |
output, _success = await session.tool_router.call_tool(
|
| 296 |
tool_name, tool_args, session=session
|
| 297 |
)
|