hdsp-jupyter-extension 2.0.27__py3-none-any.whl → 2.0.28__py3-none-any.whl

agent_server/langchain/agent_prompts/planner_prompt.py

@@ -5,15 +5,16 @@ Main Agent (Supervisor) System Prompt for Multi-Agent Mode
 PLANNER_SYSTEM_PROMPT = """당신은 작업을 조율하는 Main Agent입니다. 한국어로 응답하세요.
 
 # 핵심 원칙
-1. 3단계 이상의 복잡한 작업을 요청받은 경우에만 write_todos 로 작업 목록 관리
-2. **직접 코드, 쿼리 작성 금지** - 모든 코드/쿼리 생성은 task_tool로 서브에이전트에게 위임
-3. 서브에이전트가 반환한 코드를 적절한 도구로 실행
-4. 모든 응답 content는 2~3줄 내외로 핵심만 명확하게 전달
+1. **간단한 작업 (1-2단계)**: write_todos 사용 금지 → 바로 실행하고 종료
+2. **복잡한 작업 (3단계+)**: write_todos로 계획 → 순차 실행 → 완료 시 final_summary_tool 호출
+3. **직접 코드, 쿼리 작성 금지** - 모든 코드/쿼리 생성은 task_tool로 서브에이전트에게 위임
+4. 서브에이전트가 반환한 코드를 적절한 도구로 실행
 
 # 작업 흐름
 
 ## Step 1: 계획 수립
-3단계 이상의 복잡한 작업을 요청받은 경우에만 write_todos로 작업 목록 생성 (마지막 항목은 반드시 "작업 요약 및 다음 단계 제시")
+- **간단한 작업 (1-2단계)**: write_todos 없이 바로 실행. 완료 후 추가 도구 호출 없이 종료.
+- **복잡한 작업 (3단계+)**: write_todos로 작업 목록 생성 (실제 작업만 포함, 요약은 시스템이 자동 처리)
 
 ## Step 2: 코드/쿼리 생성 요청
 필요한 경우, task_tool을 호출하여 서브에이전트에게 위임:
@@ -25,18 +26,20 @@ PLANNER_SYSTEM_PROMPT = """당신은 작업을 조율하는 Main Agent입니다.
 | researcher | 정보 검색 | task_tool(agent_name="researcher", description="관련 문서 검색") |
 
 ## Step 3: 결과 실행/적용 (필수!)
-**task_tool을 호출 했다면, 호출 후 반드시 결과를 처리해야 함:**
+**task_tool 호출 후 반드시 결과를 처리해야 함. 코드/SQL은 자동 주입됩니다:**
 
-| 서브에이전트 | 작업 유형 | 처리 방법 | 예시 |
-|-------------|----------|----------|------|
-| python_developer | 코드 실행 (데이터 분석, 시각화) | jupyter_cell_tool | jupyter_cell_tool(code=반환된_코드) |
-| python_developer | **파일 생성/수정** | **write_file_tool 또는 multiedit_file_tool** | write_file_tool(path="script.js", content=반환된_코드) |
-| athena_query | SQL 표시 | markdown_tool | markdown_tool(content="```sql\n반환된_쿼리\n```") |
+| 서브에이전트 | 작업 유형 | 처리 방법 | 호출 방법 |
+|-------------|----------|----------|----------|
+| python_developer | 코드 실행 (데이터 분석, 시각화) | jupyter_cell_tool | jupyter_cell_tool() ← 코드 자동 주입, code 파라미터 불필요 |
+| python_developer | **파일 생성/수정** | **write_file_tool** | write_file_tool(path="파일경로") ← content 자동 주입 |
+| athena_query | SQL 표시 | markdown_tool | markdown_tool() ← SQL 자동 주입, content 파라미터 불필요 |
 | researcher | 텍스트 요약 | 직접 응답 | - |
 
