| import logging |
| from pathlib import Path |
| from typing import Optional, List, Dict, Any |
| from tqdm import tqdm |
|
|
| from vectorDB.embedder import HFEmbedder |
| from vectorDB.semantic_index import SemanticIndex |
| from vectorDB.ast_index import ASTIndex |
| from vectorDB.hybrid_retriever import HybridRetriever |
|
|
| logger = logging.getLogger(__name__) |
|
|
| class IndexManager: |
| """ |
| Manages the complete indexing pipeline: embedding, semantic indexing, AST indexing, |
| and hybrid retrieval. Handles save/load and auto-detection. |
| """ |
| |
| def __init__( |
| self, |
| embedding_model: str = "BAAI/bge-m3", |
| language: str = "sql", |
| semantic_weight: float = 0.7, |
| ast_weight: float = 0.3, |
| save_dir: str = "data/indices" |
| ): |
| self.embedding_model = embedding_model |
| self.language = language |
| self.semantic_weight = semantic_weight |
| self.ast_weight = ast_weight |
| self.save_dir = Path(save_dir) |
| |
| self.embedder = None |
| self.semantic_index = None |
| self.ast_index = None |
| self.hybrid_retriever = None |
| |
| logger.info(f"IndexManager initialized: model={embedding_model}, lang={language}") |
| |
| def build_indices( |
| self, |
| documents: List[Dict[str, Any]], |
| text_field: str = "text", |
| code_field: Optional[str] = None, |
| batch_size: int = 32, |
| force_rebuild: bool = False |
| ): |
| """ |
| Build both semantic and AST indices from documents. |
| |
| Args: |
| documents: List of document dicts |
| text_field: Field name for natural language text |
| code_field: Field name for code/SQL (if different from text_field) |
| batch_size: Batch size for embedding |
| force_rebuild: Whether to rebuild even if indices exist |
| """ |
| |
| if not force_rebuild and self._indices_exist(): |
| logger.info("Existing indices detected. Loading...") |
| self.load_indices() |
| return |
| |
| logger.info(f"Building indices for {len(documents)} documents...") |
| |
| |
| self.embedder = HFEmbedder(model_name=self.embedding_model) |
| |
| |
| texts = [doc.get(text_field, "") for doc in documents] |
| codes = [doc.get(code_field, doc.get(text_field, "")) for doc in documents] if code_field else texts |
| |
| |
| logger.info("\n=== Building Semantic Index ===") |
| embeddings = self.embedder.encode( |
| texts, |
| batch_size=batch_size, |
| show_progress=True |
| ) |
| |
| embedding_dim = embeddings.shape[1] |
| self.semantic_index = SemanticIndex(embedding_dim=embedding_dim) |
| self.semantic_index.add(embeddings, documents, show_progress=True) |
| |
| |
| logger.info("\n=== Building AST Index ===") |
| self.ast_index = ASTIndex(language=self.language) |
| self.ast_index.add(codes, documents, show_progress=True) |
| |
| |
| self.hybrid_retriever = HybridRetriever( |
| semantic_index=self.semantic_index, |
| ast_index=self.ast_index, |
| embedder=self.embedder, |
| semantic_weight=self.semantic_weight, |
| ast_weight=self.ast_weight |
| ) |
| |
| |
| self.save_indices() |
| |
| logger.info("\n✅ All indices built and saved successfully!") |
| |
| def _indices_exist(self) -> bool: |
| """Check if all index files exist.""" |
| semantic_path = self.save_dir / "semantic_index.faiss" |
| semantic_meta = self.save_dir / "semantic_index.metadata.pkl" |
| ast_path = self.save_dir / "ast_index.pkl" |
| |
| return semantic_path.exists() and semantic_meta.exists() and ast_path.exists() |
| |
| def save_indices(self): |
| """Save all indices to disk.""" |
| self.save_dir.mkdir(parents=True, exist_ok=True) |
| |
| if self.semantic_index: |
| self.semantic_index.save(self.save_dir / "semantic_index") |
| |
| if self.ast_index: |
| self.ast_index.save(self.save_dir / "ast_index.pkl") |
| |
| logger.info(f"✅ Indices saved to {self.save_dir}") |
| |
| def load_indices(self): |
| """Load all indices from disk.""" |
| if not self._indices_exist(): |
| raise FileNotFoundError(f"No indices found at {self.save_dir}") |
| |
| |
| self.embedder = HFEmbedder(model_name=self.embedding_model) |
| |
| |
| self.semantic_index = SemanticIndex.load(self.save_dir / "semantic_index") |
| |
| |
| self.ast_index = ASTIndex.load(self.save_dir / "ast_index.pkl") |
| |
| |
| self.hybrid_retriever = HybridRetriever( |
| semantic_index=self.semantic_index, |
| ast_index=self.ast_index, |
| embedder=self.embedder, |
| semantic_weight=self.semantic_weight, |
| ast_weight=self.ast_weight |
| ) |
| |
| logger.info("✅ All indices loaded successfully!") |
| |
| def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]: |
| """ |
| Search using hybrid retriever. |
| |
| Args: |
| query: Query string |
| k: Number of results |
| |
| Returns: |
| List of metadata dicts |
| """ |
| if self.hybrid_retriever is None: |
| raise RuntimeError("Indices not initialized. Call build_indices() or load_indices() first.") |
| |
| results = self.hybrid_retriever.search(query, k=k) |
| return [metadata for metadata, score in results] |
