codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/tools/jupyter_tools.py

@@ -0,0 +1,296 @@
+"""Jupyter 노트북 도구."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from .base import BaseTool, ToolResult
+
+
+class JupyterCellTool(BaseTool):
+    """Jupyter 셀 작업 도구."""
+
+    def __init__(self, kernel_manager: Any, contents_manager: Any):
+        self._kernel_manager = kernel_manager
+        self._contents_manager = contents_manager
+        self._notebook_path: Optional[str] = None
+        self._notebook_model: Optional[dict] = None
+
+    @property
+    def name(self) -> str:
+        return "jupyter_cell"
+
+    @property
+    def description(self) -> str:
+        return """Create, read, update, delete, or execute Jupyter notebook cells.
+Operations:
+- CREATE: Create a new code cell with the given code
+- READ: Read the content of a cell at the given index
+- UPDATE: Update the code of an existing cell
+- DELETE: Delete a cell at the given index
+- EXECUTE: Execute a cell and return the output"""
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "operation": {
+                    "type": "string",
+                    "enum": ["CREATE", "READ", "UPDATE", "DELETE", "EXECUTE"],
+                    "description": "The operation to perform",
+                },
+                "code": {
+                    "type": "string",
+                    "description": "The code content (for CREATE, UPDATE, EXECUTE)",
+                },
+                "cell_index": {
+                    "type": "integer",
+                    "description": "The cell index (for READ, UPDATE, DELETE, EXECUTE)",
+                },
+                "insert_after": {
+                    "type": "integer",
+                    "description": "Insert new cell after this index (for CREATE)",
+                },
+            },
+            "required": ["operation"],
+        }
+
+    @property
+    def risk_level(self) -> str:
+        return "medium"
+
+    def set_notebook(self, path: str, model: dict) -> None:
+        """현재 노트북 컨텍스트를 설정합니다."""
+        self._notebook_path = path
+        self._notebook_model = model
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        """Jupyter 셀 작업을 실행합니다."""
+        operation = kwargs.get("operation", "").upper()
+
+        if operation == "CREATE":
+            return await self._create_cell(
+                kwargs.get("code", ""), kwargs.get("insert_after")
+            )
+        elif operation == "READ":
+            return await self._read_cell(kwargs.get("cell_index", 0))
+        elif operation == "UPDATE":
+            return await self._update_cell(
+                kwargs.get("cell_index", 0), kwargs.get("code", "")
+            )
+        elif operation == "DELETE":
+            return await self._delete_cell(kwargs.get("cell_index", 0))
+        elif operation == "EXECUTE":
+            return await self._execute_cell(
+                kwargs.get("cell_index"), kwargs.get("code")
+            )
+        else:
+            return ToolResult(success=False, error=f"알 수 없는 연산: {operation}")
+
+    async def _create_cell(
+        self, code: str, insert_after: Optional[int] = None
+    ) -> ToolResult:
+        """새 셀을 생성합니다."""
+        if not self._notebook_model:
+            return ToolResult(success=False, error="노트북 컨텍스트가 설정되지 않음")
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+        new_cell = {
+            "cell_type": "code",
+            "source": code,
+            "metadata": {},
+            "outputs": [],
+            "execution_count": None,
+        }
+
+        if insert_after is not None and 0 <= insert_after < len(cells):
+            cells.insert(insert_after + 1, new_cell)
+            cell_index = insert_after + 1
+        else:
+            cells.append(new_cell)
+            cell_index = len(cells) - 1
+
+        return ToolResult(
+            success=True,
+            data={"cell_index": cell_index, "operation": "created"},
+            metadata={"cell_type": "code"},
+        )
+
+    async def _read_cell(self, cell_index: int) -> ToolResult:
+        """셀 내용을 읽습니다."""
+        if not self._notebook_model:
+            return ToolResult(success=False, error="노트북 컨텍스트가 설정되지 않음")
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+        if not 0 <= cell_index < len(cells):
+            return ToolResult(success=False, error=f"잘못된 셀 인덱스: {cell_index}")
+
+        cell = cells[cell_index]
+        return ToolResult(
+            success=True,
+            data={
+                "source": cell.get("source", ""),
+                "cell_type": cell.get("cell_type", "code"),
+                "outputs": cell.get("outputs", []),
+            },
+        )
+
+    async def _update_cell(self, cell_index: int, code: str) -> ToolResult:
+        """셀 내용을 업데이트합니다."""
+        if not self._notebook_model:
+            return ToolResult(success=False, error="노트북 컨텍스트가 설정되지 않음")
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+        if not 0 <= cell_index < len(cells):
+            return ToolResult(success=False, error=f"잘못된 셀 인덱스: {cell_index}")
+
+        cells[cell_index]["source"] = code
+        return ToolResult(
+            success=True,
+            data={"cell_index": cell_index, "operation": "updated"},
+        )
+
+    async def _delete_cell(self, cell_index: int) -> ToolResult:
+        """셀을 삭제합니다."""
+        if not self._notebook_model:
+            return ToolResult(success=False, error="노트북 컨텍스트가 설정되지 않음")
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+        if not 0 <= cell_index < len(cells):
+            return ToolResult(success=False, error=f"잘못된 셀 인덱스: {cell_index}")
+
+        del cells[cell_index]
+        return ToolResult(
+            success=True,
+            data={"cell_index": cell_index, "operation": "deleted"},
+        )
+
+    async def _execute_cell(
+        self, cell_index: Optional[int], code: Optional[str]
+    ) -> ToolResult:
+        """셀 또는 코드를 실행합니다."""
+        # 코드가 직접 제공되면 실행
+        if code:
+            return await self._execute_code(code)
+
+        # 그렇지 않으면 주어진 인덱스의 셀 실행
+        if cell_index is not None and self._notebook_model:
+            cells = self._notebook_model.get("content", {}).get("cells", [])
+            if 0 <= cell_index < len(cells):
+                cell_code = cells[cell_index].get("source", "")
+                return await self._execute_code(cell_code)
+
+        return ToolResult(
+            success=False, error="코드 또는 유효한 셀 인덱스가 제공되지 않음"
+        )
+
+    async def _execute_code(self, code: str) -> ToolResult:
+        """커널에서 코드를 실행합니다."""
+        # 실제 커널과 통합됨
+        # 현재는 플레이스홀더 반환
+        return ToolResult(
+            success=True,
+            data={
+                "output": "코드 실행 플레이스홀더",
+                "execution_count": 1,
+            },
+            metadata={"executed": True},
+        )
+
+
+class JupyterMarkdownTool(BaseTool):
+    """Jupyter 마크다운 셀 작업 도구."""
+
+    def __init__(self, contents_manager: Any):
+        self._contents_manager = contents_manager
+        self._notebook_model: Optional[dict] = None
+
+    @property
+    def name(self) -> str:
+        return "jupyter_markdown"
+
+    @property
+    def description(self) -> str:
+        return """Create or update markdown cells in the Jupyter notebook.
+Use this for adding documentation, explanations, or headers."""
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "operation": {
+                    "type": "string",
+                    "enum": ["CREATE", "UPDATE"],
+                    "description": "The operation to perform",
+                },
+                "content": {
+                    "type": "string",
+                    "description": "The markdown content",
+                },
+                "cell_index": {
+                    "type": "integer",
+                    "description": "The cell index (for UPDATE)",
+                },
+                "insert_after": {
+                    "type": "integer",
+                    "description": "Insert new cell after this index (for CREATE)",
+                },
+            },
+            "required": ["operation", "content"],
+        }
+
+    @property
+    def risk_level(self) -> str:
+        return "low"
+
+    def set_notebook(self, model: dict) -> None:
+        """현재 노트북 컨텍스트를 설정합니다."""
+        self._notebook_model = model
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        """마크다운 셀 작업을 실행합니다."""
+        operation = kwargs.get("operation", "").upper()
+        content = kwargs.get("content", "")
+
+        if not self._notebook_model:
+            return ToolResult(success=False, error="노트북 컨텍스트가 설정되지 않음")
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+
+        if operation == "CREATE":
+            new_cell = {
+                "cell_type": "markdown",
+                "source": content,
+                "metadata": {},
+            }
+
+            insert_after = kwargs.get("insert_after")
+            if insert_after is not None and 0 <= insert_after < len(cells):
+                cells.insert(insert_after + 1, new_cell)
+                cell_index = insert_after + 1
+            else:
+                cells.append(new_cell)
+                cell_index = len(cells) - 1
+
+            return ToolResult(
+                success=True,
+                data={"cell_index": cell_index, "operation": "created"},
+            )
+
+        elif operation == "UPDATE":
+            cell_index = kwargs.get("cell_index", 0)
+            if not 0 <= cell_index < len(cells):
+                return ToolResult(
+                    success=False, error=f"잘못된 셀 인덱스: {cell_index}"
+                )
+
+            cells[cell_index]["source"] = content
+            cells[cell_index]["cell_type"] = "markdown"
+            return ToolResult(
+                success=True,
+                data={"cell_index": cell_index, "operation": "updated"},
+            )
+
+        return ToolResult(success=False, error=f"알 수 없는 연산: {operation}")

jupyter_ext/tools/registry.py

@@ -0,0 +1,159 @@
+"""도구 관리 및 실행을 위한 도구 레지스트리."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from codebuff_core.core.exceptions import ToolError
+
+from .base import BaseTool, ToolResult
+from .jupyter_tools import JupyterCellTool, JupyterMarkdownTool
+from .file_tools import FileReadTool, FileWriteTool, FileListTool
+from .shell_tools import ShellExecuteTool
+from .search_tools import SearchCodeTool, WebSearchTool
+
+
+class ToolRegistry:
+    """도구 관리를 위한 레지스트리."""
+
+    # HITL 승인이 필요한 도구들
+    HIGH_RISK_TOOLS = {"shell_execute", "file_write"}
+
+    def __init__(
+        self,
+        kernel_manager: Any = None,
+        contents_manager: Any = None,
+        base_path: str = ".",
+        tavily_api_key: Optional[str] = None,
+    ):
+        self._tools: dict[str, BaseTool] = {}
+        self._base_path = base_path
+        self._kernel_manager = kernel_manager
+        self._contents_manager = contents_manager
+
+        self._register_default_tools(tavily_api_key)
+
+    def _register_default_tools(self, tavily_api_key: Optional[str] = None) -> None:
+        """모든 기본 도구를 등록합니다."""
+        # Jupyter 도구
+        self._tools["jupyter_cell"] = JupyterCellTool(
+            self._kernel_manager,
+            self._contents_manager,
+        )
+        self._tools["jupyter_markdown"] = JupyterMarkdownTool(
+            self._contents_manager,
+        )
+
+        # 파일 도구
+        self._tools["file_read"] = FileReadTool(
+            self._contents_manager,
+            self._base_path,
+        )
+        self._tools["file_write"] = FileWriteTool(
+            self._contents_manager,
+            self._base_path,
+        )
+        self._tools["file_list"] = FileListTool(
+            self._contents_manager,
+            self._base_path,
+        )
+
+        # 셸 도구
+        self._tools["shell_execute"] = ShellExecuteTool(
+            working_dir=self._base_path,
+        )
+
+        # 검색 도구
+        self._tools["search_code"] = SearchCodeTool(
+            self._base_path,
+        )
+        web_search = WebSearchTool(api_key=tavily_api_key)
+        self._tools["web_search"] = web_search
+
+    def register(self, tool: BaseTool) -> None:
+        """도구를 등록합니다."""
+        self._tools[tool.name] = tool
+
+    def unregister(self, name: str) -> None:
+        """도구 등록을 해제합니다."""
+        if name in self._tools:
+            del self._tools[name]
+
+    def get(self, name: str) -> Optional[BaseTool]:
+        """이름으로 도구를 가져옵니다."""
+        return self._tools.get(name)
+
+    def get_definitions(self) -> list[dict[str, Any]]:
+        """LLM용 모든 도구 정의를 가져옵니다."""
+        return [tool.definition for tool in self._tools.values()]
+
+    def requires_approval(self, name: str, args: dict[str, Any]) -> bool:
+        """
+        도구 실행에 HITL 승인이 필요한지 확인합니다.
+
+        Args:
+            name: 도구 이름
+            args: 도구 인자
+
+        Returns:
+            승인이 필요하면 True
+        """
+        # 고위험 도구인지 확인
+        if name in self.HIGH_RISK_TOOLS:
+            return True
+
+        # 도구의 위험 수준 확인
+        tool = self._tools.get(name)
+        if tool and tool.risk_level in ("high", "critical"):
+            return True
+
+        # 인자 기반 추가 검사
+        if name == "jupyter_cell":
+            # EXECUTE 작업에 승인 필요
+            if args.get("operation", "").upper() == "EXECUTE":
+                return True
+
+        return False
+
+    async def execute(self, name: str, args: dict[str, Any]) -> ToolResult:
+        """
+        도구를 실행합니다.
+
+        Args:
+            name: 도구 이름
+            args: 도구 인자
+
+        Returns:
+            도구 실행 결과 ToolResult
+        """
+        tool = self._tools.get(name)
+        if not tool:
+            return ToolResult(success=False, error=f"알 수 없는 도구: {name}")
+
+        try:
+            return await tool.execute(**args)
+        except ToolError:
+            raise
+        except Exception as e:
+            raise ToolError(f"도구 실행 오류: {e}", tool_name=name, cause=e) from e
+
+    def set_notebook_context(self, path: str, model: dict) -> None:
+        """Jupyter 도구에 노트북 컨텍스트를 설정합니다."""
+        jupyter_cell = self._tools.get("jupyter_cell")
+        if isinstance(jupyter_cell, JupyterCellTool):
+            jupyter_cell.set_notebook(path, model)
+
+        jupyter_md = self._tools.get("jupyter_markdown")
+        if isinstance(jupyter_md, JupyterMarkdownTool):
+            jupyter_md.set_notebook(model)
+
+    def set_tavily_api_key(self, api_key: str) -> None:
+        """웹 검색을 위한 Tavily API 키를 설정합니다."""
+        web_search = self._tools.get("web_search")
+        if isinstance(web_search, WebSearchTool):
+            web_search.set_api_key(api_key)
+
+    @property
+    def tool_names(self) -> list[str]:
+        """등록된 도구 이름 목록을 가져옵니다."""
+        return list(self._tools.keys())