-**🔴 중요: 코드 저장 도구 선택**
-- **파일 생성/수정 요청** → `write_file_tool` 또는 `multiedit_file_tool` 사용
-- **코드 실행 요청** (데이터 분석, 차트 등) → `jupyter_cell_tool` 사용
+**🔴 중요: 코드/SQL 자동 주입**
+- task_tool이 생성한 코드/SQL은 **State를 통해 자동 주입**됩니다
+- **코드를 직접 복사하거나 인자로 전달할 필요 없음** — 도구만 호출하면 됨
+- **파일 생성/수정 요청** → `write_file_tool(path=...)` 사용 (content 자동 주입)
+- **코드 실행 요청** (데이터 분석, 차트 등) → `jupyter_cell_tool()` 사용 (code 자동 주입)
 - **❌ markdown_tool은 코드 저장용이 아님!** (표시 전용)
 
 **중요**: task_tool 결과를 받은 후 바로 write_todos로 완료 처리하지 말고, 반드시 위 도구로 결과를 먼저 적용!
@@ -49,34 +52,22 @@ PLANNER_SYSTEM_PROMPT = """당신은 작업을 조율하는 Main Agent입니다.
     - **🔴 기존 todo 절대 삭제 금지**: 전체 리스트를 항상 포함하고 status만 변경
     - **🔴 상태 전환 순서 필수**: pending → in_progress → completed (건너뛰기 금지!)
     - **🔴 초기 생성 규칙**: 첫 write_todos 호출 시 첫 번째 todo만 in_progress, 나머지는 모두 pending
-      - 올바른 초기 예: [{"content": "작업1", "status": "in_progress"}, {"content": "작업2", "status": "pending"}, {"content": "작업 요약 및 다음 단계 제시", "status": "pending"}]
+      - 올바른 초기 예: [{"content": "작업1", "status": "in_progress"}, {"content": "작업2", "status": "pending"}]
       - 잘못된 초기 예: [{"content": "작업1", "status": "completed"}, ...] ← 실제 작업 없이 completed 금지!
     - **🔴 completed 전환 조건**: 실제 도구(task_tool, jupyter_cell_tool 등)로 작업 수행 후에만 completed로 변경
     - in_progress 상태는 **동시에 1개만** 허용 (completed, pending todo는 삭제하지 않고 모두 유지)
     - content에 도구(tool)명 언급 금지
-    - **[필수] 마지막 todo는 반드시 "작업 요약 및 다음 단계 제시"**
-
-# "작업 요약 및 다음 단계 제시" todo 완료 시 [필수]
-1. "작업 요약 및 다음 단계 제시"를 **in_progress**로 변경 (write_todos 호출)
-2. **반드시 final_summary_tool 호출**:
-   final_summary_tool(
-     summary="완료된 작업 요약",
-     next_items=[{"subject": "제목", "description": "설명"}, ...]
-   )
-   
-3. final_summary_tool 호출 후 "작업 요약 및 다음 단계 제시"를 **completed**로 변경
-
-- next_items 3개 이상 필수
-- **final_summary_tool 호출 없이 종료 금지**
+    - **"작업 요약" todo 추가 금지**: 실제 작업만 todo로 생성 (요약은 시스템이 자동 처리)
 
 # 도구 사용시 주의할 점
 ## 파일 위치 모를 때 탐색 순서: search_files_tool → list_workspace_tool → 재검색 → ask_user_tool 순서로!)
-## list_workspace_tool로 전체 디렉토리 파일 목록 검색 금지! 최대한 pattern 으로 drill down 해서 검색할 것 
+## list_workspace_tool로 전체 디렉토리 파일 목록 검색 금지! 최대한 pattern 으로 drill down 해서 검색할 것
 
 # 금지 사항
 - 직접 코드/SQL 작성 (반드시 task_tool 사용)
 - task_tool 없이 jupyter_cell_tool 호출
-- **task_tool 결과를 표시하지 않고 바로 완료 처리** (athena_query → markdown_tool 필수)
+- **task_tool 결과를 처리하지 않고 바로 완료** (python_developer → jupyter_cell_tool, athena_query → markdown_tool 필수)
+- jupyter_cell_tool에 code 인자를 직접 전달 (자동 주입되므로 불필요)
 - 빈 응답
 """
 

agent_server/langchain/custom_middleware.py

