hdsp-jupyter-extension 2.0.11__py3-none-any.whl → 2.0.13__py3-none-any.whl

agent_server/langchain/middleware/subagent_middleware.py

@@ -0,0 +1,329 @@
+"""
+SubAgentMiddleware
+
+Middleware that enables subagent delegation via the `task` tool.
+Based on Deep Agents library pattern (benchmarked, not installed).
+
+Key features:
+- Provides task(agent_name, description) tool for subagent invocation
+- Context isolation: subagents run in clean context
+- Synchronous execution: subagent returns result directly to caller
+- Nested subagent support: python_developer can call athena_query
+"""
+
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from langchain_core.tools import tool
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+# Global registry for subagent factories (set by AgentFactory)
+_subagent_factory = None
+_current_llm_config = None
+
+
+def set_subagent_factory(factory_func, llm_config: Dict[str, Any]):
+    """
+    Set the subagent factory function.
+    Called by AgentFactory during initialization.
+    """
+    global _subagent_factory, _current_llm_config
+    _subagent_factory = factory_func
+    _current_llm_config = llm_config
+    logger.info("SubAgentMiddleware factory initialized")
+
+
+def get_subagent_factory():
+    """Get the current subagent factory function."""
+    return _subagent_factory, _current_llm_config
+
+
+def create_task_tool(
+    caller_name: str,
+    allowed_subagents: Optional[List[str]] = None,
+):
+    """
+    Create a task tool for calling subagents.
+
+    The task tool executes subagents synchronously and returns their result.
+    Subagents like python_developer return generated code/analysis,
+    which the Main Agent can then execute if needed.
+
+    Args:
+        caller_name: Name of the agent creating this tool (for logging/validation)
+        allowed_subagents: Optional list of subagent names this agent can call.
+                          If None, all subagents are allowed (for Main Agent).
+
+    Returns:
+        A tool that can be used to delegate tasks to subagents.
+    """
+    from agent_server.langchain.subagents.base import (
+        SUBAGENT_CONFIGS,
+        get_subagent_config,
+    )
+
+    # Build description based on allowed subagents
+    if allowed_subagents:
+        available = [
+            f"- {name}: {SUBAGENT_CONFIGS[name].description}"
+            for name in allowed_subagents
+            if name in SUBAGENT_CONFIGS
+        ]
+    else:
+        # Main Agent (planner) can call non-restricted subagents OR those explicitly allowing "planner"
+        available = [
+            f"- {config.name}: {config.description}"
+            for config in SUBAGENT_CONFIGS.values()
+            if not config.callable_by or caller_name in config.callable_by
+        ]
+
+    available_str = "\n".join(available)
+
+    # Create Pydantic schema for the task tool (required for Gemini compatibility)
+    class TaskInput(BaseModel):
+        """Input schema for task tool"""
+
+        agent_name: str = Field(
+            description=f"Name of the subagent to invoke. Available: {', '.join(allowed_subagents) if allowed_subagents else 'python_developer, researcher, athena_query'}"
+        )
+        description: str = Field(
+            description="Detailed task description for the subagent (Korean preferred)"
+        )
+        context: Optional[str] = Field(
+            default=None,
+            description="Additional context for the subagent: resource info (file sizes, memory), previous code, variable state, etc."
+        )
+
+    @tool(args_schema=TaskInput)
+    def task_tool(agent_name: str, description: str, context: Optional[str] = None) -> str:
+        """
+        Delegate a task to a specialized subagent.
+
+        The subagent will execute the task and return its result (code, analysis, etc.).
+        Code execution tools (jupyter_cell_tool, write_file_tool) are handled by Main Agent.
+
+        Args:
+            agent_name: Name of the subagent to invoke
+            description: Detailed task description for the subagent
+            context: Additional context (resource info, previous code, etc.)
+
+        Returns:
+            Result from the subagent execution (string summary or generated code)
+        """
+        # Validate subagent exists
+        if agent_name not in SUBAGENT_CONFIGS:
+            return f"Error: Unknown agent '{agent_name}'. Available agents:\n{available_str}"
+
+        # Validate caller is allowed to call this subagent
+        config = get_subagent_config(agent_name)
+        if allowed_subagents and agent_name not in allowed_subagents:
+            return f"Error: '{caller_name}' cannot call '{agent_name}'. Allowed: {allowed_subagents}"
+
+        if config.callable_by and caller_name not in config.callable_by:
+            return f"Error: '{agent_name}' can only be called by: {config.callable_by}"
+
+        logger.info(
+            f"[{caller_name}] Invoking subagent '{agent_name}': {description[:100]}..."
+        )
+
+        # Import subagent event emitters
+        from agent_server.langchain.middleware.subagent_events import (
+            emit_subagent_start,
+            emit_subagent_complete,
+            set_current_subagent,
+            clear_current_subagent,
+        )
+
+        # Emit subagent start event for UI
+        emit_subagent_start(agent_name, description)
+
+        # Get the factory and config
+        factory_func, llm_config = get_subagent_factory()
+        if factory_func is None:
+            return "Error: SubAgentMiddleware not initialized. Call set_subagent_factory first."
+
+        try:
+            # Set current subagent context for tool call tracking
+            set_current_subagent(agent_name)
+
+            # Create the subagent
+            subagent = factory_func(agent_name, llm_config)
+
+            # Execute subagent synchronously with clean context
+            # The subagent runs in isolation, receiving task description + optional context
+            import uuid
+
+            subagent_thread_id = f"subagent-{agent_name}-{uuid.uuid4().hex[:8]}"
+            subagent_config = {
+                "configurable": {
+                    "thread_id": subagent_thread_id,
+                }
+            }
+
+            # Inject code history for python_developer
+            enhanced_context = context
+            if agent_name == "python_developer":
+                try:
+                    from agent_server.langchain.middleware.code_history_middleware import (
+                        get_context_with_history,
+                        get_code_history_tracker,
+                    )
+                    tracker = get_code_history_tracker()
+                    if tracker.get_entry_count() > 0:
+                        enhanced_context = get_context_with_history(context)
+                        logger.info(
+                            f"[{caller_name}] Injected code history into context "
+                            f"(entries={tracker.get_entry_count()}, "
+                            f"context_len={len(enhanced_context) if enhanced_context else 0})"
+                        )
+                except Exception as e:
+                    logger.warning(f"Failed to inject code history: {e}")
+
+            # Build the message content with optional context
+            if enhanced_context:
+                message_content = f"""## Task
+{description}
+
+## Context (provided by Main Agent)
+{enhanced_context}"""
+            else:
+                message_content = description
+
+            logger.info(f"[{caller_name}] Subagent message length: {len(message_content)}")
+
+            # Execute the subagent
+            result = subagent.invoke(
+                {"messages": [{"role": "user", "content": message_content}]},
+                config=subagent_config,
+            )
+
+            # Extract the final message from the result
+            messages = result.get("messages", [])
+            if messages:
+                final_message = messages[-1]
+                if hasattr(final_message, "content"):
+                    response = final_message.content
+                else:
+                    response = str(final_message)
+            else:
+                response = "Subagent completed but returned no messages."
+
+            logger.info(
+                f"[{caller_name}] Subagent '{agent_name}' returned: {str(response)[:200]}..."
+            )
+
+            # Extract description from python_developer response for auto-injection
+            if agent_name == "python_developer":
+                try:
+                    from agent_server.langchain.middleware.description_injector import (
+                        process_task_tool_response,
+                    )
+                    process_task_tool_response(agent_name, str(response))
+                except Exception as e:
+                    logger.warning(f"Failed to extract description: {e}")
+
+            # Emit subagent complete event for UI
+            emit_subagent_complete(agent_name, str(response)[:100])
+
+            return response
+
+        except Exception as e:
+            error_msg = f"Subagent '{agent_name}' failed: {str(e)}"
+            logger.error(error_msg, exc_info=True)
+            # Emit complete event even on error
+            emit_subagent_complete(agent_name, f"Error: {str(e)[:50]}")
+            return f"Error: {error_msg}"
+        finally:
+            # Always clear subagent context
+            clear_current_subagent()
+
+    # Update tool docstring with available agents
+    task_tool.__doc__ = f"""Delegate a task to a specialized subagent.
+
+Available agents:
+{available_str}
+
+The subagent will analyze the task and return its result.
+- python_developer: Returns generated Python code and analysis
+- researcher: Returns search results and findings
+- athena_query: Returns SQL query string
+
+For code execution (running Python, writing files), use Main Agent's tools directly.
+
+IMPORTANT: For python_developer, ALWAYS provide context with:
+- Resource info (file sizes, memory) from check_resource_tool
+- Previous code context if building on existing work
+- Variable names and their current state
+
+Args:
+    agent_name: Name of the subagent to invoke
+    description: Detailed task description for the subagent
+    context: Additional context (resource info, previous code, etc.)
+
+Returns:
+    Result from the subagent execution
+"""
+
+    return task_tool
+
+
+class SubAgentMiddleware:
+    """
+    Middleware that adds subagent delegation capability.
+
+    This middleware:
+    1. Adds the `task` tool to the agent's toolset
+    2. The task tool executes subagents synchronously
+    3. Subagent results are returned directly to the caller
+
+    Usage:
+        middleware = SubAgentMiddleware(
+            caller_name="main_agent",
+            allowed_subagents=["python_developer", "researcher"],
+        )
+
+        agent = create_agent(
+            model=llm,
+            tools=tools,
+            middleware=[middleware, ...],
+        )
+    """
+
+    def __init__(
+        self,
+        caller_name: str,
+        allowed_subagents: Optional[List[str]] = None,
+    ):
+        """
+        Initialize SubAgentMiddleware.
+
+        Args:
+            caller_name: Name of the agent using this middleware
+            allowed_subagents: List of subagents this agent can call.
+                             None means all non-restricted subagents.
+        """
+        self.caller_name = caller_name
+        self.allowed_subagents = allowed_subagents
+        self.task_tool = create_task_tool(caller_name, allowed_subagents)
+
+        logger.info(
+            f"SubAgentMiddleware initialized for '{caller_name}' "
+            f"with allowed_subagents={allowed_subagents}"
+        )
+
+    def get_tools(self) -> List[Any]:
+        """Get the tools provided by this middleware."""
+        return [self.task_tool]
+
+    def __call__(self, tools: List[Any]) -> List[Any]:
+        """
+        Add task tool to the agent's toolset.
+
+        This is called during agent creation to augment the tool list.
+        """
+        return tools + [self.task_tool]

