codebuff-agent 2.1.74__py3-none-any.whl

codebuff_core/services/rag_service.py

@@ -0,0 +1,145 @@
+"""
+RAG 서비스 구현
+
+IRAGService의 Embedded 구현.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from codebuff_core.interfaces import IRAGService
+from codebuff_core.models.rag import SearchRequest, SearchResponse
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddedRAGService(IRAGService):
+    """
+    RAG 서비스 구현.
+
+    인프로세스에서 직접 RAGManager를 사용합니다.
+    """
+
+    def __init__(self):
+        """Embedded RAG 서비스 초기화"""
+        self._rag_manager = None
+        self._initialized = False
+        logger.info("EmbeddedRAGService 생성됨 (아직 초기화되지 않음)")
+
+    async def initialize(self) -> None:
+        """RAG 관리자 초기화."""
+        if self._initialized:
+            return
+        try:
+            from codebuff_core.services.rag_manager import get_rag_manager
+
+            self._rag_manager = get_rag_manager()
+            success = await self._rag_manager.initialize()
+            self._initialized = success
+            if success:
+                logger.info("EmbeddedRAGService: RAG 초기화 완료")
+            else:
+                logger.info("EmbeddedRAGService: RAG 비활성화 (설정)")
+        except Exception as e:
+            logger.warning(f"EmbeddedRAGService: RAG 초기화 실패: {e}")
+            self._initialized = False
+
+    async def shutdown(self) -> None:
+        """RAG 서비스 종료"""
+        self._rag_manager = None
+        self._initialized = False
+        logger.info("EmbeddedRAGService 종료됨")
+
+    def is_ready(self) -> bool:
+        """RAG 서비스 준비 상태 확인"""
+        if not self._initialized or not self._rag_manager:
+            return False
+        return self._rag_manager.is_ready
+
+    async def search(self, request: SearchRequest) -> SearchResponse:
+        """지식 베이스 검색"""
+        logger.info(f"[Embedded] RAG 검색: {request.query[:100]}...")
+
+        if not self.is_ready():
+            return SearchResponse(
+                results=[],
+                total=0,
+                query=request.query,
+            )
+
+        results = await self._rag_manager.search(
+            query=request.query,
+            collection=request.collection,
+            limit=request.limit,
+            score_threshold=request.scoreThreshold,
+        )
+
+        return SearchResponse(
+            results=results.get("results", []),
+            total=results.get("total", 0),
+            query=request.query,
+        )
+
+    async def get_context_for_query(
+        self,
+        query: str,
+        detected_libraries: Optional[List[str]] = None,
+        max_results: int = 5,
+    ) -> Optional[str]:
+        """쿼리에 대한 포맷된 컨텍스트 가져오기"""
+        logger.info(f"[Embedded] 쿼리 컨텍스트 조회: {query[:100]}...")
+
+        if not self.is_ready():
+            return None
+
+        try:
+            return await self._rag_manager.get_context_for_query(
+                query=query,
+                detected_libraries=detected_libraries,
+                max_results=max_results,
+            )
+        except Exception as e:
+            logger.error(f"RAG 컨텍스트 조회 실패: {e}")
+            return None
+
+    async def get_index_status(self) -> Dict[str, Any]:
+        """현재 인덱스 상태 가져오기"""
+        logger.info("[Embedded] 인덱스 상태 조회")
+
+        if not self.is_ready():
+            return {
+                "ready": False,
+                "documentCount": 0,
+                "collections": [],
+                "message": "RAG 서비스가 초기화되지 않음",
+            }
+
+        try:
+            return await self._rag_manager.get_index_status()
+        except Exception as e:
+            logger.error(f"인덱스 상태 조회 실패: {e}")
+            return {
+                "ready": False,
+                "documentCount": 0,
+                "collections": [],
+                "message": str(e),
+            }
+
+    async def trigger_reindex(self, force: bool = False) -> Dict[str, Any]:
+        """재인덱싱 작업 트리거"""
+        logger.info(f"[Embedded] 재인덱싱 트리거 (force={force})")
+
+        if not self.is_ready():
+            return {
+                "success": False,
+                "message": "RAG 서비스가 초기화되지 않음",
+            }
+
+        try:
+            return await self._rag_manager.trigger_reindex(force=force)
+        except Exception as e:
+            logger.error(f"재인덱싱 트리거 실패: {e}")
+            return {
+                "success": False,
+                "message": str(e),
+            }

codebuff_core/services/retriever.py

@@ -0,0 +1,260 @@
+"""
+Retriever - Dense vector search implementation.
+
+Provides semantic similarity search via Qdrant vector database.
+"""
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional
+
+if TYPE_CHECKING:
+    from codebuff_core.models.rag import RAGConfig
+
+    from codebuff_core.services.vllm_embedding_service import VLLMEmbeddingService
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ChunkScoreDetails:
+    """청크별 점수 상세 정보"""
+
+    chunk_id: str
+    content: str
+    score: float  # Vector similarity score (0-1)
+    rank: int
+    metadata: Dict[str, Any]
+    passed_threshold: bool
+
+
+class DebugSearchResult(NamedTuple):
+    """디버그 검색 결과"""
+
+    chunks: List[ChunkScoreDetails]
+    search_ms: float
+    total_candidates: int
+
+
+class Retriever:
+    """
+    Dense vector retrieval using Qdrant.
+
+    Features:
+    - Dense vector search via Qdrant
+    - Metadata filtering
+    - Score thresholding
+
+    Usage:
+        retriever = Retriever(client, embedding_service, config)
+        results = await retriever.search("query", top_k=5)
+    """
+
+    def __init__(
+        self,
+        client,  # Qdrant client
+        embedding_service: "VLLMEmbeddingService",
+        config: "RAGConfig",
+    ):
+        self._client = client
+        self._embedding_service = embedding_service
+        self._config = config
+
+    async def search(
+        self,
+        query: str,
+        collection_name: Optional[str] = None,
+        top_k: Optional[int] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        score_threshold: Optional[float] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Perform dense vector search.
+
+        Args:
+            query: Search query
+            collection_name: Target collection (default from config)
+            top_k: Number of results (default from config)
+            filters: Metadata filters
+            score_threshold: Minimum score (default from config)
+
+        Returns:
+            List of results with content, score, metadata
+        """
+        effective_top_k = top_k or self._config.top_k
+        target_collection = collection_name or self._config.qdrant.collection_name
+        effective_threshold = score_threshold or self._config.score_threshold
+
+        # Generate query embedding
+        query_embedding = await self._embedding_service.embed_query(query)
+
+        # Build filter condition
+        qdrant_filter = self._build_filter(filters) if filters else None
+
+        # Dense vector search
+        try:
+            response = self._client.query_points(
+                collection_name=target_collection,
+                query=query_embedding,
+                query_filter=qdrant_filter,
+                limit=effective_top_k,
+                score_threshold=effective_threshold
+                * 0.5,  # Lower for initial retrieval
+                with_payload=True,
+                with_vectors=False,
+            )
+            results = response.points
+        except Exception as e:
+            logger.error(f"Search failed: {e}")
+            return []
+
+        if not results:
+            logger.debug(f"No results for query: {query[:50]}...")
+            return []
+
+        return self._format_results(results, effective_threshold)
+
+    def _build_filter(self, filters: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Convert filter dict to Qdrant filter format"""
+        if not filters:
+            return None
+
+        conditions = []
+        for key, value in filters.items():
+            if isinstance(value, list):
+                # Multiple values - any match
+                conditions.append(
+                    {"should": [{"key": key, "match": {"value": v}} for v in value]}
+                )
+            else:
+                conditions.append({"key": key, "match": {"value": value}})
+
+        return {"must": conditions} if conditions else None
+
+    def _format_results(
+        self, results: List, score_threshold: float
+    ) -> List[Dict[str, Any]]:
+        """Format Qdrant results to standard format"""
+        formatted = []
+        for r in results:
+            if r.score < score_threshold:
+                continue
+
+            formatted.append(
+                {
+                    "content": r.payload.get("content", ""),
+                    "score": round(r.score, 4),
+                    "metadata": {k: v for k, v in r.payload.items() if k != "content"},
+                }
+            )
+        return formatted
+
+    def search_sync(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Synchronous search wrapper for non-async contexts.
+
+        Note: Qdrant client operations are synchronous,
+        so this is just a convenience method.
+        """
+        import asyncio
+
+        try:
+            asyncio.get_running_loop()
+            logger.warning(
+                "search_sync called from async context, use search() instead"
+            )
+        except RuntimeError:
+            pass
+
+        return asyncio.run(self.search(query, top_k, filters))
+
+    async def search_with_debug(
+        self,
+        query: str,
+        collection_name: Optional[str] = None,
+        top_k: Optional[int] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        score_threshold: Optional[float] = None,
+    ) -> DebugSearchResult:
+        """
+        전체 점수 정보를 포함한 디버그 검색 수행.
+
+        Args:
+            query: Search query
+            collection_name: Target collection (default from config)
+            top_k: Number of results (default from config)
+            filters: Metadata filters
+            score_threshold: Minimum score (default from config)
+
+        Returns:
+            DebugSearchResult with detailed scoring information
+        """
+        start_time = time.perf_counter()
+
+        effective_top_k = top_k or self._config.top_k
+        target_collection = collection_name or self._config.qdrant.collection_name
+        effective_threshold = score_threshold or self._config.score_threshold
+
+        # Generate query embedding
+        query_embedding = await self._embedding_service.embed_query(query)
+
+        # Build filter condition
+        qdrant_filter = self._build_filter(filters) if filters else None
+
+        # Vector search with timing
+        try:
+            # 디버그용으로 더 많은 결과 (3배)를 낮은 threshold로 가져옴
+            response = self._client.query_points(
+                collection_name=target_collection,
+                query=query_embedding,
+                query_filter=qdrant_filter,
+                limit=effective_top_k * 3,
+                score_threshold=effective_threshold * 0.3,
+                with_payload=True,
+                with_vectors=False,
+            )
+            results = response.points
+        except Exception as e:
+            logger.error(f"Search failed: {e}")
+            return DebugSearchResult(
+                chunks=[],
+                search_ms=0.0,
+                total_candidates=0,
+            )
+
+        search_ms = (time.perf_counter() - start_time) * 1000
+
+        if not results:
+            return DebugSearchResult(
+                chunks=[],
+                search_ms=round(search_ms, 2),
+                total_candidates=0,
+            )
+
+        # Build detailed results
+        chunks = []
+        for rank, result in enumerate(results, start=1):
+            chunks.append(
+                ChunkScoreDetails(
+                    chunk_id=str(result.id),
+                    content=result.payload.get("content", ""),
+                    score=round(result.score, 4),
+                    rank=rank,
+                    metadata={
+                        k: v for k, v in result.payload.items() if k != "content"
+                    },
+                    passed_threshold=result.score >= effective_threshold,
+                )
+            )
+
+        return DebugSearchResult(
+            chunks=chunks,
+            search_ms=round(search_ms, 2),
+            total_candidates=len(results),
+        )

codebuff_core/services/vllm_embedding_service.py

@@ -0,0 +1,290 @@
+"""
+vLLM Embedding Service - Remote embedding generation using vLLM server.
+
+Features:
+- GPU-accelerated embeddings via vLLM server
+- OpenAI-compatible API interface
+- Retry logic for reliability
+- Support for large models (qwen3-embedding-8b, gte-Qwen2-7B, etc.)
+
+Prerequisites:
+- vLLM embedding server running (e.g., http://10.222.52.31:8000)
+- Model loaded on vLLM server
+"""
+
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+import httpx
+
+if TYPE_CHECKING:
+    from codebuff_core.models.rag import EmbeddingConfig
+
+logger = logging.getLogger(__name__)
+
+
+class VLLMEmbeddingService:
+    """
+    Remote embedding generation using vLLM server.
+
+    Design Principles:
+    - Stateless client (vLLM server holds the model)
+    - Retry logic for network resilience
+    - OpenAI-compatible API interface
+
+    Usage:
+        service = get_vllm_embedding_service()
+        embeddings = service.embed_texts(["text1", "text2"])
+        query_embedding = service.embed_query("search query")
+    """
+
+    _instance: Optional["VLLMEmbeddingService"] = None
+    _initialized: bool = False
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, config: Optional["EmbeddingConfig"] = None):
+        if self._initialized:
+            return
+        self._initialized = True
+
+        from codebuff_core.models.rag import EmbeddingConfig
+
+        self._config = config or EmbeddingConfig()
+
+        # 환경변수 우선 → config 파일 폴백
+        from codebuff_core.managers.config_manager import get_config_manager
+
+        cfg = get_config_manager().get_config()
+        rag = cfg.get("rag", {})
+
+        # get_config() 가 centri > env > 파일로 해석해 둔 값이다 (이슈 #113)
+        self._endpoint = rag.get("vllmEndpoint", "") or "http://localhost:8000"
+        self._model = rag.get("vllmModel", "") or "qwen3-embedding-8b"
+        self._dimension = int(rag.get("vllmDimension", 4096) or 4096)
+        rag_api_key = rag.get("vllmApiKey", "")
+        if rag_api_key in ("***", ""):
+            rag_api_key = ""
+        self._api_key = rag_api_key or cfg.get("vllm", {}).get("apiKey", "")
+
+        # HTTP client with retry (include auth header if API key provided)
+        headers = {}
+        if self._api_key:
+            headers["Authorization"] = f"Bearer {self._api_key}"
+        self._client = httpx.AsyncClient(
+            base_url=self._endpoint,
+            timeout=httpx.Timeout(None),
+            limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
+            headers=headers,
+        )
+
+        logger.info(
+            f"vLLM Embedding Service initialized: "
+            f"endpoint={self._endpoint}, model={self._model}, dim={self._dimension}, "
+            f"auth={'enabled' if self._api_key else 'disabled'}"
+        )
+
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension"""
+        return self._dimension
+
+    async def _call_vllm_api(
+        self, texts: List[str], max_retries: int = 3, timeout: Optional[float] = None
+    ) -> List[List[float]]:
+        """
+        Call vLLM embedding API with retry logic.
+
+        Args:
+            texts: List of text strings to embed
+            max_retries: Maximum number of retry attempts
+
+        Returns:
+            List of embedding vectors
+
+        Raises:
+            Exception if all retries fail
+        """
+        payload = {
+            "model": self._model,
+            "input": texts,
+            "dimensions": self._dimension,
+        }
+
+        last_error = None
+        for attempt in range(max_retries):
+            try:
+                # base_url에 이미 /v1이 포함될 수 있으므로 상대 경로 사용
+                base = self._endpoint.rstrip("/")
+                if base.endswith("/v1"):
+                    path = "/embeddings"
+                else:
+                    path = "/v1/embeddings"
+                response = await self._client.post(path, json=payload, timeout=timeout)
+                response.raise_for_status()
+
+                data = response.json()
+                # OpenRouter 등 일부 서비스는 HTTP 200으로 오면서 error body를 반환
+                if "error" in data and "data" not in data:
+                    err_msg = data["error"].get("message", str(data["error"]))
+                    err_code = data["error"].get("code", "")
+                    raise Exception(f"API error {err_code}: {err_msg}")
+                # Sort by index to ensure correct order
+                sorted_items = sorted(data["data"], key=lambda x: x["index"])
+                embeddings = [item["embedding"] for item in sorted_items]
+                return embeddings
+
+            except httpx.HTTPStatusError as e:
+                last_error = e
+                logger.warning(
+                    f"vLLM API HTTP error (attempt {attempt + 1}/{max_retries}): "
+                    f"{e.response.status_code} - {e.response.text}"
+                )
+            except httpx.RequestError as e:
+                last_error = e
+                logger.warning(
+                    f"vLLM API connection error (attempt {attempt + 1}/{max_retries}): {e}"
+                )
+                # connection pool 오염 시 client 재생성
+                try:
+                    await self._client.aclose()
+                except Exception:
+                    pass
+                headers = {}
+                if self._api_key:
+                    headers["Authorization"] = f"Bearer {self._api_key}"
+                self._client = httpx.AsyncClient(
+                    base_url=self._endpoint,
+                    timeout=httpx.Timeout(None),
+                    limits=httpx.Limits(
+                        max_keepalive_connections=5, max_connections=10
+                    ),
+                    headers=headers,
+                )
+            except Exception as e:
+                last_error = e
+                logger.error(f"Unexpected error calling vLLM API: {e}")
+                break
+
+        # 에러 메시지에 응답 본문 포함 (토큰 초과 등 상세 메시지 전달)
+        detail = str(last_error)
+        if hasattr(last_error, "response") and hasattr(last_error.response, "text"):
+            detail = f"{last_error} | body={last_error.response.text[:300]}"
+        raise Exception(f"vLLM embed failed after {max_retries} attempts: {detail}")
+
+    async def embed_texts(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate embeddings for a list of texts (documents/passages).
+
+        Args:
+            texts: List of text strings to embed
+
+        Returns:
+            List of embedding vectors (as lists of floats)
+        """
+        if not texts:
+            return []
+
+        try:
+            return await self._call_vllm_api(texts)
+        except Exception as e:
+            logger.error(f"Failed to generate embeddings via vLLM: {e}")
+            raise
+
+    async def embed_query(self, query: str) -> List[float]:
+        """
+        Generate embedding for a single query.
+
+        Args:
+            query: Query string
+
+        Returns:
+            Embedding vector as list of floats
+        """
+        if not query:
+            raise ValueError("Query cannot be empty")
+
+        try:
+            embeddings = await self._call_vllm_api([query], timeout=10.0)
+            return embeddings[0]
+        except Exception as e:
+            logger.error(f"Failed to generate query embedding via vLLM: {e}")
+            raise
+
+    async def embed_batch(
+        self, texts: List[str], batch_size: Optional[int] = None
+    ) -> List[List[float]]:
+        """
+        Generate embeddings with batching for large document sets.
+
+        Args:
+            texts: List of text strings to embed
+            batch_size: Override default batch size (for vLLM, can handle large batches)
+
+        Returns:
+            List of embedding vectors
+        """
+        if not texts:
+            return []
+
+        # vLLM can handle large batches efficiently
+        effective_batch_size = batch_size or 100
+        all_embeddings = []
+
+        for i in range(0, len(texts), effective_batch_size):
+            batch = texts[i : i + effective_batch_size]
+            embeddings = await self._call_vllm_api(batch)
+            all_embeddings.extend(embeddings)
+
+        return all_embeddings
+
+    def get_model_info(self) -> dict:
+        """Get information about the vLLM embedding service"""
+        return {
+            "backend": "vllm",
+            "endpoint": self._endpoint,
+            "model_name": self._model,
+            "dimension": self._dimension,
+        }
+
+    async def close(self):
+        """Close HTTP client connection"""
+        await self._client.aclose()
+
+
+# ============ Singleton Accessor ============
+
+_vllm_embedding_service: Optional[VLLMEmbeddingService] = None
+
+
+def get_vllm_embedding_service(
+    config: Optional["EmbeddingConfig"] = None,
+) -> VLLMEmbeddingService:
+    """
+    Get the singleton VLLMEmbeddingService instance.
+
+    Args:
+        config: Optional EmbeddingConfig (only used on first call)
+
+    Returns:
+        VLLMEmbeddingService singleton instance
+    """
+    global _vllm_embedding_service
+    if _vllm_embedding_service is None:
+        _vllm_embedding_service = VLLMEmbeddingService(config)
+    return _vllm_embedding_service
+
+
+def reset_vllm_embedding_service() -> None:
+    """
+    Reset the singleton instance (for testing purposes).
+    """
+    global _vllm_embedding_service
+    if _vllm_embedding_service is not None:
+        _vllm_embedding_service._initialized = False
+        _vllm_embedding_service = None
+    VLLMEmbeddingService._instance = None
+    VLLMEmbeddingService._initialized = False

codebuff_core/tests/__init__.py

@@ -0,0 +1 @@
+"""Codebuff Core - Tests"""