jupyter_ext/tools/search_tools.py

@@ -0,0 +1,230 @@
+"""코드 및 웹 검색 도구."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Any, Optional
+
+import httpx
+
+from .base import BaseTool, ToolResult
+
+
+class SearchCodeTool(BaseTool):
+    """파일에서 코드 검색 도구."""
+
+    def __init__(self, base_path: str = "."):
+        self._base_path = Path(base_path)
+
+    @property
+    def name(self) -> str:
+        return "search_code"
+
+    @property
+    def description(self) -> str:
+        return """Search for patterns in code files.
+Uses regex patterns to find matches across files.
+Returns matching lines with file paths and line numbers."""
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "pattern": {
+                    "type": "string",
+                    "description": "Regex pattern to search for",
+                },
+                "path": {
+                    "type": "string",
+                    "description": "Directory to search in (relative to workspace)",
+                    "default": ".",
+                },
+                "file_pattern": {
+                    "type": "string",
+                    "description": "Glob pattern for files to search (e.g., '*.py')",
+                    "default": "*",
+                },
+                "max_results": {
+                    "type": "integer",
+                    "description": "Maximum number of results to return",
+                    "default": 50,
+                },
+                "case_sensitive": {
+                    "type": "boolean",
+                    "description": "Case-sensitive search",
+                    "default": True,
+                },
+            },
+            "required": ["pattern"],
+        }
+
+    @property
+    def risk_level(self) -> str:
+        return "low"
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        """코드 패턴을 검색합니다."""
+        pattern = kwargs.get("pattern", "")
+        path = kwargs.get("path", ".")
+        file_pattern = kwargs.get("file_pattern", "*")
+        max_results = kwargs.get("max_results", 50)
+        case_sensitive = kwargs.get("case_sensitive", True)
+
+        if not pattern:
+            return ToolResult(success=False, error="패턴이 필요합니다")
+
+        try:
+            flags = 0 if case_sensitive else re.IGNORECASE
+            regex = re.compile(pattern, flags)
+        except re.error as e:
+            return ToolResult(success=False, error=f"잘못된 정규식 패턴: {e}")
+
+        search_path = self._base_path / path
+        if not search_path.exists():
+            return ToolResult(success=False, error=f"경로를 찾을 수 없음: {path}")
+
+        results = []
+        try:
+            for file_path in search_path.rglob(file_pattern):
+                if not file_path.is_file():
+                    continue
+
+                # 바이너리 파일과 숨김 디렉토리 건너뛰기
+                if any(part.startswith(".") for part in file_path.parts):
+                    continue
+
+                try:
+                    content = file_path.read_text(encoding="utf-8", errors="ignore")
+                    for line_num, line in enumerate(content.splitlines(), 1):
+                        if regex.search(line):
+                            rel_path = file_path.relative_to(self._base_path)
+                            results.append(
+                                {
+                                    "file": str(rel_path),
+                                    "line": line_num,
+                                    "content": line.strip(),
+                                }
+                            )
+
+                            if len(results) >= max_results:
+                                break
+                except Exception:
+                    # 읽을 수 없는 파일 건너뛰기
+                    continue
+
+                if len(results) >= max_results:
+                    break
+
+            return ToolResult(
+                success=True,
+                data={"matches": results, "pattern": pattern},
+                metadata={
+                    "count": len(results),
+                    "truncated": len(results) >= max_results,
+                },
+            )
+
+        except Exception as e:
+            return ToolResult(success=False, error=f"검색 오류: {str(e)}")
+
+
+class WebSearchTool(BaseTool):
+    """Tavily API를 사용한 웹 검색 도구."""
+
+    def __init__(self, api_key: Optional[str] = None):
+        self._api_key = api_key
+        self._client = httpx.AsyncClient(timeout=30.0)
+
+    @property
+    def name(self) -> str:
+        return "web_search"
+
+    @property
+    def description(self) -> str:
+        return """Search the web for information.
+Returns relevant search results with titles, URLs, and snippets.
+Useful for finding current information, documentation, or examples."""
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The search query",
+                },
+                "max_results": {
+                    "type": "integer",
+                    "description": "Maximum number of results to return",
+                    "default": 5,
+                },
+            },
+            "required": ["query"],
+        }
+
+    @property
+    def risk_level(self) -> str:
+        return "low"
+
+    def set_api_key(self, api_key: str) -> None:
+        """Tavily API 키를 설정합니다."""
+        self._api_key = api_key
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        """웹 검색을 실행합니다."""
+        query = kwargs.get("query", "")
+        max_results = kwargs.get("max_results", 5)
+
+        if not query:
+            return ToolResult(success=False, error="쿼리가 필요합니다")
+
+        if not self._api_key:
+            return ToolResult(success=False, error="Tavily API 키가 설정되지 않음")
+
+        try:
+            response = await self._client.post(
+                "https://api.tavily.com/search",
+                json={
+                    "api_key": self._api_key,
+                    "query": query,
+                    "max_results": max_results,
+                    "include_answer": True,
+                },
+            )
+            response.raise_for_status()
+
+            data = response.json()
+            results = []
+
+            for result in data.get("results", []):
+                results.append(
+                    {
+                        "title": result.get("title", ""),
+                        "url": result.get("url", ""),
+                        "snippet": result.get("content", ""),
+                    }
+                )
+
+            return ToolResult(
+                success=True,
+                data={
+                    "results": results,
+                    "answer": data.get("answer"),
+                    "query": query,
+                },
+                metadata={"count": len(results)},
+            )
+
+        except httpx.HTTPStatusError as e:
+            return ToolResult(
+                success=False, error=f"검색 API 오류: {e.response.status_code}"
+            )
+        except Exception as e:
+            return ToolResult(success=False, error=f"검색 오류: {str(e)}")
+
+    async def close(self):
+        """HTTP 클라이언트를 닫습니다."""
+        await self._client.aclose()