agent_server/langchain/prompts.py

@@ -1,67 +1,62 @@
 """
 Prompt templates for LangChain agent.
-
-Contains system prompts, JSON schema for fallback tool calling,
-and middleware-specific prompts.
 """
 
-DEFAULT_SYSTEM_PROMPT = """You are an expert Python data scientist and Jupyter notebook assistant. Respond in Korean only.
-
-# Core Rules
-1. Be concise (≤4 lines unless detail requested)
-2. ALWAYS call a tool in every response - never respond with text only
-3. ALWAYS include a brief Korean explanation before tool calls
-
-# Task Workflow
-
-## Simple Tasks (1-2 steps)
-Execute directly without todos.
-
-## Complex Tasks (3+ steps)
-1. Create todos with write_todos (all items in Korean)
-2. ALWAYS include "작업 요약 및 다음단계 제시" as the LAST item
-3. After each tool result: check todos → call next tool → repeat
-4. **Final todo ("작업 요약 및 다음단계 제시")**:
-   - FIRST: Output summary JSON in your content (REQUIRED!)
-   - THEN: Call write_todos to mark all as completed
-   - Both must be in the SAME response
-
-### Summary JSON Format (MUST output before marking complete)
-```json
-{"summary": "실행된 작업 요약", "next_items": [{"subject": "제목", "description": "설명"}]}
-```
-Suggest 3-5 next items. **You CANNOT mark "작업 요약" as completed without outputting this JSON first.**
-
-# Mandatory Checks
-
-## Resource Check (BEFORE data operations)
-Call `check_resource_tool` FIRST when:
-- Loading files (.csv, .parquet, .json, .xlsx, .pickle, .h5, .feather)
-- Using pandas/polars/dask for dataframes
-- Training ML models
-
-# Tool Usage
-
-## File Search (execute_command_tool)
-```bash
-find . -iname '*filename*.csv' 2>/dev/null     # Find by name
-grep -rn 'pattern' --include='*.py' .           # Search contents
-```
-
-## File Reading (read_file_tool)
-- Large files: `read_file_tool(path, limit=100)` first
-- Use `offset` for pagination
-- Small files (<500 lines): Read without limit
-
-## Code Output
-- For plots/charts: Use English labels only
-- Use LSP tools for error checking and symbol lookup
-- Use multiedit_file_tool for multiple changes
-
-# Forbidden
-- Empty responses (no tool call AND no content)
-- Tool calls without Korean explanation
-- Stopping with pending/in_progress todos
+DEFAULT_SYSTEM_PROMPT = """You are an expert Python data scientist and Jupyter notebook assistant.
+
+# 핵심 규칙
+1. 한국어로 응답하세요
+2. 간결하게 (4줄 이하, 상세 요청 시 예외)
+3. 모든 응답에 도구를 호출하세요 (텍스트만 응답 금지)
+4. **[필수] 도구 호출 시 반드시 텍스트 설명을 함께 출력**
+   - 예: "파일 구조를 확인하겠습니다." + tool_call
+   - 설명 없이 tool_call만 하면 안 됨
+
+# 작업 흐름
+1. 간단한 작업 (1-2단계): 바로 실행
+2. 복잡한 작업 (3단계+): write_todos로 계획 → 순차 실행
+
+# write_todos 규칙 [필수]
+- 한국어로 작성
+- **🔴 기존 todo 절대 삭제 금지**: 전체 리스트를 항상 포함하고 status만 변경
+  - 잘못된 예: [{"content": "작업 요약", "status": "completed"}] ← 기존 todo 삭제됨!
+  - 올바른 예: [{"content": "기존 작업1", "status": "completed"}, {"content": "기존 작업2", "status": "completed"}, {"content": "작업 요약", "status": "completed"}]
+- **일괄 업데이트**: 연속 완료된 todo는 한 번의 write_todos 호출로 처리
+- in_progress는 **1개만** 유지
+- **[필수] 마지막 todo는 반드시 "작업 요약 및 다음 단계 제시"로 생성**
+- **🔴 [실행 순서 필수]**: "작업 요약 및 다음 단계 제시"는 **반드시 가장 마지막에 실행**
+  - 다른 모든 todo가 completed 상태가 된 후에만 이 todo를 in_progress로 변경
+  - 비슷한 이름의 다른 작업(보고서 검토, 결과 정리 등)과 혼동 금지
+- **[중요] "작업 요약 및 다음 단계 제시"는 summary JSON 출력 후에만 completed 표시**
+
+# 모든 작업 완료 후 [필수]
+마지막 todo "작업 요약 및 다음 단계 제시"를 completed로 변경한 후,
+**텍스트 응답으로** 아래 JSON을 출력하세요 (todo content가 아님!):
+{"summary": "완료된 작업 요약", "next_items": [{"subject": "제목", "description": "설명"}]}
+- next_items 3개 이상 필수
+- **summary JSON 없이 종료 금지**
+- **주의**: JSON은 todo 항목이 아닌 일반 텍스트 응답으로 출력
+
+# 도구 사용
+- check_resource_tool: 대용량 파일/데이터프레임 작업 전 필수
+- read_file_tool: 대용량 파일은 limit=100으로 먼저 확인
+- jupyter_cell_tool: 차트 라벨은 영어로
+- **파일 수정 후**: diagnostics_tool로 오류 확인 필수
+
+# 사용자 입력 요청 [중요]
+- **ask_user_tool**: 사용자 응답이 필요할 때 사용 (파일 업로드, 선택, 정보 요청)
+  - 파일 업로드: ask_user_tool(question="kaggle.json 파일을 업로드해 주세요", input_type="file")
+  - 선택 요청: ask_user_tool(question="모델 선택?", options=["Logistic", "RandomForest"])
+  - 정보 요청: ask_user_tool(question="API 키를 입력해 주세요", input_type="text")
+- **markdown_tool**: 정보 출력용 (사용자 응답 불필요)
+- ⚠️ 사용자 응답이 필요하면 markdown_tool 대신 반드시 ask_user_tool 사용!
+
+# 금지 사항
+- 빈 응답 (도구 호출도 없고 내용도 없음)
+- 설명 없이 도구만 호출
+- pending/in_progress todo 남기고 종료
+- "작업 요약 및 다음 단계 제시" todo 없이 todo 리스트 생성
+- **🔴 다른 pending todo가 있는데 "작업 요약 및 다음 단계 제시"를 먼저 실행** (순서 위반)
 """
 
 JSON_TOOL_SCHEMA = """Respond with ONLY valid JSON:
