codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/handlers/rag_admin_handler.py

@@ -0,0 +1,689 @@
+"""
+RAG Admin REST API Handlers.
+
+Admin 모드에서만 접근 가능한 RAG 관리 엔드포인트.
+RAGCollectionManager를 통해 3개 컬렉션 관리.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from typing import Any, Dict, Optional
+
+from jupyter_server.base.handlers import APIHandler
+
+from codebuff_core.legacy_paths import resolve_home
+
+logger = logging.getLogger(__name__)
+
+
+def _is_admin() -> bool:
+    """Admin 모드 판별."""
+    sagemaker_space = os.environ.get("SAGEMAKER_SPACE_NAME")
+    hdsp_prj_id = os.environ.get("HDSP_PRJ_ID")
+    if sagemaker_space is None and hdsp_prj_id is None:
+        return True  # dev
+    if hdsp_prj_id == "local_dev":
+        return True  # local dev (jl.sh sets this)
+    if sagemaker_space and "pl2wadmprj" in sagemaker_space:
+        return True
+    if hdsp_prj_id == "pl2wadmprj":
+        return True
+    return False
+
+
+def _parse_body(handler: APIHandler) -> Dict[str, Any]:
+    """요청 body JSON 파싱."""
+    if not handler.request.body:
+        return {}
+    return json.loads(handler.request.body.decode("utf-8"))
+
+
+def _get_manager():
+    """RAGCollectionManager 싱글턴 획득."""
+
+    # 싱글턴 패턴: settings에 저장
+    # 초기화는 RAGAdminConfigHandler.post 에서 수행
+    return _rag_collection_manager
+
+
+_rag_collection_manager: Optional[Any] = None
+
+
+def _ensure_manager(config: Optional[Dict[str, Any]] = None):
+    """RAGCollectionManager 초기화 (lazy)."""
+    global _rag_collection_manager
+    if _rag_collection_manager is not None:
+        return _rag_collection_manager
+
+    from codebuff_core.services.rag_collections import RAGCollectionManager
+
+    # RAGManager에 이미 Qdrant 클라이언트가 있으면 재사용 (로컬 모드 잠금 충돌 방지)
+    from codebuff_core.services.rag_manager import get_rag_manager
+
+    rag_mgr = get_rag_manager()
+    if rag_mgr._client is not None:
+        client = rag_mgr._client
+    else:
+        try:
+            from qdrant_client import QdrantClient
+        except ImportError:
+            raise ImportError(
+                "qdrant-client가 설치되지 않았습니다. pip install qdrant-client"
+            )
+
+        from codebuff_core.managers.config_manager import get_config_manager
+
+        cfg = get_config_manager().get_config()
+        rag_cfg = cfg.get("rag", {})
+
+        # 환경변수 우선 → config 파일 폴백
+        mode = os.environ.get("QDRANT_MODE", "") or rag_cfg.get("qdrantMode", "local")
+        if mode == "server":
+            url = os.environ.get("QDRANT_URL", "") or rag_cfg.get(
+                "qdrantUrl", "http://localhost:6333"
+            )
+            api_key = os.environ.get("QDRANT_API_KEY", "") or rag_cfg.get(
+                "qdrantApiKey"
+            )
+            client = (
+                QdrantClient(url=url, api_key=api_key)
+                if api_key
+                else QdrantClient(url=url)
+            )
+        else:
+            configured = os.environ.get("QDRANT_PATH", "") or rag_cfg.get(
+                "qdrantPath", ""
+            )
+            local_path = (
+                os.path.expanduser(configured)
+                if configured
+                else str(resolve_home("qdrant"))
+            )
+            os.makedirs(local_path, exist_ok=True)
+            client = QdrantClient(path=local_path)
+
+    # 임베딩 서비스 (리모트 전용)
+    from codebuff_core.services.vllm_embedding_service import (
+        get_vllm_embedding_service,
+    )
+
+    embedding = get_vllm_embedding_service()
+
+    from codebuff_core.services.rag_collections import register_manager
+
+    _rag_collection_manager = RAGCollectionManager(
+        qdrant_client=client,
+        embedding_service=embedding,
+        llm_config=config,
+    )
+    register_manager(_rag_collection_manager)
+    return _rag_collection_manager
+
+
+def reset_manager():
+    """매니저 리셋 (설정 변경 시). 임베딩 서비스 싱글톤도 함께 리셋."""
+    global _rag_collection_manager
+    _rag_collection_manager = None
+    # 임베딩 서비스 싱글톤 리셋 (dimension 등 설정 변경 반영)
+    from codebuff_core.services.vllm_embedding_service import (
+        reset_vllm_embedding_service,
+    )
+
+    reset_vllm_embedding_service()
+
+
+# ============ Config Handler ============
+
+
+class RAGAdminConfigHandler(APIHandler):
+    """GET/POST /codebuff/rag/admin/config
+
+    Qdrant 연결 + 임베딩 + 컬렉션 매핑 설정 조회/저장.
+    """
+
+    async def get(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            from codebuff_core.managers.config_manager import get_config_manager
+
+            cfg = get_config_manager().get_config()
+            rag = cfg.get("rag", {})
+
+            # 임베딩 접속 설정은 읽기 전용이다 — get_config() 가 centri > 시크릿 env >
+            # pod 파일로 해석한 값을 보여 줄 뿐, 여기서 고치지 않는다 (이슈 #113)
+            self.write(
+                {
+                    "qdrantMode": rag.get("qdrantMode", "local"),
+                    "qdrantUrl": rag.get("qdrantUrl", "http://localhost:6333"),
+                    "qdrantApiKey": "***" if rag.get("qdrantApiKey") else "",
+                    "qdrantPath": rag.get("qdrantPath") or str(resolve_home("qdrant")),
+                    "embeddingEndpoint": rag.get("vllmEndpoint", ""),
+                    "embeddingApiKey": "***" if rag.get("vllmApiKey") else "",
+                    "embeddingModel": rag.get("vllmModel", ""),
+                    "embeddingDimension": rag.get("vllmDimension", 4096),
+                    "collections": rag.get("collections", {}),
+                    "chunking": rag.get("chunking", {}),
+                }
+            )
+        except Exception as e:
+            logger.error("RAGAdminConfig GET error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            from codebuff_core.managers.config_manager import get_config_manager
+
+            cm = get_config_manager()
+            cfg = cm.get_config()
+            rag = cfg.get("rag", {})
+
+            # 업데이트 가능한 필드. 임베딩 접속 설정(주소·모델·차원·키)은 받지 않는다 —
+            # centri 가 정본이고 키는 시크릿 env 다 (이슈 #113)
+            for key in (
+                "qdrantMode",
+                "qdrantUrl",
+                "qdrantPath",
+                "collections",
+                "chunking",
+            ):
+                if key in body:
+                    rag[key] = body[key]
+
+            # API 키 필드: 마스킹 값(***) 이면 기존 값 유지
+            if body.get("qdrantApiKey") not in (None, "", "***"):
+                rag["qdrantApiKey"] = body["qdrantApiKey"]
+
+            cfg["rag"] = rag
+            cm.save_config(cfg)
+            reset_manager()  # 설정 변경 시 매니저 리셋
+            self.write({"success": True})
+        except Exception as e:
+            logger.error("RAGAdminConfig POST error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Connection Test ============
+
+
+class RAGAdminTestConnectionHandler(APIHandler):
+    """POST /codebuff/rag/admin/test-connection
+
+    Qdrant 연결 테스트. Body: { "type": "qdrant" }
+    임베딩 연결 확인은 없앴다 — 접속 설정은 centri·env 가 정하고 화면에서 고치지 않는다 (이슈 #113).
+    """
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            test_type = body.get("type", "qdrant")
+
+            if test_type == "qdrant":
+                mgr = _ensure_manager()
+                stats = await mgr.get_collection_stats()
+                self.write({"success": True, "stats": stats})
+            else:
+                self.set_status(400)
+                self.write({"error": f"Unknown type: {test_type}"})
+        except Exception as e:
+            logger.error("Connection test error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"success": False, "error": str(e)})
+
+
+# ============ Stats Handler ============
+
+
+class RAGAdminStatsHandler(APIHandler):
+    """GET /codebuff/rag/admin/stats"""
+
+    async def get(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            mgr = _ensure_manager()
+            stats = await mgr.get_collection_stats()
+            self.write({"stats": stats})
+        except Exception as e:
+            logger.error("RAGAdminStats error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Upsert Handler ============
+
+
+class RAGAdminUpsertHandler(APIHandler):
+    """POST /codebuff/rag/admin/upsert
+
+    Body:
+      - collection: "enterprise_tables" | "bdp_tables" | "guide_docs" | "ds_guide_docs"
+      - filePath: CSV 또는 디렉토리 경로
+      - joinFilePath: (bdp_tables만)
+      - joinKey: (bdp_tables만)
+      - mapping: 컬럼 매핑 (테이블 컬렉션만)
+      - chunking: 청킹 설정 (guide 컬렉션만)
+    """
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            collection = body.get("collection", "")
+            mgr = _ensure_manager()
+
+            # 컬렉션 ensure
+            await mgr.ensure_collections(mgr._embedding.dimension)
+
+            if collection in ("enterprise_tables", "bdp_tables"):
+                result = await mgr.upsert_table_collection(
+                    collection_name=collection,
+                    file_path=body["filePath"],
+                    mapping=body.get("mapping", {}),
+                    join_file_path=body.get("joinFilePath"),
+                    join_key=body.get("joinKey"),
+                )
+            elif collection in ("guide_docs", "ds_guide_docs"):
+                result = await mgr.upsert_guide_directory(
+                    dir_path=body["filePath"],
+                    chunking_config=body.get("chunking"),
+                    collection_name=collection,
+                )
+            elif collection == "faq_docs":
+                result = await mgr.upsert_faq_collection(
+                    file_path=body["filePath"],
+                )
+            else:
+                self.set_status(400)
+                self.write({"error": f"Unknown collection: {collection}"})
+                return
+
+            self.write(result)
+        except Exception as e:
+            logger.error("RAGAdminUpsert error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Delete Collection ============
+
+
+class RAGAdminDeleteHandler(APIHandler):
+    """DELETE /codebuff/rag/admin/collection/{name}"""
+
+    async def delete(self, collection_name: str):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            mgr = _ensure_manager()
+            await mgr.delete_collection(collection_name)
+            self.write({"success": True, "collection": collection_name})
+        except Exception as e:
+            logger.error("RAGAdminDelete error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Preview Handler ============
+
+
+class RAGAdminPreviewHandler(APIHandler):
+    """POST /codebuff/rag/admin/preview
+
+    CSV 헤더 파싱 / 디렉토리 파일 목록 반환.
+    Body: { "filePath": "...", "type": "csv_headers" | "directory_files" }
+    """
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            file_path = body.get("filePath", "")
+            preview_type = body.get("type", "csv_headers")
+            mgr = _ensure_manager()
+
+            if preview_type == "csv_headers":
+                headers = mgr.parse_csv_headers(file_path)
+                self.write({"headers": headers})
+            elif preview_type == "directory_files":
+                files = mgr.list_directory_files(file_path)
+                self.write({"files": files})
+            else:
+                self.set_status(400)
+                self.write({"error": f"Unknown type: {preview_type}"})
+        except Exception as e:
+            logger.error("RAGAdminPreview error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Dry Run ============
+
+
+class RAGAdminDryRunHandler(APIHandler):
+    """POST /codebuff/rag/admin/dry-run
+
+    적재 전 프리뷰.
+    Body: { "collection": "...", "filePath": "...", "mapping": {...}, ... }
+    """
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            collection = body.get("collection", "")
+            mgr = _ensure_manager()
+
+            if collection in ("enterprise_tables", "bdp_tables"):
+                result = await mgr.dry_run_table(
+                    file_path=body["filePath"],
+                    mapping=body.get("mapping", {}),
+                    join_file_path=body.get("joinFilePath"),
+                    join_key=body.get("joinKey"),
+                )
+            elif collection in ("guide_docs", "ds_guide_docs"):
+                result = await mgr.dry_run_guide(
+                    dir_path=body["filePath"],
+                    chunking_config=body.get("chunking"),
+                )
+            else:
+                self.set_status(400)
+                self.write({"error": f"Unknown collection: {collection}"})
+                return
+
+            self.write(result)
+        except Exception as e:
+            logger.error("RAGAdminDryRun error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Validate CSV ============
+
+
+class RAGAdminValidateHandler(APIHandler):
+    """POST /codebuff/rag/admin/validate
+
+    CSV 유효성 검증.
+    Body: { "filePath": "...", "mapping": {...} }
+    """
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            mgr = _ensure_manager()
+            result = mgr.validate_csv_mapping(
+                file_path=body["filePath"],
+                mapping=body.get("mapping", {}),
+            )
+            self.write(result)
+        except Exception as e:
+            logger.error("RAGAdminValidate error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Browse Collection ============
+
+
+class RAGAdminBrowseHandler(APIHandler):
+    """GET /codebuff/rag/admin/browse/{collection}
+
+    포인트 목록 페이징.
+    Query params: offset, limit, filter_field, filter_value
+    """
+
+    async def get(self, collection: str):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+
+            offset_str = self.get_argument("offset", "")
+            offset = offset_str if offset_str else None
+            limit = int(self.get_argument("limit", "10"))
+
+            filters_json = self.get_argument("filters", "")
+            filters = None
+            if filters_json:
+                import json as _json
+
+                try:
+                    filters = _json.loads(filters_json)
+                except Exception:
+                    pass
+
+            mgr = _ensure_manager()
+            result = await mgr.browse_collection(
+                collection=collection,
+                offset=offset,
+                limit=limit,
+                filters=filters,
+            )
+            self.write(result)
+        except Exception as e:
+            logger.error("RAGAdminBrowse error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+    async def delete(self, collection: str):
+        """개별 포인트 삭제."""
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            point_id = body.get("pointId", "")
+            if not point_id:
+                self.set_status(400)
+                self.write({"error": "pointId required"})
+                return
+            mgr = _ensure_manager()
+            await mgr.delete_point(collection, point_id)
+            self.write({"success": True})
+        except Exception as e:
+            logger.error("RAGAdminBrowse DELETE error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Search Test ============
+
+
+class RAGAdminSearchTestHandler(APIHandler):
+    """POST /codebuff/rag/admin/search-test
+
+    Body: { "query": "...", "collection": null | "...", "topK": 5 }
+    """
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            body = _parse_body(self)
+            collection = body.get("collection")
+            query = body.get("query", "")
+            top_k = body.get("topK", 5)
+
+            mgr = _ensure_manager()
+            results = await mgr.search_test(
+                query=query,
+                collection=collection,
+                top_k=top_k,
+            )
+            self.write({"results": results})
+        except Exception as e:
+            logger.error("RAGAdminSearchTest error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ History ============
+
+
+class RAGAdminHistoryHandler(APIHandler):
+    """GET /codebuff/rag/admin/history"""
+
+    async def get(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            mgr = _ensure_manager()
+            history = mgr.get_history()
+            self.write({"history": history})
+        except Exception as e:
+            logger.error("RAGAdminHistory error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Progress Handler ============
+
+
+class RAGAdminProgressHandler(APIHandler):
+    """GET /codebuff/rag/admin/progress
+
+    적재 진행 상황 조회 (폴링).
+    """
+
+    async def get(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            mgr = _get_manager()
+            if mgr is None:
+                self.write({"status": "idle"})
+                return
+            self.write(mgr.get_progress() or {"status": "idle"})
+        except Exception as e:
+            logger.error("RAGAdminProgress error: %s", e, exc_info=True)
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+class RAGAdminCancelHandler(APIHandler):
+    """POST /codebuff/rag/admin/cancel — 적재 중지."""
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            mgr = _get_manager()
+            if mgr is None:
+                self.write({"success": False, "error": "Manager not initialized"})
+                return
+            cancelled = mgr.cancel_upsert()
+            self.write({"success": cancelled})
+        except Exception as e:
+            self.set_status(500)
+            self.write({"error": str(e)})
+
+
+# ============ Guide Pipeline (3-step) ============
+
+
+class RAGAdminGuidePipelineHandler(APIHandler):
+    """POST /codebuff/rag/admin/guide-pipeline — 3단계 파이프라인 실행 (SSE).
+
+    Body: {"sourceDir": "...", "collection": "guide_docs", "step": 1|2|3}
+    """
+
+    async def post(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+
+            body = _parse_body(self)
+            source_dir = body.get("sourceDir", "")
+            collection = body.get("collection", "guide_docs")
+            step = body.get("step", 1)
+            mgr = _ensure_manager()
+
+            self.set_header("Content-Type", "text/event-stream")
+            self.set_header("Cache-Control", "no-cache")
+
+            if step == 1:
+                gen = mgr.guide_pipeline_step1(source_dir)
+            elif step == 2:
+                gen = mgr.guide_pipeline_step2(source_dir)
+            elif step == 3:
+                gen = mgr.guide_pipeline_step3(source_dir, collection)
+            else:
+                self.set_status(400)
+                self.write({"error": f"Invalid step: {step}"})
+                return
+
+            async for event in gen:
+                self.write(f"data: {json.dumps(event, ensure_ascii=False)}\n\n")
+                self.flush()
+
+        except Exception as e:
+            logger.error("GuidePipeline error: %s", e, exc_info=True)
+            self.write(f"data: {json.dumps({'status': 'error', 'error': str(e)})}\n\n")
+            self.flush()
+
+
+class RAGAdminGuidePipelineStatusHandler(APIHandler):
+    """GET /codebuff/rag/admin/guide-pipeline/status"""
+
+    async def get(self):
+        try:
+            if not _is_admin():
+                self.set_status(403)
+                self.write({"error": "Admin only"})
+                return
+            source_dir = self.get_argument("sourceDir", "")
+            if not source_dir:
+                self.set_status(400)
+                self.write({"error": "sourceDir required"})
+                return
+            mgr = _ensure_manager()
+            self.write(mgr.guide_pipeline_status(source_dir))
+        except Exception as e:
+            self.set_status(500)
+            self.write({"error": str(e)})