""" Semantic Search Module - Vector Search + BM25 + Hybrid Search Mendukung: - Vector Embeddings dengan berbagai model - BM25 Full-text Search - Hybrid Search (kombinasi keduanya) - Knowledge Base Management """ import os import json import pickle import numpy as np from typing import List, Dict, Tuple, Optional, Any from pathlib import Path from dataclasses import dataclass, field from collections import Counter import math try: from sklearn.feature_extraction.text import TfidfVectorizer HAS_SKLEARN = True except ImportError: HAS_SKLEARN = False try: import gensim.downloader as api from gensim.models import KeyedVectors HAS_GENSIM = True except ImportError: HAS_GENSIM = False try: from sentence_transformers import SentenceTransformer HAS_SENTENCE_TRANSFORMERS = True except ImportError: HAS_SENTENCE_TRANSFORMERS = False from config import CACHE_DIR, DATA_DIR from utils import console, Theme, debug_logger, error_logger @dataclass class SearchResult: """Hasil pencarian dengan score dan metadata""" text: str score: float source: str = "unknown" type: str = "text" metadata: Dict[str, Any] = field(default_factory=dict) index: int = -1 class BM25Search: """ BM25 (Best Matching 25) - Lexical/Full-Text Search Lebih baik untuk exact term matching dan keyword search """ def __init__(self, documents: List[str] = None, k1: float = 1.5, b: float = 0.75): self.k1 = k1 # Term frequency saturation parameter self.b = b # Length normalization parameter self.documents = documents or [] self.inverted_index: Dict[str, List[int]] = {} self.doc_lengths: List[int] = [] self.idf: Dict[str, float] = {} self.avg_doc_length = 0.0 if self.documents: self._build_index() def _tokenize(self, text: str) -> List[str]: """Tokenize text dengan cleaning""" text = text.lower() # Hapus karakter spesial tapi pertahankan spasi import re tokens = re.findall(r'\b\w+\b', text) # Hapus stopwords umum stopwords = {'dan', 'atau', 'yang', 'di', 'ke', 'dari', 'untuk', 'dengan', 'adalah', 'pada', 'dalam', 'atas', 'oleh', 'sebagai'} return [t for t in tokens if t not in stopwords and len(t) > 2] def _build_index(self) -> None: """Build inverted index dari documents""" self.inverted_index = {} self.doc_lengths = [] for doc_idx, doc in enumerate(self.documents): tokens = self._tokenize(doc) self.doc_lengths.append(len(tokens)) seen = set() for token in tokens: if token not in seen: if token not in self.inverted_index: self.inverted_index[token] = [] self.inverted_index[token].append(doc_idx) seen.add(token) self.avg_doc_length = sum(self.doc_lengths) / len(self.doc_lengths) if self.doc_lengths else 0 # Calculate IDF n_docs = len(self.documents) for token, doc_indices in self.inverted_index.items(): self.idf[token] = math.log((n_docs - len(doc_indices) + 0.5) / (len(doc_indices) + 0.5) + 1) def add_documents(self, documents: List[str]) -> None: """Tambah documents dan rebuild index""" self.documents.extend(documents) self._build_index() def search(self, query: str, top_k: int = 5) -> List[SearchResult]: """Search menggunakan BM25""" query_tokens = self._tokenize(query) if not query_tokens or not self.documents: return [] scores = [0.0] * len(self.documents) for token in query_tokens: if token not in self.idf: continue idf_score = self.idf[token] for doc_idx in self.inverted_index.get(token, []): # Count token frequency dalam document tf = self._tokenize(self.documents[doc_idx]).count(token) # BM25 formula numerator = idf_score * tf * (self.k1 + 1) denominator = tf + self.k1 * ( 1 - self.b + self.b * (self.doc_lengths[doc_idx] / self.avg_doc_length) ) scores[doc_idx] += numerator / denominator # Get top-k results top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k] results = [] for idx in top_indices: if scores[idx] > 0: results.append(SearchResult( text=self.documents[idx], score=scores[idx], source="bm25", type="keyword_match", index=idx )) return results class VectorSearch: """ Vector Search - Semantic Search menggunakan embeddings Lebih baik untuk semantic similarity dan meaning-based search """ def __init__(self, model_name: str = "all-MiniLM-L6-v2", use_gpu: bool = False): self.model_name = model_name self.model = None self.embeddings = None self.documents = [] self.use_gpu = use_gpu and HAS_SENTENCE_TRANSFORMERS self._load_model() def _load_model(self) -> None: """Load embedding model""" if not HAS_SENTENCE_TRANSFORMERS: console.print(Theme.warning("sentence-transformers tidak terinstall. Skipping vector search.")) return try: self.model = SentenceTransformer(self.model_name) if self.use_gpu: self.model.to('cuda') console.print(Theme.success(f"Embedding model loaded: {self.model_name}")) except Exception as e: console.print(Theme.error(f"Gagal load embedding model: {e}")) error_logger.error(f"Embedding model error: {e}") def add_documents(self, documents: List[str]) -> None: """Add documents dan generate embeddings""" if self.model is None: console.print(Theme.warning("Model not available")) return try: self.documents = documents self.embeddings = self.model.encode(documents, convert_to_numpy=True) console.print(Theme.success(f"Indexed {len(documents)} documents")) except Exception as e: error_logger.error(f"Embedding error: {e}") def search(self, query: str, top_k: int = 5) -> List[SearchResult]: """Search dengan semantic similarity""" if self.model is None or self.embeddings is None or len(self.documents) == 0: return [] try: query_embedding = self.model.encode([query], convert_to_numpy=True)[0] # Compute cosine similarity from sklearn.metrics.pairwise import cosine_similarity similarities = cosine_similarity([query_embedding], self.embeddings)[0] # Get top-k top_indices = np.argsort(similarities)[::-1][:top_k] results = [] for idx in top_indices: score = float(similarities[idx]) if score > 0: # Only positive scores results.append(SearchResult( text=self.documents[idx], score=score, source="vector", type="semantic", index=idx )) return results except Exception as e: error_logger.error(f"Search error: {e}") return [] class HybridSearch: """ Hybrid Search - Kombinasi BM25 + Vector Search Menggabungkan strengths dari lexical dan semantic search """ def __init__(self, embedding_model: str = "all-MiniLM-L6-v2"): self.bm25 = BM25Search() self.vector_search = VectorSearch(embedding_model) self.documents = [] self.bm25_weight = 0.4 # Weight untuk BM25 self.vector_weight = 0.6 # Weight untuk Vector Search def add_documents(self, documents: List[str]) -> None: """Add documents ke both search engines""" self.documents = documents self.bm25.add_documents(documents) self.vector_search.add_documents(documents) def search(self, query: str, top_k: int = 5) -> List[SearchResult]: """Hybrid search combining both methods""" bm25_results = self.bm25.search(query, top_k=top_k) vector_results = self.vector_search.search(query, top_k=top_k) # Combine scores combined_scores = {} for result in bm25_results: idx = result.index normalized_score = result.score / (max(r.score for r in bm25_results) or 1) combined_scores[idx] = combined_scores.get(idx, 0) + normalized_score * self.bm25_weight for result in vector_results: idx = result.index normalized_score = result.score # Already normalized [0, 1] combined_scores[idx] = combined_scores.get(idx, 0) + normalized_score * self.vector_weight # Sort and return top-k sorted_indices = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)[:top_k] results = [] for idx, score in sorted_indices: results.append(SearchResult( text=self.documents[idx], score=score, source="hybrid", type="combined", index=idx )) return results def set_weights(self, bm25_weight: float, vector_weight: float) -> None: """Adjust weights antara BM25 dan Vector search""" total = bm25_weight + vector_weight self.bm25_weight = bm25_weight / total self.vector_weight = vector_weight / total console.print(Theme.info(f"Weights set: BM25={self.bm25_weight:.2f}, Vector={self.vector_weight:.2f}")) class KnowledgeBase: """ Knowledge Base Manager - Manajemen document repository dengan search capabilities """ def __init__(self, kb_dir: str = None, embedding_model: str = "all-MiniLM-L6-v2"): self.kb_dir = Path(kb_dir or os.path.join(DATA_DIR, "knowledge_base")) self.kb_dir.mkdir(exist_ok=True) self.documents: Dict[str, Dict[str, Any]] = {} self.hybrid_search = HybridSearch(embedding_model) self.metadata_file = self.kb_dir / "metadata.json" self.embeddings_file = self.kb_dir / "embeddings.pkl" self._load_knowledge_base() def add_document(self, doc_id: str, text: str, metadata: Dict[str, Any] = None) -> None: """Add document dengan metadata""" self.documents[doc_id] = { 'text': text, 'metadata': metadata or {}, 'added_at': __import__('time').time() } self._rebuild_search_index() self._save_knowledge_base() def add_documents(self, documents: List[Dict[str, Any]]) -> None: """Add multiple documents""" for doc in documents: doc_id = doc.get('id', f"doc_{len(self.documents)}") text = doc.get('text', '') metadata = doc.get('metadata', {}) self.add_document(doc_id, text, metadata) def search(self, query: str, top_k: int = 5, method: str = "hybrid") -> List[Dict[str, Any]]: """Search dalam knowledge base""" if method == "hybrid": results = self.hybrid_search.search(query, top_k) elif method == "vector": results = self.hybrid_search.vector_search.search(query, top_k) elif method == "bm25": results = self.hybrid_search.bm25.search(query, top_k) else: results = [] enriched_results = [] for result in results: # Find corresponding document for doc_id, doc_data in self.documents.items(): if doc_data['text'] == result.text: enriched_results.append({ 'id': doc_id, 'text': result.text, 'score': result.score, 'source': result.source, 'type': result.type, 'metadata': doc_data['metadata'] }) break return enriched_results def _rebuild_search_index(self) -> None: """Rebuild search indexes dari all documents""" texts = [doc['text'] for doc in self.documents.values()] if texts: self.hybrid_search.add_documents(texts) def _save_knowledge_base(self) -> None: """Save KB to disk""" try: with open(self.metadata_file, 'w', encoding='utf-8') as f: json.dump(self.documents, f, ensure_ascii=False, indent=2, default=str) except Exception as e: error_logger.error(f"KB save error: {e}") def _load_knowledge_base(self) -> None: """Load KB from disk""" try: if self.metadata_file.exists(): with open(self.metadata_file, 'r', encoding='utf-8') as f: self.documents = json.load(f) self._rebuild_search_index() console.print(Theme.success(f"Loaded {len(self.documents)} documents from KB")) except Exception as e: error_logger.error(f"KB load error: {e}") def delete_document(self, doc_id: str) -> bool: """Delete document dari KB""" if doc_id in self.documents: del self.documents[doc_id] self._rebuild_search_index() self._save_knowledge_base() return True return False def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]: """Get specific document""" return self.documents.get(doc_id) def list_documents(self) -> List[Dict[str, Any]]: """List semua documents dengan metadata""" return [ {'id': doc_id, **doc_data} for doc_id, doc_data in self.documents.items() ] def get_stats(self) -> Dict[str, Any]: """Get KB statistics""" return { 'total_documents': len(self.documents), 'total_text_length': sum(len(doc['text']) for doc in self.documents.values()), 'average_doc_length': np.mean([len(doc['text']) for doc in self.documents.values()]) if self.documents else 0, }