@@ -70,61 +65,61 @@ JSON_TOOL_SCHEMA = """Respond with ONLY valid JSON:
 Tools:
 - jupyter_cell_tool: {"code": "<python>"}
 - markdown_tool: {"content": "<markdown>"}
-- write_todos: {"todos": [{"content": "한국어 내용", "status": "pending|in_progress|completed"}]}
+- ask_user_tool: {"question": "<질문>", "options": ["선택1", "선택2"], "input_type": "text|file|selection"}
+- write_todos: {"todos": [{"content": "내용", "status": "pending|in_progress|completed"}]}
 - read_file_tool: {"path": "<path>", "offset": 0, "limit": 500}
-- write_file_tool: {"path": "<path>", "content": "<content>", "overwrite": false}
-- search_notebook_cells_tool: {"pattern": "<regex>"}
+- write_file_tool: {"path": "<path>", "content": "<content>"}
+- multiedit_file_tool: {"path": "<path>", "edits": [{"old_string": "...", "new_string": "..."}]}
 - execute_command_tool: {"command": "<cmd>"}
 - check_resource_tool: {"files": ["<path>"], "dataframes": ["<var>"]}
 
 No markdown wrapping. JSON only."""
 
-TODO_LIST_SYSTEM_PROMPT = """
-# Todo Rules
-
-## New Message = Fresh Start
-- Each user message is a NEW task
-- Ignore completion status from chat history
-- Execute ALL current todos from scratch
-
-## Structure
-All todo items must be in Korean. Always end with:
-- 작업 요약 및 다음단계 제시  ← 필수 마지막 항목!
-
-## Workflow
-1. Find pending/in_progress todo
-2. Execute it NOW in THIS response
-3. Mark completed
-4. Repeat until all done
-
-## 🔴 Final Todo ("작업 요약 및 다음단계 제시") - CRITICAL
-When executing this todo, you MUST:
-1. Output the summary JSON in your content FIRST:
-   {"summary": "작업 내용 요약", "next_items": [{"subject": "...", "description": "..."}]}
-2. THEN call write_todos to mark all as completed
-3. If you don't output the JSON, the todo will NOT be marked as completed!
-
-## Completion Check
-- ✅ Done: Executed in THIS response
-- ❌ Not done: Only visible in chat history
-- ❌ "작업 요약" cannot be completed without outputting summary JSON
-
-## Forbidden
-- Marking complete without executing in THIS response
-- Marking "작업 요약" complete without outputting JSON summary
-- Todos without "작업 요약 및 다음단계 제시" as final item
+# Merged into DEFAULT_SYSTEM_PROMPT
+TODO_LIST_SYSTEM_PROMPT = ""
+
+TODO_LIST_TOOL_DESCRIPTION = """Todo 리스트 관리 도구.
+
+사용 시점:
+- 3단계 이상의 복잡한 작업
+- 진행 상황 추적이 필요할 때
+
+규칙:
+- in_progress는 1개만
+- 완료 즉시 completed로 변경
+- **[필수] 마지막 todo는 반드시 "작업 요약 및 다음 단계 제시"로 생성**
+- **🔴 [실행 순서]**: todo는 반드시 리스트 순서대로 실행하고, "작업 요약 및 다음 단계 제시"는 맨 마지막에 실행
+- 이 "작업 요약 및 다음 단계 제시" todo 에서는 전체 작업 요약과 다음 단계를 제시하는 내용을 JSON 형태로 출력:
+  {"summary": "완료 요약", "next_items": [{"subject": "...", "description": "..."}]}
+  (next_items 3개 이상 필수)
 """
 
-TODO_LIST_TOOL_DESCRIPTION = """Update task list for tracking progress.
-This tool ONLY tracks status - does NOT execute tasks.
-After calling: immediately call next action tool (unless ALL completed)."""
+# List of tools available to the agent
+TOOL_LIST = [
+    "jupyter_cell_tool",
+    "markdown_tool",
+    "ask_user_tool",
+    "write_todos",
+    "read_file_tool",
+    "write_file_tool",
+    "multiedit_file_tool",
+    "search_notebook_cells_tool",
+    "execute_command_tool",
+    "check_resource_tool",
+    "list_workspace_tool",
+    "search_files_tool",
+]
 
-# Non-HITL tools that execute immediately without user approval
+# Tools that don't require HITL (Human-in-the-Loop) approval
 NON_HITL_TOOLS = {
     "markdown_tool",
     "markdown",
     "read_file_tool",
     "read_file",
+    "list_files_tool",
+    "list_files",
+    "search_workspace_tool",
+    "search_workspace",
     "search_notebook_cells_tool",
     "search_notebook_cells",
     "write_todos",