@@ -12,7 +12,9 @@ import uuid
 from typing import Any, Dict, Optional
 
 from json_repair import repair_json
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain.agents.middleware import AgentMiddleware
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langgraph.types import Command
 
 from agent_server.langchain.logging_utils import (
     _format_middleware_marker,
@@ -25,6 +27,92 @@ from agent_server.langchain.prompts import JSON_TOOL_SCHEMA, NON_HITL_TOOLS
 logger = logging.getLogger(__name__)
 
 
+# ---------------------------------------------------------------------------
+# TodoActiveMiddleware — manages todo_active state field
+# ---------------------------------------------------------------------------
+
+
+class TodoActiveMiddleware(AgentMiddleware):
+    """Middleware that manages the `todo_active` state field.
+
+    Intercepts write_todos and final_summary_tool calls to set/clear
+    the todo_active flag in LangGraph state via Command.
+
+    - write_todos called → todo_active = True
+    - final_summary_tool called → todo_active = False
+
+    This flag is checked by handle_empty_response and continuation_control
+    middlewares to decide whether to force continuation or let the LLM
+    terminate naturally (for simple 1-2 step tasks).
+    """
+
+    def wrap_tool_call(self, request, handler):
+        """Intercept tool calls to manage todo_active state."""
+        result = handler(request)
+        tool_name = request.tool_call.get("name", "")
+
+        if tool_name == "write_todos":
+            return self._wrap_with_todo_active(request, result, active=True)
+        elif tool_name in ("final_summary_tool", "final_summary"):
+            return self._wrap_with_todo_active(request, result, active=False)
+
+        return result
+
+    def _wrap_with_todo_active(self, request, result, active: bool):
+        """Wrap tool result in a Command that updates todo_active state.
+
+        Handles two cases:
+        1. Result is already a Command (e.g., from TodoListMiddleware) → merge
+        2. Result is a ToolMessage → wrap in new Command
+        """
+        try:
+            if isinstance(result, Command):
+                # Merge todo_active into existing Command's update dict
+                existing_update = (
+                    result.update if hasattr(result, "update") and result.update else {}
+                )
+                merged_update = {**existing_update, "todo_active": active}
+                logger.info(
+                    "[TodoActive] Merged todo_active=%s into Command for tool '%s'",
+                    active,
+                    request.tool_call.get("name", ""),
+                )
+                return Command(update=merged_update)
+            elif isinstance(result, ToolMessage):
+                # Wrap ToolMessage in a new Command
+                logger.info(
+                    "[TodoActive] Wrapped ToolMessage in Command with todo_active=%s for tool '%s'",
+                    active,
+                    request.tool_call.get("name", ""),
+                )
+                return Command(
+                    update={
+                        "todo_active": active,
+                        "messages": [result],
+                    }
+                )
+            else:
+                # Unknown result type — wrap as ToolMessage
+                tool_call_id = request.tool_call.get("id", "")
+                content = str(result) if result else ""
+                logger.info(
+                    "[TodoActive] Wrapped unknown result type (%s) in Command with todo_active=%s",
+                    type(result).__name__,
+                    active,
+                )
+                return Command(
+                    update={
+                        "todo_active": active,
+                        "messages": [
+                            ToolMessage(content=content, tool_call_id=tool_call_id)
+                        ],
+                    }
+                )
+        except Exception as e:
+            logger.warning("[TodoActive] Failed to set todo_active=%s: %s", active, e)
+            return result
+
+
 def parse_json_tool_call(text) -> Optional[Dict[str, Any]]:
     """Parse JSON tool call from text response.
 
@@ -262,6 +350,31 @@ def create_handle_empty_response_middleware(wrap_model_call):
     def handle_empty_response(request, handler):
         max_retries = 2
 
+        # Guard: If final_summary_tool was already called, stop the agent immediately.
+        # This is independent of todo status (LLM may call final_summary before
+        # marking all todos as completed).
+        todo_active = request.state.get("todo_active", False)
+        if not todo_active:
+            messages = request.messages
+            # Find last REAL HumanMessage index
+            _last_human = -1
+            for _i, _msg in enumerate(messages):
+                _mtype = getattr(_msg, "type", "") or type(_msg).__name__
+                if _mtype in ("human", "HumanMessage"):
+                    _mcontent = getattr(_msg, "content", "") or ""
+                    if not _mcontent.startswith("[SYSTEM]"):
+                        _last_human = _i
+            _msgs_after = (
+                messages[_last_human + 1 :] if _last_human >= 0 else messages[-10:]
+            )
+            for _msg in _msgs_after:
+                _name = getattr(_msg, "name", "") or ""
+                if _name in ("final_summary_tool", "final_summary"):
+                    logger.info(
+                        "final_summary_tool already executed and todo_active=False - stopping agent (no LLM call)"
+                    )
+                    return AIMessage(content="", tool_calls=[])
+
         # Check if all todos are completed - if so, return empty response to stop agent
         # Method 1: Check state.todos
         todos = request.state.get("todos", [])
@@ -297,8 +410,15 @@ def create_handle_empty_response_middleware(wrap_model_call):
                     else messages[-10:]
                 )
                 for msg in messages_to_check:
+                    # Check ToolMessage name for final_summary_tool
+                    msg_name = getattr(msg, "name", "") or ""
+                    if msg_name in ("final_summary_tool", "final_summary"):
+                        summary_exists = True
+                        break
                     content = getattr(msg, "content", "") or ""
-                    if '"summary"' in content and '"next_items"' in content:
+                    if ('"summary"' in content and '"next_items"' in content) or (
+                        "'summary'" in content and "'next_items'" in content
+                    ):
                         summary_exists = True
                         break
 
@@ -343,8 +463,15 @@ def create_handle_empty_response_middleware(wrap_model_call):
         messages = request.messages
         summary_exists = False
         for msg in messages[-15:]:
+            # Check ToolMessage name for final_summary_tool
+            msg_name = getattr(msg, "name", "") or ""
+            if msg_name in ("final_summary_tool", "final_summary"):
+                summary_exists = True
+                break
             msg_content = getattr(msg, "content", "") or ""
-            if '"summary"' in msg_content and '"next_items"' in msg_content:
+            if ('"summary"' in msg_content and '"next_items"' in msg_content) or (
+                "'summary'" in msg_content and "'next_items'" in msg_content
+            ):
                 summary_exists = True
                 break
             if any(
@@ -583,6 +710,14 @@ def create_handle_empty_response_middleware(wrap_model_call):
 
             # Invalid response - retry with JSON schema prompt
             if response_message and attempt < max_retries:
+                # todo_active=False → LLM can terminate naturally (simple tasks)
+                todo_active = request.state.get("todo_active", False)
+                if not todo_active:
+                    logger.info(
+                        "todo_active=False - skipping retry, allowing LLM natural termination"
+                    )
+                    return response
+
                 reason = "text-only" if has_content else "empty"
 
                 json_prompt = _build_json_prompt(request, response_message, has_content)
@@ -776,23 +911,38 @@ def _build_json_prompt(request, response_message, has_content):
             f"Example: {example_json}"
         )
     elif not todos:
-        # No todos yet = new task starting, LLM must create todos or call a tool
-        # This happens when LLM returns empty response at the start of a new task
-        logger.info("No todos exist yet - forcing retry to create todos or call tool")
-        return (
-            f"{JSON_TOOL_SCHEMA}\n\n"
-            f"Your response was empty. You MUST call a tool to proceed.\n"
-            f"한국어로 응답하고, write_todos로 작업 목록을 만들거나 jupyter_cell_tool/read_file_tool을 호출하세요.\n"
-            f'Example: {{"tool": "write_todos", "arguments": {{"todos": [{{"content": "데이터 분석", "status": "in_progress"}}]}}}}'
+        # No todos → simple task (1-2 steps), don't force write_todos creation
+        # This was the DIRECT CAUSE of the simple-task infinite loop:
+        # LLM completes simple task → empty response → forced to create todos → loop
+        logger.info(
+            "No todos exist - simple task, skipping retry (no write_todos forcing)"
         )
+        return None  # Signal to skip retry — LLM terminates naturally
     else:
-        # Todos exist but all completed - ask for summary
-        logger.info("All todos completed but response empty - asking for summary")
+        # Todos exist but all completed
+        # Check if final_summary_tool was already called in message history
+        messages = getattr(request, "messages", [])
+        final_summary_already_called = any(
+            getattr(msg, "name", "") in ("final_summary_tool", "final_summary")
+            for msg in messages
+        )
+        if final_summary_already_called:
+            logger.info(
+                "All todos completed and final_summary_tool already called - "
+                "signaling skip (no more retries needed)"
+            )
+            return None  # Signal to skip retry and synthesize completion
+
+        logger.info(
+            "All todos completed but response empty - asking for final_summary_tool"
+        )
         return (
             f"{JSON_TOOL_SCHEMA}\n\n"
-            f"All tasks completed. Call markdown_tool to provide a summary in Korean.\n"
-            f"한국어로 작업 요약을 작성하세요.\n"
-            f'Example: {{"tool": "markdown_tool", "arguments": {{"content": "작업이 완료되었습니다."}}}}'
+            f"All tasks completed. Call final_summary_tool to provide a summary.\n"
+            f"final_summary_tool(summary='완료된 작업 요약', "
+            f"next_items=[{{'subject': '제목', 'description': '설명'}}, ...]) "
+            f"(next_items 3개 이상 필수).\n"
+            f"텍스트로 JSON을 출력하지 말고, 반드시 도구 호출로 실행하세요."
         )
 
 
@@ -1020,8 +1170,31 @@ def create_normalize_tool_args_middleware(wrap_model_call, tools=None):
                                 tool_call["args"], dict
                             ):
                                 args = tool_call["args"]
-                                # Normalize list arguments to strings for str-typed params
+                                # Normalize non-string arguments for str-typed params
                                 for key, value in args.items():
+                                    # Convert dict to string/None for str-typed params
+                                    # LLM sometimes sends {} instead of null for Optional[str]
+                                    if key in string_params and isinstance(value, dict):
+                                        if not value:  # Empty dict {}
+                                            logger.info(
+                                                "Converted empty dict to None for '%s' in tool '%s'",
+                                                key,
+                                                tool_name,
+                                            )
+                                            args[key] = None
+                                        else:
+                                            # Non-empty dict → JSON string
+                                            json_str = json.dumps(
+                                                value, ensure_ascii=False
+                                            )
+                                            logger.info(
+                                                "Converted dict to JSON string for '%s' in tool '%s': %s",
+                                                key,
+                                                tool_name,
+                                                json_str[:100],
+                                            )
+                                            args[key] = json_str
+
                                     if key in string_params and isinstance(value, list):
                                         # Join list items into a single string
                                         text_parts = []
@@ -1150,10 +1323,18 @@ def create_continuation_control_middleware(wrap_model_call):
             else messages[-15:]
         )
         for msg in messages_to_check:
+            # Check if this is a ToolMessage from final_summary_tool
+            msg_name = getattr(msg, "name", "") or ""
+            if msg_name in ("final_summary_tool", "final_summary"):
+                return True
+
             msg_content = getattr(msg, "content", "") or ""
-            # Check for summary JSON
+            # Check for summary JSON (double quotes)
             if '"summary"' in msg_content and '"next_items"' in msg_content:
                 return True
+            # Check for summary Python str (single quotes from tool output)
+            if "'summary'" in msg_content and "'next_items'" in msg_content:
+                return True
             # Check for markdown summary (common patterns)
             if any(
                 kw in msg_content
@@ -1203,6 +1384,24 @@ def create_continuation_control_middleware(wrap_model_call):
                         pass
 
                 if tool_name in NON_HITL_TOOLS:
+                    # GUARD: Skip forcing when final_summary_tool already ran
+                    if tool_name in ("final_summary_tool", "final_summary"):
+                        logger.info(
+                            "final_summary_tool already executed - "
+                            "skipping continuation (preventing infinite loop)"
+                        )
+                        return handler(request)
+
+                    # GUARD: todo_active=False → simple task, skip continuation
+                    todo_active = request.state.get("todo_active", False)
+                    if not todo_active:
+                        logger.info(
+                            "todo_active=False after tool '%s' - "
+                            "simple task, skipping continuation",
+                            tool_name,
+                        )
+                        return handler(request)
+
                     todos = request.state.get("todos", [])
 
                     last_real_human_idx = _find_last_real_human_idx(messages)
@@ -1237,36 +1436,60 @@ def create_continuation_control_middleware(wrap_model_call):
                         tool_name,
                     )
 
-                    # Skip continuation injection for write_todos
-                    # This prevents auto-continuation to next task after completing one
-                    # Agent will decide next action based on its own reasoning
-                    if tool_name == "write_todos":
+                    # === State-based branching: todos 유무로 분기 ===
+                    #
+                    # (1) todos 없음 → 간단한 1~2단계 작업 → continuation 불필요
+                    # (2) todos 있음 + 미완료 → 다음 작업 유도
+                    # (3) todos 있음 + 전부 완료 → final_summary_tool 호출 유도
+                    #
+                    if not todos:
+                        # No todos in state → simple task (1~2 steps)
+                        # Don't inject any continuation — LLM finishes naturally.
                         logger.info(
-                            "Skipping continuation prompt after write_todos - "
-                            "agent decides next action (pending: %d)",
-                            len(pending_todos) if pending_todos else 0,
+                            "No todos in state after tool: %s - "
+                            "simple task, skipping continuation",
+                            tool_name,
                         )
-                        # Don't inject continuation - let agent naturally continue or stop
                     elif pending_todos:
-                        pending_list = ", ".join(
-                            t.get("content", "")[:30] for t in pending_todos[:3]
-                        )
-                        continuation = (
-                            f"Tool '{tool_name}' completed. "
-                            f"Continue with pending tasks: {pending_list}. "
-                            f"Call jupyter_cell_tool or the next appropriate tool."
-                        )
-                        new_messages = list(messages) + [
-                            HumanMessage(content=f"[SYSTEM] {continuation}")
-                        ]
-                        request = request.override(messages=new_messages)
+                        # Todos exist with pending items → guide to next task
+                        if tool_name == "write_todos":
+                            # write_todos with pending items → agent manages its own flow
+                            logger.info(
+                                "write_todos with %d pending todos - "
+                                "agent manages own flow",
+                                len(pending_todos),
+                            )
+                        else:
+                            pending_list = ", ".join(
+                                t.get("content", "")[:30] for t in pending_todos[:3]
+                            )
+                            continuation = (
+                                f"Tool '{tool_name}' completed. "
+                                f"Continue with pending tasks: {pending_list}. "
+                                f"Call jupyter_cell_tool or the next appropriate tool."
+                            )
+                            new_messages = list(messages) + [
+                                HumanMessage(content=f"[SYSTEM] {continuation}")
+                            ]
+                            request = request.override(messages=new_messages)
                     else:
+                        # All todos completed → prompt for final_summary_tool
+                        logger.info(
+                            "All %d todos completed after tool: %s - "
+                            "prompting for final_summary_tool",
+                            len(todos),
+                            tool_name,
+                        )
                         continuation = (
-                            f"Tool '{tool_name}' completed. "
-                            f"Create a todo list with write_todos if needed."
+                            "[SYSTEM] 모든 작업이 완료되었습니다. "
+                            "반드시 final_summary_tool을 호출하여 작업 요약과 다음 단계를 제시하세요. "
+                            "final_summary_tool(summary='완료된 작업 요약', "
+                            "next_items=[{'subject': '제목', 'description': '설명'}, ...]) "
+                            "(next_items 3개 이상 필수). "
+                            "텍스트로 JSON을 출력하지 말고, 반드시 도구 호출로 실행하세요."
                         )
                         new_messages = list(messages) + [
-                            HumanMessage(content=f"[SYSTEM] {continuation}")
+                            HumanMessage(content=continuation)
                         ]
                         request = request.override(messages=new_messages)
 
@@ -1287,8 +1510,10 @@ def create_continuation_control_middleware(wrap_model_call):
                 if isinstance(p, (str, dict))
             )
 
-        # Check if content contains summary JSON pattern
-        has_summary_json = '"summary"' in content and '"next_items"' in content
+        # Check if content contains summary JSON pattern (double or single quotes)
+        has_summary_json = ('"summary"' in content and '"next_items"' in content) or (
+            "'summary'" in content and "'next_items'" in content
+        )
 
         if has_summary_json:
             tool_calls = getattr(response_message, "tool_calls", []) or []