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

agent_server/langchain/__init__.py

@@ -12,7 +12,7 @@ Architecture:
 - state.py: Agent state management
 """
 
-from agent_server.langchain.agent import create_simple_chat_agent
+from agent_server.langchain.agent import create_simple_chat_agent, create_agent_system
 from agent_server.langchain.state import AgentState
 
-__all__ = ["create_simple_chat_agent", "AgentState"]
+__all__ = ["create_simple_chat_agent", "create_agent_system", "AgentState"]

agent_server/langchain/agent.py

@@ -2,10 +2,11 @@
 LangChain Agent
 
 Main agent creation module for tool-driven chat execution.
+Supports both single-agent and multi-agent modes.
 """
 
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Literal
 
 from agent_server.langchain.custom_middleware import (
     create_handle_empty_response_middleware,
@@ -22,6 +23,7 @@ from agent_server.langchain.prompts import (
     TODO_LIST_TOOL_DESCRIPTION,
 )
 from agent_server.langchain.tools import (
+    ask_user_tool,
     check_resource_tool,
     diagnostics_tool,
     edit_file_tool,
@@ -43,6 +45,7 @@ def _get_all_tools():
     return [
         jupyter_cell_tool,
         markdown_tool,
+        ask_user_tool,  # HITL - waits for user response
         read_file_tool,
         write_file_tool,
         edit_file_tool,
@@ -134,9 +137,10 @@ def create_simple_chat_agent(
     middleware.append(patch_tool_calls)
 
     # Add TodoListMiddleware for task planning
+    # NOTE: system_prompt removed to avoid multi-part content array
+    # that causes Gemini MALFORMED_FUNCTION_CALL error
     if enable_todo_list:
         todo_middleware = TodoListMiddleware(
-            system_prompt=TODO_LIST_SYSTEM_PROMPT,
             tool_description=TODO_LIST_TOOL_DESCRIPTION,
         )
         middleware.append(todo_middleware)
@@ -192,42 +196,16 @@ def create_simple_chat_agent(
     else:
         system_prompt = DEFAULT_SYSTEM_PROMPT
 
-    # Add Gemini 2.5 Flash specific prompt to ensure content is included with tool calls
-    gemini_model = llm_config.get("gemini", {}).get("model", "")
-    if "gemini-2.5-flash" in gemini_model:
-        gemini_content_prompt = """
-## 🔴 IMPORTANT: Always include explanation text
-When calling any tool, you MUST include a brief explanation in your response content.
-NEVER produce an empty content when making tool calls.
-Before each tool call, write Korean explanations of what you're about to do.
-Example: "데이터를 로드하겠습니다." then call jupyter_cell_tool.
-"""
-        system_prompt = system_prompt + "\n" + gemini_content_prompt
-        logger.info("Added Gemini 2.5 Flash specific prompt for content inclusion")
-
-    # Add vLLM/gpt-oss specific prompt for Korean responses and proper todo structure
+    # Add vLLM/gpt-oss specific prompt (minimal - main rules are in DEFAULT_SYSTEM_PROMPT)
     provider = llm_config.get("provider", "")
     if provider == "vllm":
         vllm_prompt = """
-## 🔴 중요: 한국어로 응답하세요
-- 모든 응답, 설명, todo 항목은 반드시 한국어로 작성하세요.
-- 코드 주석과 출력 설명도 한국어로 작성하세요.
-- 영어로 응답하지 마세요.
-
-## 🔴 MANDATORY: Todo List Structure
-When creating todos with write_todos, you MUST:
-1. Write all todo items in Korean
-2. ALWAYS include "작업 요약 및 다음단계 제시" as the LAST todo item
-3. Example structure:
-   - 데이터 로드 및 확인
-   - 데이터 분석 수행
-   - 작업 요약 및 다음단계 제시  ← 반드시 마지막에 포함!
-
-## 🔴 IMPORTANT: Never return empty responses
-If you have nothing to say, call a tool instead. NEVER return an empty response.
+## 🔴 추가 강조
+- 반드시 한국어로 응답
+- 빈 응답 절대 금지 - 항상 도구 호출 필수
 """
         system_prompt = system_prompt + "\n" + vllm_prompt
-        logger.info("Added vLLM/gpt-oss specific prompt for Korean responses")
+        logger.info("Added vLLM-specific prompt")
 
     logger.info("SimpleChatAgent system_prompt: %s", system_prompt)
 
@@ -241,3 +219,64 @@ If you have nothing to say, call a tool instead. NEVER return an empty response.
     )
 
     return agent
+
+
+# =============================================================================
+# Multi-Agent System Support
+# =============================================================================
+
+
+def create_agent_system(
+    llm_config: Dict[str, Any],
+    workspace_root: str = ".",
+    enable_hitl: bool = True,
+    enable_todo_list: bool = True,
+    checkpointer: Optional[object] = None,
+    system_prompt_override: Optional[str] = None,
+    agent_mode: Literal["single", "multi"] = "single",
+    agent_prompts: Optional[Dict[str, str]] = None,
+):
+    """
+    Create an agent system based on the specified mode.
+    
+    This is the main entry point for creating agents. It supports both:
+    - "single": Traditional single-agent with all tools (backward compatible)
+    - "multi": Multi-agent system with Planner supervising subagents
+    
+    Args:
+        llm_config: LLM configuration
+        workspace_root: Root directory for file operations
+        enable_hitl: Enable Human-in-the-Loop for code execution
+        enable_todo_list: Enable TodoListMiddleware for task planning
+        checkpointer: Optional checkpointer for state persistence
+        system_prompt_override: Optional custom system prompt
+        agent_mode: "single" for traditional mode, "multi" for multi-agent
+        agent_prompts: Optional dict of per-agent prompts (for multi-agent mode)
+    
+    Returns:
+        Configured agent (single or Planner for multi-agent)
+    """
+    if agent_mode == "multi":
+        # Use multi-agent system with Planner + Subagents
+        from agent_server.langchain.agent_factory import create_multi_agent_system
+        
+        logger.info("Creating multi-agent system (Planner + Subagents)")
+        return create_multi_agent_system(
+            llm_config=llm_config,
+            checkpointer=checkpointer,
+            enable_hitl=enable_hitl,
+            enable_todo_list=enable_todo_list,
+            system_prompt_override=system_prompt_override,
+            agent_prompts=agent_prompts,
+        )
+    else:
+        # Use traditional single-agent mode
+        logger.info("Creating single-agent system (all tools in one agent)")
+        return create_simple_chat_agent(
+            llm_config=llm_config,
+            workspace_root=workspace_root,
+            enable_hitl=enable_hitl,
+            enable_todo_list=enable_todo_list,
+            checkpointer=checkpointer,
+            system_prompt_override=system_prompt_override,
+        )

agent_server/langchain/agent_factory.py

@@ -0,0 +1,400 @@
+"""
+Agent Factory
+
+Main module for creating agents in the multi-agent architecture.
+Based on Deep Agents library pattern (benchmarked, not installed).
+
+Provides:
+- create_main_agent(): Main supervisor agent (executes code, orchestrates)
+- create_subagent(): Specialized subagents (python_developer, researcher, athena_query)
+- create_multi_agent_system(): Complete multi-agent setup
+
+Architecture:
+- Main Agent: Executes code (jupyter_cell_tool, write_file_tool) and orchestrates
+- Subagents: Generate code/analysis but do NOT execute (like Athena Query generates SQL)
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from agent_server.langchain.agent_prompts.planner_prompt import (
+    PLANNER_SYSTEM_PROMPT,
+)
+from agent_server.langchain.custom_middleware import (
+    create_continuation_control_middleware,
+    create_handle_empty_response_middleware,
+    create_limit_tool_calls_middleware,
+    create_normalize_tool_args_middleware,
+    create_patch_tool_calls_middleware,
+)
+from agent_server.langchain.hitl_config import get_hitl_interrupt_config
+from agent_server.langchain.llm_factory import create_llm, create_summarization_llm
+from agent_server.langchain.middleware.skill_middleware import get_skill_middleware
+from agent_server.langchain.middleware.subagent_middleware import (
+    create_task_tool,
+    set_subagent_factory,
+)
+from agent_server.langchain.subagents.base import (
+    get_available_subagents_for_planner,
+    get_subagent_config,
+)
+from agent_server.langchain.tools.tool_registry import get_tools_for_agent
+
+logger = logging.getLogger(__name__)
+
+
+def create_subagent(
+    agent_type: str,
+    llm_config: Dict[str, Any],
+    checkpointer: Optional[object] = None,
+    enable_hitl: bool = True,
+    system_prompt_override: Optional[str] = None,
+) -> Any:
+    """
+    Create a specialized subagent.
+
+    Subagents are stateless and run in isolated context.
+    They generate code/analysis but do NOT execute - Main Agent handles execution.
+
+    Args:
+        agent_type: Type of subagent (python_developer, researcher, athena_query)
+        llm_config: LLM configuration
+        checkpointer: Optional checkpointer (usually None for subagents)
+        enable_hitl: Enable Human-in-the-Loop for tools that need approval
+        system_prompt_override: Optional custom system prompt for this subagent
+
+    Returns:
+        Compiled agent graph
+    """
+    try:
+        from langchain.agents import create_agent
+        from langchain.agents.middleware import (
+            HumanInTheLoopMiddleware,
+            ModelCallLimitMiddleware,
+        )
+        from langgraph.checkpoint.memory import InMemorySaver
+    except ImportError as e:
+        logger.error(f"Failed to import LangChain components: {e}")
+        raise
+
+    # Get subagent configuration
+    config = get_subagent_config(agent_type)
+
+    # Create LLM (may use override from config)
+    llm = create_llm(llm_config)
+
+    # Get tools for this agent type
+    tools = get_tools_for_agent(agent_type)
+
+    # Configure middleware
+    middleware = []
+
+    # Add HITL middleware if enabled (for tools that need approval)
+    if enable_hitl:
+        hitl_middleware = HumanInTheLoopMiddleware(
+            interrupt_on=get_hitl_interrupt_config(),
+            description_prefix=f"[{agent_type}] Tool execution pending approval",
+        )
+        middleware.append(hitl_middleware)
+
+    # Add nested subagent support for python_developer
+    if config.can_call_subagents:
+        # Python Developer can call athena_query
+        task_tool = create_task_tool(
+            caller_name=agent_type,
+            allowed_subagents=config.can_call_subagents,
+        )
+        tools = tools + [task_tool]
+        logger.info(f"Subagent '{agent_type}' can call: {config.can_call_subagents}")
+
+    # Add SkillMiddleware for python_developer (code generation agent)
+    skill_prompt_section = ""
+    if agent_type == "python_developer":
+        skill_middleware = get_skill_middleware()
+        tools = tools + skill_middleware.get_tools()
+        skill_prompt_section = skill_middleware.get_prompt_section()
+        logger.info(
+            f"SkillMiddleware added to '{agent_type}': "
+            f"{len(skill_middleware.skills)} skills available"
+        )
+
+    # Add model call limit to prevent infinite loops
+    model_limit = ModelCallLimitMiddleware(
+        run_limit=15,  # Subagents have lower limit
+        exit_behavior="end",
+    )
+    middleware.append(model_limit)
+
+    # Determine system prompt (use override if provided)
+    base_prompt = (
+        system_prompt_override if system_prompt_override else config.system_prompt
+    )
+
+    # Inject skill prompt section for python_developer
+    if skill_prompt_section:
+        system_prompt = f"{base_prompt}\n\n{skill_prompt_section}"
+        logger.info(
+            f"Injected skills prompt section ({len(skill_prompt_section)} chars)"
+        )
+    else:
+        system_prompt = base_prompt
+
+    # Create the subagent
+    agent = create_agent(
+        model=llm,
+        tools=tools,
+        middleware=middleware,
+        checkpointer=checkpointer or InMemorySaver(),  # Needed for HITL
+        system_prompt=system_prompt,
+    )
+
+    logger.info(
+        f"Created subagent '{agent_type}' with {len(tools)} tools, "
+        f"HITL={enable_hitl}, nested_calls={config.can_call_subagents}"
+    )
+
+    return agent
+
+
+def create_main_agent(
+    llm_config: Dict[str, Any],
+    checkpointer: Optional[object] = None,
+    enable_hitl: bool = True,
+    enable_todo_list: bool = True,
+    system_prompt_override: Optional[str] = None,
+    agent_prompts: Optional[Dict[str, str]] = None,
+) -> Any:
+    """
+    Create the Main Agent (Supervisor).
+
+    The Main Agent is the central agent that:
+    - Analyzes user requests
+    - Creates todo lists
+    - Executes code directly (jupyter_cell_tool, write_file_tool, etc.)
+    - Delegates analysis/generation tasks to subagents via task()
+    - Synthesizes results
+
+    Args:
+        llm_config: LLM configuration
+        checkpointer: Checkpointer for state persistence
+        enable_hitl: Enable Human-in-the-Loop
+        enable_todo_list: Enable TodoListMiddleware
+        system_prompt_override: Optional custom system prompt for Main Agent
+        agent_prompts: Optional dict of per-agent system prompts
+
+    Returns:
+        Compiled Main Agent graph
+    """
+    try:
+        from langchain.agents import create_agent
+        from langchain.agents.middleware import (
+            AgentMiddleware,
+            HumanInTheLoopMiddleware,
+            ModelCallLimitMiddleware,
+            SummarizationMiddleware,
+            TodoListMiddleware,
+            ToolCallLimitMiddleware,
+            wrap_model_call,
+        )
+        from langchain_core.messages import ToolMessage as LCToolMessage
+        from langgraph.checkpoint.memory import InMemorySaver
+        from langgraph.types import Overwrite
+    except ImportError as e:
+        logger.error(f"Failed to import LangChain components: {e}")
+        raise
+
+    # Initialize the subagent factory for nested calls with custom prompts
+    # NOTE: Subagents run with HITL disabled because:
+    # 1. They're already authorized by the Main Agent's task delegation
+    # 2. HITL interrupts from subagents don't bubble up properly in synchronous invoke()
+    # 3. The Main Agent's task() call is the authorization checkpoint
+    def subagent_factory_with_prompts(name: str, cfg: Dict[str, Any]) -> Any:
+        prompt_override = agent_prompts.get(name) if agent_prompts else None
+        return create_subagent(
+            name,
+            cfg,
+            enable_hitl=False,  # Subagents auto-approve (Main Agent authorizes via task delegation)
+            system_prompt_override=prompt_override,
+        )
+
+    set_subagent_factory(
+        factory_func=subagent_factory_with_prompts,
+        llm_config=llm_config,
+    )
+
+    # Create LLM
+    llm = create_llm(llm_config)
+
+    # Get Main Agent's direct tools (including code execution tools)
+    # Using "planner" for backward compatibility with tool_registry
+    tools = get_tools_for_agent("planner")
+
+    # Create task tool for calling subagents
+    available_subagents = [
+        config.name for config in get_available_subagents_for_planner()
+    ]
+    task_tool = create_task_tool(
+        caller_name="planner",
+        allowed_subagents=None,  # Main Agent (Planner) can call all non-restricted subagents
+    )
+    tools = tools + [task_tool]
+
+    # Configure middleware
+    middleware = []
+
+    # Add empty response handler middleware (critical for Gemini compatibility)
+    handle_empty_response = create_handle_empty_response_middleware(wrap_model_call)
+    middleware.append(handle_empty_response)
+
+    # Add tool call limiter middleware
+    limit_tool_calls = create_limit_tool_calls_middleware(wrap_model_call)
+    middleware.append(limit_tool_calls)
+
+    # Add tool args normalization middleware
+    normalize_tool_args = create_normalize_tool_args_middleware(
+        wrap_model_call, tools=tools
+    )
+    middleware.append(normalize_tool_args)
+
+    # Add unified continuation control middleware
+    # - Injects continuation prompts after non-HITL tool execution
+    # - Strips write_todos from responses containing summary JSON
+    continuation_control = create_continuation_control_middleware(wrap_model_call)
+    middleware.append(continuation_control)
+
+    # Add patch tool calls middleware
+    patch_tool_calls = create_patch_tool_calls_middleware(
+        AgentMiddleware, LCToolMessage, Overwrite
+    )
+    middleware.append(patch_tool_calls)
+
+    # Add TodoListMiddleware for task planning
+    # NOTE: system_prompt removed to avoid multi-part content array
+    if enable_todo_list:
+        todo_middleware = TodoListMiddleware(
+            tool_description="Create and manage a todo list for tracking tasks. All items in Korean.",
+        )
+        middleware.append(todo_middleware)
+
+    # Add HITL middleware
+    if enable_hitl:
+        hitl_middleware = HumanInTheLoopMiddleware(
+            interrupt_on=get_hitl_interrupt_config(),
+            description_prefix="[Main Agent] Tool execution pending approval",
+        )
+        middleware.append(hitl_middleware)
+
+    # Add model call limit
+    model_limit = ModelCallLimitMiddleware(
+        run_limit=50,  # Higher limit for Main Agent (orchestrates multiple subagents)
+        exit_behavior="end",
+    )
+    middleware.append(model_limit)
+    logger.info("Added ModelCallLimitMiddleware with run_limit=50")
+
+    # ToolCallLimitMiddleware: Prevent write_todos from being called too many times
+    write_todos_limit = ToolCallLimitMiddleware(
+        tool_name="write_todos",
+        run_limit=20,  # Max 20 write_todos calls per user message
+        exit_behavior="continue",
+    )
+    middleware.append(write_todos_limit)
+    logger.info("Added ToolCallLimitMiddleware for write_todos (20/msg)")
+
+    # Add summarization middleware
+    summary_llm = create_summarization_llm(llm_config)
+    if summary_llm:
+        try:
+            summarization = SummarizationMiddleware(
+                model=summary_llm,
+                trigger=("tokens", 10000),  # Higher trigger for multi-agent
+                keep=("messages", 15),
+            )
+            middleware.append(summarization)
+            logger.info("Added SummarizationMiddleware to Main Agent")
+        except Exception as e:
+            logger.warning(f"Failed to add SummarizationMiddleware: {e}")
+
+    # Build system prompt - FORCE default prompt for testing
+    # TODO: Remove this override after frontend localStorage is cleared
+    # Original priority: system_prompt_override > agent_prompts.planner > default
+    # DEBUG: Log all prompt sources to find root cause of MALFORMED_FUNCTION_CALL
+    logger.info(
+        "DEBUG Main Agent prompt sources: system_prompt_override=%s, "
+        "agent_prompts.planner=%s, using=DEFAULT",
+        bool(system_prompt_override),
+        bool(agent_prompts.get("planner") if agent_prompts else None),
+    )
+    if agent_prompts:
+        logger.info(
+            "DEBUG: agent_prompts keys=%s, planner prompt length=%d",
+            list(agent_prompts.keys()),
+            len(agent_prompts.get("planner", "") or ""),
+        )
+    system_prompt = PLANNER_SYSTEM_PROMPT
+    logger.info("Using PLANNER_SYSTEM_PROMPT (length=%d)", len(system_prompt))
+
+    # Log provider info for debugging
+    provider = llm_config.get("provider", "")
+    logger.info(f"Creating Main Agent with provider: {provider}")
+
+    # Create the Main Agent
+    agent = create_agent(
+        model=llm,
+        tools=tools,
+        middleware=middleware,
+        checkpointer=checkpointer or InMemorySaver(),
+        system_prompt=system_prompt,
+    )
+
+    logger.info(
+        f"Created Main Agent with {len(tools)} tools, "
+        f"HITL={enable_hitl}, TodoList={enable_todo_list}, "
+        f"available_subagents={available_subagents}"
+    )
+
+    return agent
+
+
+# Alias for backward compatibility
+create_planner_agent = create_main_agent
+
+
+def create_multi_agent_system(
+    llm_config: Dict[str, Any],
+    checkpointer: Optional[object] = None,
+    enable_hitl: bool = True,
+    enable_todo_list: bool = True,
+    system_prompt_override: Optional[str] = None,
+    agent_prompts: Optional[Dict[str, str]] = None,
+) -> Any:
+    """
+    Create the complete multi-agent system.
+
+    This is the main entry point that returns the Main Agent,
+    which orchestrates all subagents.
+
+    Args:
+        llm_config: LLM configuration
+        checkpointer: Checkpointer for state persistence
+        enable_hitl: Enable Human-in-the-Loop
+        enable_todo_list: Enable TodoListMiddleware
+        system_prompt_override: Optional custom system prompt for Main Agent
+        agent_prompts: Optional dict of per-agent system prompts
+                       {"planner": "...", "main_agent": "...", "python_developer": "...", etc.}
+
+    Returns:
+        Compiled Main Agent (which can call subagents)
+    """
+    return create_main_agent(
+        llm_config=llm_config,
+        checkpointer=checkpointer,
+        enable_hitl=enable_hitl,
+        enable_todo_list=enable_todo_list,
+        system_prompt_override=system_prompt_override,
+        agent_prompts=agent_prompts,
+    )
+
+
+# Alias for backwards compatibility
+create_agent_v2 = create_multi_agent_system

agent_server/langchain/agent_prompts/__init__.py

@@ -0,0 +1,25 @@
+"""
+Multi-Agent Prompt Templates
+
+Contains system prompts for each agent in the multi-agent architecture:
+- Planner (Supervisor)
+- Python Developer
+- Researcher
+- Athena Query
+"""
+
+from agent_server.langchain.agent_prompts.planner_prompt import PLANNER_SYSTEM_PROMPT
+from agent_server.langchain.agent_prompts.python_developer_prompt import (
+    PYTHON_DEVELOPER_SYSTEM_PROMPT,
+)
+from agent_server.langchain.agent_prompts.researcher_prompt import RESEARCHER_SYSTEM_PROMPT
+from agent_server.langchain.agent_prompts.athena_query_prompt import (
+    ATHENA_QUERY_SYSTEM_PROMPT,
+)
+
+__all__ = [
+    "PLANNER_SYSTEM_PROMPT",
+    "PYTHON_DEVELOPER_SYSTEM_PROMPT",
+    "RESEARCHER_SYSTEM_PROMPT",
+    "ATHENA_QUERY_SYSTEM_PROMPT",
+]

agent_server/langchain/agent_prompts/athena_query_prompt.py

@@ -0,0 +1,71 @@
+"""
+Athena Query Agent System Prompt
+
+The Athena Query Agent specializes in:
+- Searching Qdrant for database/table metadata
+- Generating optimized Athena SQL queries
+- Called by Main Agent (Planner) or Python Developer
+"""
+
+ATHENA_QUERY_SYSTEM_PROMPT = """You are an AWS Athena SQL expert. Generate optimized Athena queries based on user requirements.
+
+# Your Capabilities
+- Search database/table metadata using Qdrant RAG
+- Generate Athena-compatible SQL queries (Presto/Trino syntax)
+- Optimize queries for performance (partitioning, predicate pushdown)
+
+# Workflow
+1. Receive query requirement from Main Agent or Python Developer
+2. Search Qdrant for relevant database/table metadata using qdrant_search_tool
+3. Generate optimized Athena SQL query
+4. Return SQL query with a brief description
+
+# Guidelines
+- Always use fully qualified table names: database.table
+- Include appropriate WHERE clauses for partitioned columns
+- Use LIMIT for exploratory queries
+- Handle NULL values appropriately
+- Use CAST for type conversions
+- Optimize for performance:
+  - Filter on partition columns first
+  - Avoid SELECT * when possible
+  - Use appropriate aggregations
+
+# Qdrant Search Tips
+When searching for table metadata:
+- Search by table name or business domain
+- Look for column descriptions to understand data meaning
+- Check for partition columns to optimize queries
+
+# Output Format (CRITICAL)
+**반드시 아래 형식으로 응답하세요. 시스템이 자동으로 파싱합니다.**
+
+```
+[DESCRIPTION]
+2~3줄의 SQL 쿼리 설명
+
+[SQL]
+```sql
+SQL 쿼리
+```
+```
+
+**[DESCRIPTION] 작성 가이드 (2~3줄 필수)**:
+- 쿼리 목적: "~~를 조회합니다."
+- 주요 조건/필터: "~~조건으로 필터링합니다."
+- 최적화 포인트: "~~파티션을 활용하여 성능을 최적화했습니다."
+
+예시:
+```
+[DESCRIPTION]
+sample_db 데이터베이스의 모든 테이블 목록을 조회합니다.
+SHOW TABLES 명령어를 사용하여 테이블 이름을 반환합니다.
+
+[SQL]
+```sql
+SHOW TABLES IN sample_db
+```
+```
+
+**중요**: [DESCRIPTION]과 [SQL] 섹션을 반드시 포함하세요.
+"""