codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/tests/test_resource_tiers.py

@@ -0,0 +1,298 @@
+"""Tests for resource tier classification and guidance generation."""
+
+from __future__ import annotations
+
+import pytest
+
+from ..resource_usage import (
+    _classify_cpu,
+    _classify_memory,
+    _classify_single_gpu,
+    _classify_gpu,
+    classify_resource_tiers,
+    generate_tier_guidance,
+)
+
+
+# ---------------------------------------------------------------------------
+# CPU tier classification
+# ---------------------------------------------------------------------------
+
+
+class TestClassifyCPU:
+    @pytest.mark.parametrize(
+        "cores,expected",
+        [
+            (1, "low"),
+            (2, "low"),
+            (3, "low"),
+            (4, "medium"),
+            (7, "medium"),
+            (8, "high"),
+            (16, "high"),
+            (31, "high"),
+            (32, "very_high"),
+            (56, "very_high"),
+            (128, "very_high"),
+            (2.5, "low"),  # fractional (K8s pod)
+            (None, "unknown"),
+        ],
+    )
+    def test_cpu_tiers(self, cores, expected):
+        assert _classify_cpu(cores) == expected
+
+
+# ---------------------------------------------------------------------------
+# Memory tier classification
+# ---------------------------------------------------------------------------
+
+
+class TestClassifyMemory:
+    @pytest.mark.parametrize(
+        "total_gb,expected",
+        [
+            (2, "constrained"),
+            (3.9, "constrained"),
+            (4, "limited"),
+            (15, "limited"),
+            (16, "comfortable"),
+            (63, "comfortable"),
+            (64, "large"),
+            (128, "large"),
+            (255, "large"),
+            (256, "very_large"),
+            (900, "very_large"),
+            (None, "unknown"),
+        ],
+    )
+    def test_memory_tiers(self, total_gb, expected):
+        assert _classify_memory(total_gb) == expected
+
+
+# ---------------------------------------------------------------------------
+# GPU tier classification
+# ---------------------------------------------------------------------------
+
+
+class TestClassifyGPU:
+    @pytest.mark.parametrize(
+        "name,vram_mb,expected",
+        [
+            ("NVIDIA L4", 24576, "entry"),
+            ("Tesla T4", 16384, "entry"),
+            ("Tesla V100-SXM2-16GB", 16384, "mid"),
+            ("Tesla V100-SXM2-32GB", 32768, "high"),
+            ("NVIDIA A100-SXM4-40GB", 40960, "high"),
+            ("NVIDIA A100-SXM4-80GB", 81920, "ultra"),
+            ("NVIDIA H100 80GB HBM3", 81920, "ultra"),
+            ("NVIDIA H200", 144000, "ultra"),
+            ("NVIDIA RTX 3060", 12288, "entry"),
+            ("NVIDIA RTX 4090", 24576, "high"),
+            ("Unknown GPU", 8192, "entry"),
+            ("Unknown GPU", 32768, "high"),
+            ("Unknown GPU", 49152, "ultra"),
+        ],
+    )
+    def test_single_gpu_tiers(self, name, vram_mb, expected):
+        assert _classify_single_gpu(name, vram_mb) == expected
+
+    def test_no_gpu(self):
+        result = _classify_gpu([])
+        assert result["tier"] == "none"
+        assert result["count"] == 0
+
+    def test_single_gpu(self):
+        gpus = [{"name": "NVIDIA A100-SXM4-80GB", "memory_total_mb": 81920}]
+        result = _classify_gpu(gpus)
+        assert result["tier"] == "ultra"
+        assert result["count"] == 1
+        assert not result["multi_gpu"]
+
+    def test_multi_gpu(self):
+        gpus = [
+            {"name": "NVIDIA H200", "memory_total_mb": 144000},
+            {"name": "NVIDIA H200", "memory_total_mb": 144000},
+            {"name": "NVIDIA H200", "memory_total_mb": 144000},
+            {"name": "NVIDIA H200", "memory_total_mb": 144000},
+        ]
+        result = _classify_gpu(gpus)
+        assert result["tier"] == "ultra"
+        assert result["count"] == 4
+        assert result["multi_gpu"]
+
+
+# ---------------------------------------------------------------------------
+# classify_resource_tiers integration
+# ---------------------------------------------------------------------------
+
+
+class TestClassifyResourceTiers:
+    def _make_data(self, cores=8, total_gb=32, avail_gb=24, gpus=None):
+        return {
+            "cpu": {"cores": cores, "usage_percent": 30.0},
+            "memory": {
+                "total_gb": total_gb,
+                "available_gb": avail_gb,
+                "used_gb": total_gb - avail_gb,
+            },
+            "gpus": gpus or [],
+            "gpu_status": "available" if gpus else "not_detected",
+        }
+
+    def test_mid_workstation(self):
+        tiers = classify_resource_tiers(
+            self._make_data(cores=7, total_gb=128, avail_gb=100)
+        )
+        assert tiers["cpu_tier"] == "medium"
+        assert tiers["mem_tier"] == "large"
+        assert tiers["gpu"]["tier"] == "none"
+
+    def test_high_server(self):
+        gpus = [
+            {"name": "NVIDIA A100-SXM4-80GB", "memory_total_mb": 81920},
+            {"name": "NVIDIA A100-SXM4-80GB", "memory_total_mb": 81920},
+        ]
+        tiers = classify_resource_tiers(
+            self._make_data(cores=56, total_gb=900, avail_gb=850, gpus=gpus)
+        )
+        assert tiers["cpu_tier"] == "very_high"
+        assert tiers["mem_tier"] == "very_large"
+        assert tiers["gpu"]["tier"] == "ultra"
+        assert tiers["gpu"]["multi_gpu"]
+
+    def test_low_container(self):
+        tiers = classify_resource_tiers(
+            self._make_data(cores=2, total_gb=4, avail_gb=2)
+        )
+        assert tiers["cpu_tier"] == "low"
+        assert tiers["mem_tier"] == "limited"
+
+
+# ---------------------------------------------------------------------------
+# Guidance generation — differentiation tests
+# ---------------------------------------------------------------------------
+
+
+class TestGenerateTierGuidance:
+    def _make_data(self, cores=8, total_gb=32, avail_gb=24, gpus=None):
+        return {
+            "cpu": {"cores": cores, "usage_percent": 30.0},
+            "memory": {
+                "total_gb": total_gb,
+                "available_gb": avail_gb,
+                "used_gb": total_gb - avail_gb,
+            },
+            "gpus": gpus or [],
+            "gpu_status": "available" if gpus else "not_detected",
+        }
+
+    def test_scenario_a_mid_workstation_with_l4(self):
+        """7코어/128GB/L4 → Polars 권장, joblib, LoRA"""
+        gpus = [{"name": "NVIDIA L4", "memory_total_mb": 24576}]
+        tiers = classify_resource_tiers(
+            self._make_data(cores=7, total_gb=128, avail_gb=100, gpus=gpus)
+        )
+        guidance = generate_tier_guidance(tiers)
+
+        assert "joblib" in guidance
+        assert "n_jobs=6" in guidance or "n_jobs" in guidance
+        assert "Polars" in guidance
+        assert "LoRA" in guidance or "entry" in guidance
+
+    def test_scenario_b_high_server_with_a100(self):
+        """56코어/900GB/A100x2 → Dask/Ray, cuDF, DDP"""
+        gpus = [
+            {"name": "NVIDIA A100-SXM4-80GB", "memory_total_mb": 81920},
+            {"name": "NVIDIA A100-SXM4-80GB", "memory_total_mb": 81920},
+        ]
+        tiers = classify_resource_tiers(
+            self._make_data(cores=56, total_gb=900, avail_gb=850, gpus=gpus)
+        )
+        guidance = generate_tier_guidance(tiers)
+
+        assert "Dask" in guidance or "Ray" in guidance
+        assert "cuDF" in guidance or "cudf" in guidance.lower()
+        assert "DDP" in guidance or "FSDP" in guidance or "멀티GPU" in guidance
+        assert "서버급" in guidance
+
+    def test_scenario_c_high_server_no_gpu(self):
+        """56코어/900GB/GPU없음 → Dask/Ray CPU, Polars 멀티스레드"""
+        tiers = classify_resource_tiers(
+            self._make_data(cores=56, total_gb=900, avail_gb=850)
+        )
+        guidance = generate_tier_guidance(tiers)
+
+        assert "Dask" in guidance or "Ray" in guidance
+        assert "Polars" in guidance
+        assert "GPU 미감지" in guidance
+        assert "device='cpu'" in guidance or "CPU" in guidance
+
+    def test_scenario_d_low_container(self):
+        """2코어/4GB/GPU없음 → chunked, 병렬처리 비권장"""
+        tiers = classify_resource_tiers(
+            self._make_data(cores=2, total_gb=4, avail_gb=2)
+        )
+        guidance = generate_tier_guidance(tiers)
+
+        assert "저사양" in guidance
+        assert "병렬처리" in guidance
+        assert "dtypes" in guidance or "chunked" in guidance
+
+    def test_scenarios_produce_different_guidance(self):
+        """핵심: 서로 다른 환경은 확실히 다른 가이드를 생성해야 함."""
+        gpus_l4 = [{"name": "NVIDIA L4", "memory_total_mb": 24576}]
+        gpus_a100 = [
+            {"name": "NVIDIA A100-SXM4-80GB", "memory_total_mb": 81920},
+            {"name": "NVIDIA A100-SXM4-80GB", "memory_total_mb": 81920},
+        ]
+
+        tiers_a = classify_resource_tiers(
+            self._make_data(cores=7, total_gb=128, avail_gb=100, gpus=gpus_l4)
+        )
+        tiers_b = classify_resource_tiers(
+            self._make_data(cores=56, total_gb=900, avail_gb=850, gpus=gpus_a100)
+        )
+
+        guidance_a = generate_tier_guidance(tiers_a)
+        guidance_b = generate_tier_guidance(tiers_b)
+
+        # Must be substantially different
+        assert guidance_a != guidance_b
+
+        # Scenario A should NOT contain server-grade terms
+        assert "서버급" not in guidance_a
+        assert "Spark" not in guidance_a
+
+        # Scenario B should NOT contain low-tier terms
+        assert "저사양" not in guidance_b
+        assert "LoRA" not in guidance_b
+
+    def test_mode_codegen_differs_from_review(self):
+        """codegen 모드는 다른 헤더를 사용."""
+        tiers = classify_resource_tiers(self._make_data())
+        review = generate_tier_guidance(tiers, mode="review")
+        codegen = generate_tier_guidance(tiers, mode="codegen")
+
+        assert "환경 맞춤 기술 스택 권장" in review
+        assert "기술 스택 지침" in codegen
+
+    def test_memory_pressure_warning(self):
+        """가용 메모리가 전체의 30% 미만일 때 경고 출력."""
+        tiers = classify_resource_tiers(
+            self._make_data(cores=16, total_gb=128, avail_gb=10)
+        )
+        guidance = generate_tier_guidance(tiers)
+
+        assert "메모리 압박" in guidance
+
+    def test_h200_multi_gpu_guidance(self):
+        """4x H200 → tensor parallelism, FSDP"""
+        gpus = [{"name": "NVIDIA H200", "memory_total_mb": 144000}] * 4
+        tiers = classify_resource_tiers(
+            self._make_data(cores=56, total_gb=900, avail_gb=850, gpus=gpus)
+        )
+        guidance = generate_tier_guidance(tiers)
+
+        assert "FSDP" in guidance or "DeepSpeed" in guidance
+        assert "멀티GPU" in guidance
+        assert "tensor" in guidance.lower() or "pipeline" in guidance.lower()

