hdsp-jupyter-extension 2.0.7__py3-none-any.whl → 2.0.8__py3-none-any.whl

agent_server/core/embedding_service.py

@@ -4,13 +4,14 @@ Local Embedding Service - Wraps sentence-transformers for local embedding genera
 Features:
 - Zero external API calls (data sovereignty)
 - Lazy model loading (only when first needed)
-- Thread-safe singleton pattern
+- Thread-safe singleton pattern with async support
 - Configurable model and device
 - E5 model prefix handling for optimal performance
 
 Default model: intfloat/multilingual-e5-small (384 dimensions, Korean support)
 """
 
+import asyncio
 import logging
 from typing import TYPE_CHECKING, List, Optional
 
@@ -55,51 +56,59 @@ class EmbeddingService:
         self._model = None
         self._dimension: Optional[int] = None
         self._is_e5_model: bool = False
+        self._load_lock = asyncio.Lock()  # Thread-safe lazy loading
 
-    @property
-    def model(self):
-        """Lazy load the embedding model"""
-        if self._model is None:
-            self._load_model()
-        return self._model
-
-    def _load_model(self) -> None:
-        """Load the sentence-transformers model"""
-        try:
-            from sentence_transformers import SentenceTransformer
-        except ImportError:
-            raise ImportError(
-                "sentence-transformers is required for RAG. "
-                "Install with: pip install sentence-transformers"
-            )
-
-        model_name = self._config.get_model_name()
-        device = self._config.get_device()
-
-        logger.info(f"Loading embedding model: {model_name} on {device}")
-
-        try:
-            self._model = SentenceTransformer(
-                model_name, device=device, cache_folder=self._config.cache_folder
-            )
-            self._dimension = self._model.get_sentence_embedding_dimension()
-
-            # Check if E5 model (requires special prefix)
-            self._is_e5_model = "e5" in model_name.lower()
+    async def _ensure_model_loaded(self):
+        """Lazy load the embedding model (thread-safe, async)"""
+        if self._model is not None:
+            return
 
-            logger.info(
-                f"Embedding model loaded successfully. "
-                f"Dimension: {self._dimension}, E5 model: {self._is_e5_model}"
-            )
-        except Exception as e:
-            logger.error(f"Failed to load embedding model: {e}")
-            raise
+        async with self._load_lock:
+            # Double-check after acquiring lock
+            if self._model is not None:
+                return
+
+            try:
+                from sentence_transformers import SentenceTransformer
+            except ImportError:
+                raise ImportError(
+                    "sentence-transformers is required for RAG. "
+                    "Install with: pip install sentence-transformers"
+                )
+
+            model_name = self._config.get_model_name()
+            device = self._config.get_device()
+
+            logger.info(f"Loading embedding model: {model_name} on {device}")
+
+            try:
+                # Load model in separate thread to avoid blocking event loop
+                self._model = await asyncio.to_thread(
+                    SentenceTransformer,
+                    model_name,
+                    device=device,
+                    cache_folder=self._config.cache_folder,
+                )
+                self._dimension = self._model.get_sentence_embedding_dimension()
+
+                # Check if E5 model (requires special prefix)
+                self._is_e5_model = "e5" in model_name.lower()
+
+                logger.info(
+                    f"Embedding model loaded successfully. "
+                    f"Dimension: {self._dimension}, E5 model: {self._is_e5_model}"
+                )
+            except Exception as e:
+                logger.error(f"Failed to load embedding model: {e}")
+                raise
 
     @property
     def dimension(self) -> int:
-        """Get embedding dimension (triggers model load if needed)"""
+        """Get embedding dimension (must be loaded first)"""
         if self._dimension is None:
-            _ = self.model  # Trigger lazy load
+            raise RuntimeError(
+                "Embedding dimension not available. Model not loaded yet."
+            )
         return self._dimension
 
     def _prepare_texts(self, texts: List[str], is_query: bool = False) -> List[str]:
@@ -116,7 +125,7 @@ class EmbeddingService:
         prefix = "query: " if is_query else "passage: "
         return [prefix + text for text in texts]
 
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
+    async def embed_texts(self, texts: List[str]) -> List[List[float]]:
         """
         Generate embeddings for a list of texts (documents/passages).
 
@@ -129,11 +138,15 @@ class EmbeddingService:
         if not texts:
             return []
 
+        await self._ensure_model_loaded()
+
         # Prepare texts with prefix if E5 model
         prepared_texts = self._prepare_texts(texts, is_query=False)
 
         try:
-            embeddings = self.model.encode(
+            # Run in separate thread to avoid blocking event loop
+            embeddings = await asyncio.to_thread(
+                self._model.encode,
                 prepared_texts,
                 batch_size=self._config.batch_size,
                 show_progress_bar=len(texts) > 100,
@@ -145,7 +158,7 @@ class EmbeddingService:
             logger.error(f"Failed to generate embeddings: {e}")
             raise
 
-    def embed_query(self, query: str) -> List[float]:
+    async def embed_query(self, query: str) -> List[float]:
         """
         Generate embedding for a single query.
 
@@ -160,11 +173,15 @@ class EmbeddingService:
         if not query:
             raise ValueError("Query cannot be empty")
 
+        await self._ensure_model_loaded()
+
         # Prepare query with prefix if E5 model
         prepared_query = self._prepare_texts([query], is_query=True)[0]
 
         try:
-            embedding = self.model.encode(
+            # Run in separate thread to avoid blocking event loop
+            embedding = await asyncio.to_thread(
+                self._model.encode,
                 prepared_query,
                 convert_to_numpy=True,
                 normalize_embeddings=self._config.normalize_embeddings,
@@ -174,7 +191,7 @@ class EmbeddingService:
             logger.error(f"Failed to generate query embedding: {e}")
             raise
 
-    def embed_batch(
+    async def embed_batch(
         self, texts: List[str], batch_size: Optional[int] = None
     ) -> List[List[float]]:
         """
@@ -190,11 +207,15 @@ class EmbeddingService:
         if not texts:
             return []
 
+        await self._ensure_model_loaded()
+
         prepared_texts = self._prepare_texts(texts, is_query=False)
         effective_batch_size = batch_size or self._config.batch_size
 
         try:
-            embeddings = self.model.encode(
+            # Run in separate thread to avoid blocking event loop
+            embeddings = await asyncio.to_thread(
+                self._model.encode,
                 prepared_texts,
                 batch_size=effective_batch_size,
                 show_progress_bar=True,

agent_server/core/rag_manager.py

@@ -88,13 +88,24 @@ class RAGManager:
             self._client = self._create_qdrant_client()
             logger.info("Qdrant client initialized")
 
-            # 2. Initialize embedding service
-            from agent_server.core.embedding_service import get_embedding_service
+            # 2. Initialize embedding service (local or vLLM backend)
+            import os
+            embedding_backend = os.environ.get("HDSP_EMBEDDING_BACKEND", "local").lower()
 
-            self._embedding_service = get_embedding_service(self._config.embedding)
-            logger.info(
-                f"Embedding service initialized (dim={self._embedding_service.dimension})"
-            )
+            if embedding_backend == "vllm":
+                from agent_server.core.vllm_embedding_service import get_vllm_embedding_service
+                self._embedding_service = get_vllm_embedding_service(self._config.embedding)
+                logger.info(
+                    f"vLLM Embedding service initialized (dim={self._embedding_service.dimension})"
+                )
+            else:
+                from agent_server.core.embedding_service import get_embedding_service
+                self._embedding_service = get_embedding_service(self._config.embedding)
+                # Load model to get dimension
+                await self._embedding_service._ensure_model_loaded()
+                logger.info(
+                    f"Local Embedding service initialized (dim={self._embedding_service.dimension})"
+                )
 
             # 3. Ensure collection exists
             await self._ensure_collection()
@@ -151,26 +162,29 @@ class RAGManager:
             )
 
         cfg = self._config.qdrant
+        mode = cfg.get_mode()  # Use get_mode() for env override
 
-        if cfg.mode == "local":
+        if mode == "local":
             # Local file-based storage
             local_path = cfg.get_local_path()
             Path(local_path).mkdir(parents=True, exist_ok=True)
             logger.info(f"Initializing Qdrant in local mode: {local_path}")
             return QdrantClient(path=local_path)
 
-        elif cfg.mode == "server":
+        elif mode == "server":
             # Docker or external server
-            logger.info(f"Connecting to Qdrant server: {cfg.url}")
-            return QdrantClient(url=cfg.url)
+            url = cfg.get_url()  # Use get_url() for env override
+            logger.info(f"Connecting to Qdrant server: {url}")
+            return QdrantClient(url=url)
 
-        elif cfg.mode == "cloud":
+        elif mode == "cloud":
             # Qdrant Cloud
+            url = cfg.get_url()  # Use get_url() for env override
             logger.info("Connecting to Qdrant Cloud")
-            return QdrantClient(url=cfg.url, api_key=cfg.api_key)
+            return QdrantClient(url=url, api_key=cfg.api_key)
 
         else:
-            raise ValueError(f"Unknown Qdrant mode: {cfg.mode}")
+            raise ValueError(f"Unknown Qdrant mode: {mode}")
 
     async def _ensure_collection(self) -> None:
         """Create collection if it doesn't exist."""
@@ -274,7 +288,7 @@ class RAGManager:
                 )
 
                 if chunks:
-                    self._index_chunks(chunks, file_path)
+                    await self._index_chunks(chunks, file_path)
                     indexed += 1
                     self._index_stats["total_documents"] += 1
                     self._index_stats["total_chunks"] += len(chunks)
@@ -345,13 +359,13 @@ class RAGManager:
         else:
             return "general"
 
-    def _index_chunks(self, chunks: List[Dict], file_path: Path) -> None:
+    async def _index_chunks(self, chunks: List[Dict], file_path: Path) -> None:
         """Index document chunks to Qdrant."""
         from qdrant_client.models import PointStruct
 
         # Generate embeddings
         texts = [c["content"] for c in chunks]
-        embeddings = self._embedding_service.embed_texts(texts)
+        embeddings = await self._embedding_service.embed_texts(texts)
 
         # Add content hash to all chunks
         file_hash = self._compute_file_hash(file_path)
@@ -430,7 +444,7 @@ class RAGManager:
             )
 
             if chunks:
-                self._index_chunks(chunks, file_path)
+                await self._index_chunks(chunks, file_path)
                 logger.info(f"Reindexed: {file_path}")
         except Exception as e:
             logger.error(f"Failed to reindex {file_path}: {e}")

agent_server/core/retriever.py

@@ -84,21 +84,23 @@ class Retriever:
         effective_threshold = score_threshold or self._config.score_threshold
 
         # Generate query embedding
-        query_embedding = self._embedding_service.embed_query(query)
+        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:
-            results = self._client.search(
+            response = self._client.query_points(
                 collection_name=self._config.qdrant.collection_name,
-                query_vector=query_embedding,
+                query=query_embedding,
                 query_filter=qdrant_filter,
                 limit=effective_top_k,
-                score_threshold=effective_threshold
-                * 0.5,  # Lower for initial retrieval
+                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 []
@@ -193,7 +195,7 @@ class Retriever:
         effective_threshold = score_threshold or self._config.score_threshold
 
         # Generate query embedding
-        query_embedding = self._embedding_service.embed_query(query)
+        query_embedding = await self._embedding_service.embed_query(query)
 
         # Build filter condition
         qdrant_filter = self._build_filter(filters) if filters else None
@@ -201,13 +203,16 @@ class Retriever:
         # Vector search with timing
         try:
             # 디버그용으로 더 많은 결과 (3배)를 낮은 threshold로 가져옴
-            results = self._client.search(
+            response = self._client.query_points(
                 collection_name=self._config.qdrant.collection_name,
-                query_vector=query_embedding,
+                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(

agent_server/core/vllm_embedding_service.py

@@ -0,0 +1,243 @@
+"""
+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
+import os
+from typing import TYPE_CHECKING, List, Optional
+
+import httpx
+import time
+
+if TYPE_CHECKING:
+    from hdsp_agent_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 hdsp_agent_core.models.rag import EmbeddingConfig
+
+        self._config = config or EmbeddingConfig()
+
+        # vLLM configuration from environment variables
+        self._endpoint = os.environ.get("HDSP_VLLM_ENDPOINT", "http://localhost:8000")
+        self._model = os.environ.get("HDSP_VLLM_MODEL", "qwen3-embedding-8b")
+        self._dimension = int(os.environ.get("HDSP_VLLM_DIMENSION", "8192"))
+
+        # HTTP client with retry
+        self._client = httpx.AsyncClient(
+            base_url=self._endpoint,
+            timeout=httpx.Timeout(30.0),
+            limits=httpx.Limits(max_keepalive_connections=5, max_connections=10)
+        )
+
+        logger.info(
+            f"vLLM Embedding Service initialized: "
+            f"endpoint={self._endpoint}, model={self._model}, dim={self._dimension}"
+        )
+
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension"""
+        return self._dimension
+
+    async def _call_vllm_api(self, texts: List[str], max_retries: int = 3) -> 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,
+        }
+
+        last_error = None
+        for attempt in range(max_retries):
+            try:
+                response = await self._client.post("/v1/embeddings", json=payload)
+                response.raise_for_status()
+
+                data = response.json()
+                # 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}"
+                )
+            except Exception as e:
+                last_error = e
+                logger.error(f"Unexpected error calling vLLM API: {e}")
+                break
+
+        raise Exception(f"Failed to connect to vLLM after {max_retries} attempts: {last_error}")
+
+    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])
+            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

agent_server/langchain/agent.py

@@ -23,12 +23,16 @@ from agent_server.langchain.prompts import (
 )
 from agent_server.langchain.tools import (
     check_resource_tool,
+    diagnostics_tool,
+    edit_file_tool,
     execute_command_tool,
     final_answer_tool,
     jupyter_cell_tool,
     list_files_tool,
     markdown_tool,
+    multiedit_file_tool,
     read_file_tool,
+    references_tool,
     search_notebook_cells_tool,
     search_workspace_tool,
     write_file_tool,
@@ -45,11 +49,15 @@ def _get_all_tools():
         final_answer_tool,
         read_file_tool,
         write_file_tool,
+        edit_file_tool,
+        multiedit_file_tool,
         list_files_tool,
         search_workspace_tool,
         search_notebook_cells_tool,
         execute_command_tool,
         check_resource_tool,
+        diagnostics_tool,
+        references_tool,
     ]