agents: prompt and explain-error

Author: echarles

Makefile

@@ -55,21 +55,21 @@ repl:
 # Adjust --max-requests based on your Azure tier (CLI default: 4; lower if you hit rate limits)
 prompt:
 	jupyter-ai-agents prompt \
-		--url http://localhost:8888 \
-		--token MY_TOKEN \
-		--model azure-openai:gpt-4o-mini \
+		--verbose \
+		--mcp-servers http://localhost:8888/mcp \
+		--model anthropic:claude-sonnet-4-20250514 \
 		--path notebook.ipynb \
-		--max-requests 4 \
+		--max-requests 20 \
 		--max-tool-calls 10 \
 		--input "Create a matplotlib example"
 
 explain-error:
 	jupyter-ai-agents explain-error \
-		--url http://localhost:8888 \
-		--token MY_TOKEN \
-		--model azure-openai:gpt-4o-mini \
+		--verbose \
+		--mcp-servers http://localhost:8888/mcp \
+		--model anthropic:claude-sonnet-4-20250514 \
 		--path notebook.ipynb \
-		--max-requests 4 \
+		--max-requests 20 \
 		--max-tool-calls 10
 
 publish-pypi: # publish the pypi package

dev/content/notebook.ipynb

@@ -3,7 +3,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "29a9d2af-9c57-449d-bbcb-2df77dfe9dfe",
+   "id": "6e785ae0-34d9-462e-aa9a-22d7b40e15d8",
    "metadata": {},
    "outputs": [],
    "source": []

jupyter_ai_agents/agents/cli/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2024-2025 Datalayer, Inc.
+#
+# BSD 3-Clause License
+
+"""Pydantic AI agents to explain error."""
+
+from jupyter_ai_agents.agents.explain_error.explain_error_agent import create_explain_error_agent, run_explain_error_agent
+
+__all__ = [
+    "create_explain_error_agent",
+    "run_explain_error_agent",
+]

jupyter_ai_agents/agents/explain_error/__init__.py

@@ -40,7 +40,9 @@
 Your response should:
 1. Briefly explain what caused the error
 2. Insert a corrected code cell at the appropriate location
-3. Execute it to verify the fix works"""
+3. Execute it to verify the fix works
+
+IMPORTANT: When you have completed the fix successfully, provide a final text response summarizing what you did WITHOUT making any more tool calls. This signals completion and allows the program to exit properly."""
 
 
 class ExplainErrorAgentDeps:
@@ -95,8 +97,11 @@ def create_explain_error_agent(
     if error_cell_index != -1:
         system_prompt += f"\n\nError occurred at cell index: {error_cell_index}"
 
-    # Add reminder to be efficient
-    system_prompt += "\n\nBe efficient: Fix the error in as few steps as possible."
+    # Add reminder to be efficient and complete properly
+    system_prompt += (
+        "\n\nBe efficient: Fix the error in as few steps as possible. "
+        "When the fix is complete, provide a final summary WITHOUT tool calls to signal completion."
+    )
 
     # Create agent with MCP toolset
     agent = Agent(
@@ -117,7 +122,9 @@ async def run_explain_error_agent(
     notebook_content: str = "",
     error_info: dict[str, Any] | None = None,
     error_cell_index: int = -1,
+    notebook_path: str = "",
     max_tool_calls: int = 10,
+    max_requests: int = 3,
 ) -> str:
     """
     Run the explain error agent to analyze and fix an error.
@@ -128,40 +135,71 @@ async def run_explain_error_agent(
         notebook_content: Content of notebook cells
         error_info: Additional error information
         error_cell_index: Index where error occurred
+        notebook_path: Path to the notebook file
         max_tool_calls: Maximum number of tool calls to prevent excessive API usage
+        max_requests: Maximum number of API requests (default: 3)
     
     Returns:
         Agent's response with explanation and fix
     """
     import asyncio
-    from pydantic_ai import UsageLimits
+    import os
+    from pydantic_ai import UsageLimitExceeded, UsageLimits
 
     deps = ExplainErrorAgentDeps(
         notebook_content=notebook_content,
         error_info=error_info,
         error_cell_index=error_cell_index,
     )
 
-    logger.info(f"Running explain error agent for error: {error_description[:50]}... (max_tool_calls={max_tool_calls})")
+    logger.info(f"Running explain error agent for error: {error_description[:50]}... (max_tool_calls={max_tool_calls}, max_requests={max_requests})")
+    
+    # Prepend notebook connection instruction if path is provided
+    if notebook_path:
+        notebook_name = os.path.splitext(os.path.basename(notebook_path))[0]
+        
+        # Prepend instruction to connect to the notebook first
+        enhanced_description = (
+            f"First, use the use_notebook tool to connect to the notebook at path '{notebook_path}' "
+            f"with notebook_name '{notebook_name}' and mode 'connect'. "
+            f"Then, analyze and fix this error: {error_description}"
+        )
+        logger.info(f"Enhanced input to connect to notebook: {notebook_path}")
+    else:
+        enhanced_description = error_description
 
     try:
         # Create usage limits to prevent excessive API calls
         # Use strict limits: fewer requests to avoid rate limiting
         usage_limits = UsageLimits(
             tool_calls_limit=max_tool_calls,
-            request_limit=3,  # Strict limit: max 3 API requests to avoid rate limiting
+            request_limit=max_requests,  # Strict limit to avoid rate limiting
         )
 
         # Add timeout to prevent hanging on retries
         result = await asyncio.wait_for(
-            agent.run(error_description, deps=deps, usage_limits=usage_limits),
+            agent.run(enhanced_description, deps=deps, usage_limits=usage_limits),
             timeout=120.0  # 2 minute timeout
         )
         logger.info("Explain error agent completed successfully")
         return result.response
     except asyncio.TimeoutError:
         logger.error("Explain error agent timed out after 120 seconds")
         return "Error: Operation timed out. The agent may have hit rate limits or is taking too long."
+    except UsageLimitExceeded as e:
+        logger.error(f"Explain error agent hit usage limits: {e}")
+        return (
+            "Error: Reached the configured usage limits.\n"
+            f"Increase --max-requests (currently {max_requests}) or --max-tool-calls (currently {max_tool_calls}) "
+            "if your model provider allows more usage."
+        )
+    except UsageLimitExceeded as e:
+        logger.error(f"Explain error agent hit usage limits: {e}")
+        return (
+            "Error: Reached the configured usage limits.\n"
+            f"Increase --max-requests (currently {max_requests}) or --max-tool-calls (currently {max_tool_calls}) "
+            "if your model provider allows more usage."
+        )
     except Exception as e:
         logger.error(f"Error running explain error agent: {e}", exc_info=True)
         return f"Error: {str(e)}"

jupyter_ai_agents/agents/cli/explain_error_agent.py renamed to jupyter_ai_agents/agents/explain_error/explain_error_agent.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2024-2025 Datalayer, Inc.
+#
+# BSD 3-Clause License
+
+"""Pydantic AI agents for prompt handling."""
+
+from jupyter_ai_agents.agents.prompt.prompt_agent import create_prompt_agent, run_prompt_agent
+
+__all__ = [
+    "create_prompt_agent",
+    "run_prompt_agent",
+]

jupyter_ai_agents/agents/prompt/__init__.py

@@ -37,7 +37,9 @@
 3. Write the main code
 4. Add markdown explanations
 
-Execute each code cell as you create it to ensure it works properly."""
+Execute each code cell as you create it to ensure it works properly.
+
+IMPORTANT: When you have completed the task successfully, provide a final text response summarizing what you did WITHOUT making any more tool calls. This signals completion and allows the program to exit properly."""
 
 
 class PromptAgentDeps:
@@ -83,8 +85,12 @@ def create_prompt_agent(
         if notebook_context.get('current_cell_index', -1) != -1:
             system_prompt += f"\n\nUser instruction was given at cell index: {notebook_context['current_cell_index']}"
 
-    # Add reminder to be efficient
-    system_prompt += "\n\nBe efficient: Complete the task in as few steps as possible. Don't over-verify or re-check unnecessarily."
+    # Add reminder to be efficient and complete properly
+    system_prompt += (
+        "\n\nBe efficient: Complete the task in as few steps as possible. "
+        "Don't over-verify or re-check unnecessarily. "
+        "When the task is done, provide a final summary WITHOUT tool calls to signal completion."
+    )
 
     # Create agent with MCP toolset
     agent = Agent(
@@ -112,20 +118,36 @@ async def run_prompt_agent(
     Args:
         agent: The configured prompt agent
         user_input: User's instruction/prompt
-        notebook_context: Optional notebook context
+        notebook_context: Optional notebook context (should include 'notebook_path')
         max_tool_calls: Maximum number of tool calls to prevent excessive API usage
     max_requests: Maximum number of API requests (default: 4, lower if needed for strict rate limits)
     
     Returns:
         Agent's response
     """
     import asyncio
+    import os
     from pydantic_ai import UsageLimitExceeded, UsageLimits
 
     deps = PromptAgentDeps(notebook_context)
 
     logger.info(f"Running prompt agent with input: {user_input[:50]}... (max_tool_calls={max_tool_calls}, max_requests={max_requests})")
 
+    # Prepend notebook connection instruction if path is provided
+    if notebook_context and notebook_context.get('notebook_path'):
+        notebook_path = notebook_context['notebook_path']
+        notebook_name = os.path.splitext(os.path.basename(notebook_path))[0]
+        
+        # Prepend instruction to connect to the notebook first
+        enhanced_input = (
+            f"First, use the use_notebook tool to connect to the notebook at path '{notebook_path}' "
+            f"with notebook_name '{notebook_name}' and mode 'connect'. "
+            f"Then, {user_input}"
+        )
+        logger.info(f"Enhanced input to connect to notebook: {notebook_path}")
+    else:
+        enhanced_input = user_input
+    
     # Warn about low limits
     if max_requests <= 2:
         logger.warning(
@@ -144,7 +166,7 @@ async def run_prompt_agent(
 
         # Add timeout to prevent hanging on retries
         result = await asyncio.wait_for(
-            agent.run(user_input, deps=deps, usage_limits=usage_limits),
+            agent.run(enhanced_input, deps=deps, usage_limits=usage_limits),
             timeout=120.0  # 2 minute timeout
         )
         logger.info("Prompt agent completed successfully")