jupyter_ext/tools/__init__.py

@@ -0,0 +1,15 @@
+"""Tool system for Codebuff."""
+
+from .base import BaseTool, ToolResult
+from .registry import ToolRegistry
+from .context import NotebookContextManager, NotebookContext, CellInfo
+
+__all__ = [
+    "BaseTool",
+    "ToolResult",
+    "ToolRegistry",
+    # Notebook context management
+    "NotebookContextManager",
+    "NotebookContext",
+    "CellInfo",
+]

jupyter_ext/tools/base.py

@@ -0,0 +1,90 @@
+"""기본 도구 인터페이스."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+
+@dataclass
+class ToolResult:
+    """도구 실행 결과."""
+
+    success: bool
+    data: Optional[Any] = None
+    error: Optional[str] = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """딕셔너리로 변환."""
+        result = {"success": self.success}
+        if self.data is not None:
+            result["data"] = self.data
+        if self.error is not None:
+            result["error"] = self.error
+        if self.metadata:
+            result["metadata"] = self.metadata
+        return result
+
+
+class BaseTool(ABC):
+    """도구의 추상 기본 클래스."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """도구 이름."""
+        pass
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """도구 설명."""
+        pass
+
+    @property
+    @abstractmethod
+    def parameters(self) -> dict[str, Any]:
+        """도구 매개변수 스키마 (JSON Schema 형식)."""
+        pass
+
+    @property
+    def risk_level(self) -> str:
+        """
+        도구의 위험 수준.
+
+        Returns:
+            다음 중 하나: "low", "medium", "high", "critical"
+        """
+        return "low"
+
+    @property
+    def definition(self) -> dict[str, Any]:
+        """
+        OpenAI 호환 도구 정의를 가져옵니다.
+
+        Returns:
+            도구 정의 딕셔너리
+        """
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters,
+            },
+        }
+
+    @abstractmethod
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        """
+        도구를 실행합니다.
+
+        Args:
+            **kwargs: 도구 인자
+
+        Returns:
+            실행 결과가 담긴 ToolResult
+        """
+        pass

jupyter_ext/tools/context.py

@@ -0,0 +1,268 @@
+"""Jupyter 도구를 위한 노트북 컨텍스트 관리."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+
+@dataclass
+class CellInfo:
+    """노트북 셀 정보."""
+
+    index: int
+    cell_type: str  # "code" 또는 "markdown"
+    source: str
+    outputs: list[dict[str, Any]] = field(default_factory=list)
+    execution_count: Optional[int] = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class NotebookContext:
+    """현재 노트북에 대한 컨텍스트 정보."""
+
+    path: str
+    cells: list[CellInfo] = field(default_factory=list)
+    kernel_name: Optional[str] = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def cell_count(self) -> int:
+        """노트북의 셀 수를 가져옵니다."""
+        return len(self.cells)
+
+    def get_cell(self, index: int) -> Optional[CellInfo]:
+        """인덱스로 셀을 가져옵니다."""
+        if 0 <= index < len(self.cells):
+            return self.cells[index]
+        return None
+
+    def get_code_cells(self) -> list[CellInfo]:
+        """모든 코드 셀을 가져옵니다."""
+        return [c for c in self.cells if c.cell_type == "code"]
+
+    def get_markdown_cells(self) -> list[CellInfo]:
+        """모든 마크다운 셀을 가져옵니다."""
+        return [c for c in self.cells if c.cell_type == "markdown"]
+
+
+class NotebookContextManager:
+    """
+    노트북 컨텍스트 관리자로 노트북 상태에 대한 접근을 제공합니다.
+
+    이 클래스는 Jupyter 서버의 노트북 모델과 LangChain Agent 도구 사이를
+    연결하여 도구가 노트북 상태를 읽고 수정할 수 있게 합니다.
+    """
+
+    def __init__(
+        self,
+        kernel_manager: Any = None,
+        contents_manager: Any = None,
+    ):
+        """
+        컨텍스트 관리자를 초기화합니다.
+
+        Args:
+            kernel_manager: 코드 실행을 위한 Jupyter 커널 관리자
+            contents_manager: 파일 작업을 위한 Jupyter 콘텐츠 관리자
+        """
+        self._kernel_manager = kernel_manager
+        self._contents_manager = contents_manager
+        self._current_notebook: Optional[NotebookContext] = None
+        self._notebook_model: Optional[dict[str, Any]] = None
+
+    @property
+    def current_notebook(self) -> Optional[NotebookContext]:
+        """현재 노트북 컨텍스트를 가져옵니다."""
+        return self._current_notebook
+
+    @property
+    def notebook_model(self) -> Optional[dict[str, Any]]:
+        """원시 노트북 모델을 가져옵니다 (직접 조작용)."""
+        return self._notebook_model
+
+    @property
+    def kernel_manager(self) -> Any:
+        """커널 관리자를 가져옵니다."""
+        return self._kernel_manager
+
+    @property
+    def contents_manager(self) -> Any:
+        """콘텐츠 관리자를 가져옵니다."""
+        return self._contents_manager
+
+    def set_notebook(self, path: str, model: dict[str, Any]) -> None:
+        """
+        Jupyter 노트북 모델에서 현재 노트북 컨텍스트를 설정합니다.
+
+        Args:
+            path: 노트북 경로
+            model: Jupyter 노트북 모델 딕셔너리
+        """
+        self._notebook_model = model
+
+        # 모델에서 셀 파싱
+        content = model.get("content", {})
+        raw_cells = content.get("cells", [])
+        cells = []
+
+        for i, cell in enumerate(raw_cells):
+            cells.append(
+                CellInfo(
+                    index=i,
+                    cell_type=cell.get("cell_type", "code"),
+                    source=cell.get("source", ""),
+                    outputs=cell.get("outputs", []),
+                    execution_count=cell.get("execution_count"),
+                    metadata=cell.get("metadata", {}),
+                )
+            )
+
+        # 노트북 메타데이터 추출
+        nb_metadata = content.get("metadata", {})
+        kernel_info = nb_metadata.get("kernelspec", {})
+
+        self._current_notebook = NotebookContext(
+            path=path,
+            cells=cells,
+            kernel_name=kernel_info.get("name"),
+            metadata=nb_metadata,
+        )
+
+    def clear(self) -> None:
+        """현재 노트북 컨텍스트를 지웁니다."""
+        self._current_notebook = None
+        self._notebook_model = None
+
+    def update_cell(self, index: int, source: str) -> bool:
+        """
+        셀의 소스 코드를 업데이트합니다.
+
+        Args:
+            index: 셀 인덱스
+            source: 새 소스 코드
+
+        Returns:
+            성공하면 True
+        """
+        if not self._notebook_model or not self._current_notebook:
+            return False
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+        if not 0 <= index < len(cells):
+            return False
+
+        cells[index]["source"] = source
+
+        # 컨텍스트 업데이트
+        if 0 <= index < len(self._current_notebook.cells):
+            self._current_notebook.cells[index].source = source
+
+        return True
+
+    def insert_cell(
+        self,
+        source: str,
+        cell_type: str = "code",
+        insert_after: Optional[int] = None,
+    ) -> int:
+        """
+        새 셀을 삽입합니다.
+
+        Args:
+            source: 셀 소스 내용
+            cell_type: "code" 또는 "markdown"
+            insert_after: 이 인덱스 뒤에 삽입 (None = 끝에 추가)
+
+        Returns:
+            새 셀의 인덱스
+        """
+        if not self._notebook_model:
+            raise ValueError("노트북 모델이 설정되지 않음")
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+
+        new_cell = {
+            "cell_type": cell_type,
+            "source": source,
+            "metadata": {},
+        }
+
+        if cell_type == "code":
+            new_cell["outputs"] = []
+            new_cell["execution_count"] = None
+
+        if insert_after is not None and 0 <= insert_after < len(cells):
+            cells.insert(insert_after + 1, new_cell)
+            new_index = insert_after + 1
+        else:
+            cells.append(new_cell)
+            new_index = len(cells) - 1
+
+        # 컨텍스트 업데이트
+        if self._current_notebook:
+            new_cell_info = CellInfo(
+                index=new_index,
+                cell_type=cell_type,
+                source=source,
+            )
+
+            if insert_after is not None and 0 <= insert_after < len(
+                self._current_notebook.cells
+            ):
+                self._current_notebook.cells.insert(insert_after + 1, new_cell_info)
+                # 이후 셀 인덱스 재설정
+                for i in range(insert_after + 2, len(self._current_notebook.cells)):
+                    self._current_notebook.cells[i].index = i
+            else:
+                self._current_notebook.cells.append(new_cell_info)
+
+        return new_index
+
+    def delete_cell(self, index: int) -> bool:
+        """
+        셀을 삭제합니다.
+
+        Args:
+            index: 삭제할 셀 인덱스
+
+        Returns:
+            성공하면 True
+        """
+        if not self._notebook_model or not self._current_notebook:
+            return False
+
+        cells = self._notebook_model.get("content", {}).get("cells", [])
+        if not 0 <= index < len(cells):
+            return False
+
+        del cells[index]
+
+        # 컨텍스트 업데이트
+        if 0 <= index < len(self._current_notebook.cells):
+            del self._current_notebook.cells[index]
+            # 이후 셀 인덱스 재설정
+            for i in range(index, len(self._current_notebook.cells)):
+                self._current_notebook.cells[i].index = i
+
+        return True
+
+    def get_summary(self) -> dict[str, Any]:
+        """
+        현재 노트북 상태의 요약을 가져옵니다.
+
+        Returns:
+            노트북 요약이 담긴 딕셔너리
+        """
+        if not self._current_notebook:
+            return {"status": "no_notebook"}
+
+        nb = self._current_notebook
+        return {
+            "path": nb.path,
+            "cell_count": nb.cell_count,
+            "code_cells": len(nb.get_code_cells()),
+            "markdown_cells": len(nb.get_markdown_cells()),
+            "kernel": nb.kernel_name,
+        }