hanuscode 1.0.0__py3-none-any.whl

hanus/search/semantic.py

@@ -0,0 +1,596 @@
+# hanus/search/semantic.py
+"""
+Búsqueda semántica de código usando embeddings.
+
+Permite buscar código por concepto, no solo por texto exacto.
+"""
+from __future__ import annotations
+import json
+import re
+import hashlib
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Any, Tuple
+from collections import defaultdict
+import os
+
+
+@dataclass
+class CodeChunk:
+    """Un fragmento de código indexable."""
+    id: str
+    file_path: str
+    content: str
+    start_line: int
+    end_line: int
+    chunk_type: str  # function, class, module, etc.
+    name: Optional[str] = None
+    embedding: Optional[List[float]] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict:
+        return {
+            "id": self.id,
+            "file_path": self.file_path,
+            "content": self.content[:500],  # Truncar para serialización
+            "start_line": self.start_line,
+            "end_line": self.end_line,
+            "chunk_type": self.chunk_type,
+            "name": self.name,
+            "metadata": self.metadata,
+        }
+
+
+@dataclass
+class SearchResult:
+    """Resultado de búsqueda."""
+    chunk: CodeChunk
+    score: float
+    match_type: str  # "semantic", "keyword", "exact"
+    highlights: List[str] = field(default_factory=list)
+
+    def to_dict(self) -> Dict:
+        return {
+            "file_path": self.chunk.file_path,
+            "name": self.chunk.name,
+            "content": self.chunk.content[:300],
+            "start_line": self.chunk.start_line,
+            "end_line": self.chunk.end_line,
+            "chunk_type": self.chunk.chunk_type,
+            "score": self.score,
+            "match_type": self.match_type,
+            "highlights": self.highlights,
+        }
+
+
+class SemanticSearch:
+    """
+    Búsqueda semántica de código.
+
+    Features:
+    - Búsqueda por concepto usando embeddings
+    - Fallback a búsqueda fuzzy si no hay embeddings
+    - Índice de código por proyecto
+    - Cache de resultados
+    """
+
+    # Extensiones soportadas
+    CODE_EXTENSIONS = {'.py', '.js', '.ts', '.jsx', '.tsx', '.java', '.go', '.rs', '.c', '.cpp', '.h', '.hpp', '.rb', '.php'}
+
+    def __init__(
+        self,
+        project_root: Path,
+        use_embeddings: bool = True,
+        index_dir: Optional[Path] = None
+    ):
+        self.project_root = project_root
+        self.use_embeddings = use_embeddings
+        self.index_dir = index_dir or (Path.home() / ".hanus" / "search_index")
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+
+        # Índice de chunks
+        self._chunks: Dict[str, CodeChunk] = {}
+        self._file_chunks: Dict[str, List[str]] = defaultdict(list)  # file -> chunk_ids
+
+        # Modelo de embeddings (lazy load)
+        self._model = None
+        self._embedding_dim = 384  # Default para sentence-transformers
+
+        # Índice invertido para búsqueda keyword
+        self._keyword_index: Dict[str, List[str]] = defaultdict(list)  # word -> chunk_ids
+
+        # Cache de embeddings
+        self._embedding_cache: Dict[str, List[float]] = {}
+
+        # Cargar índice existente
+        self._load_index()
+
+        # Intentar cargar modelo de embeddings
+        if use_embeddings:
+            self._init_embedding_model()
+
+    def index_project(self, force_reindex: bool = False) -> Dict[str, int]:
+        """
+        Indexa todo el proyecto.
+
+        Args:
+            force_reindex: Forzar reindexación completa
+
+        Returns:
+            Estadísticas de indexación
+        """
+        stats = {"files_indexed": 0, "chunks_created": 0, "errors": 0}
+
+        for ext in self.CODE_EXTENSIONS:
+            for file_path in self.project_root.rglob(f"*{ext}"):
+                # Skip directorios comunes
+                if any(skip in str(file_path) for skip in ['node_modules', '__pycache__', '.git', 'venv', 'env', 'dist', 'build']):
+                    continue
+
+                try:
+                    chunks = self._index_file(file_path)
+                    stats["files_indexed"] += 1
+                    stats["chunks_created"] += len(chunks)
+                except Exception as e:
+                    stats["errors"] += 1
+
+        # Guardar índice
+        self._save_index()
+
+        return stats
+
+    def search(
+        self,
+        query: str,
+        limit: int = 10,
+        min_score: float = 0.3,
+        file_filter: Optional[str] = None
+    ) -> List[SearchResult]:
+        """
+        Busca código por concepto o keyword.
+
+        Args:
+            query: Consulta de búsqueda
+            limit: Máximo de resultados
+            min_score: Score mínimo
+            file_filter: Filtro de archivo (glob pattern)
+
+        Returns:
+            Lista de resultados
+        """
+        results = []
+
+        # 1. Búsqueda semántica si hay embeddings
+        if self.use_embeddings and self._model:
+            semantic_results = self._semantic_search(query, limit * 2)
+            results.extend(semantic_results)
+
+        # 2. Búsqueda keyword/fuzzy
+        keyword_results = self._keyword_search(query, limit * 2)
+        results.extend(keyword_results)
+
+        # 3. Combinar y deduplicar
+        seen_chunks = set()
+        combined = []
+        for result in results:
+            if result.chunk.id not in seen_chunks:
+                seen_chunks.add(result.chunk.id)
+                combined.append(result)
+
+        # 4. Filtrar por archivo
+        if file_filter:
+            combined = [r for r in combined if self._matches_filter(r.chunk.file_path, file_filter)]
+
+        # 5. Ordenar por score y limitar
+        combined.sort(key=lambda x: -x.score)
+        combined = [r for r in combined if r.score >= min_score]
+
+        return combined[:limit]
+
+    def search_similar(
+        self,
+        chunk_id: str,
+        limit: int = 5
+    ) -> List[SearchResult]:
+        """
+        Busca código similar a un chunk específico.
+
+        Args:
+            chunk_id: ID del chunk base
+            limit: Máximo de resultados
+
+        Returns:
+            Lista de resultados similares
+        """
+        if chunk_id not in self._chunks:
+            return []
+
+        base_chunk = self._chunks[chunk_id]
+
+        if not base_chunk.embedding:
+            return []
+
+        # Buscar por similitud de embedding
+        results = []
+        for cid, chunk in self._chunks.items():
+            if cid == chunk_id:
+                continue
+
+            if chunk.embedding:
+                similarity = self._cosine_similarity(base_chunk.embedding, chunk.embedding)
+                if similarity > 0.5:
+                    results.append(SearchResult(
+                        chunk=chunk,
+                        score=similarity,
+                        match_type="semantic"
+                    ))
+
+        results.sort(key=lambda x: -x.score)
+        return results[:limit]
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Obtiene estadísticas del índice."""
+        chunks_with_embeddings = sum(1 for c in self._chunks.values() if c.embedding is not None)
+
+        return {
+            "total_chunks": len(self._chunks),
+            "files_indexed": len(self._file_chunks),
+            "chunks_with_embeddings": chunks_with_embeddings,
+            "embedding_model": "sentence-transformers" if self._model else "none",
+            "keyword_index_size": len(self._keyword_index),
+        }
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # MÉTODOS PRIVADOS
+    # ══════════════════════════════════════════════════════════════════════════
+
+    def _init_embedding_model(self):
+        """Inicializa el modelo de embeddings."""
+        try:
+            from sentence_transformers import SentenceTransformer
+            self._model = SentenceTransformer('all-MiniLM-L6-v2')
+            self._embedding_dim = 384
+        except ImportError:
+            print("[SemanticSearch] sentence-transformers not available, using keyword search only")
+            self._model = None
+        except Exception as e:
+            print(f"[SemanticSearch] Error loading model: {e}")
+            self._model = None
+
+    def _index_file(self, file_path: Path) -> List[CodeChunk]:
+        """Indexa un archivo."""
+        chunks = []
+
+        try:
+            content = file_path.read_text(encoding="utf-8", errors="replace")
+        except Exception:
+            return chunks
+
+        # Dividir en chunks por función/clase
+        code_chunks = self._split_into_chunks(content, str(file_path))
+
+        for chunk_data in code_chunks:
+            chunk_id = self._generate_chunk_id(file_path, chunk_data)
+
+            # Crear embedding si el modelo está disponible
+            embedding = None
+            if self._model:
+                embedding = self._get_embedding(chunk_data["content"])
+
+            chunk = CodeChunk(
+                id=chunk_id,
+                file_path=str(file_path.relative_to(self.project_root)),
+                content=chunk_data["content"],
+                start_line=chunk_data["start_line"],
+                end_line=chunk_data["end_line"],
+                chunk_type=chunk_data["type"],
+                name=chunk_data.get("name"),
+                embedding=embedding,
+            )
+
+            self._chunks[chunk_id] = chunk
+            self._file_chunks[str(file_path)].append(chunk_id)
+            chunks.append(chunk)
+
+            # Actualizar índice keyword
+            self._update_keyword_index(chunk)
+
+        return chunks
+
+    def _split_into_chunks(self, content: str, file_path: str) -> List[Dict]:
+        """Divide el código en chunks indexables."""
+        chunks = []
+        lines = content.split('\n')
+        ext = Path(file_path).suffix.lower()
+
+        if ext == '.py':
+            chunks = self._split_python(content, lines)
+        elif ext in ('.js', '.ts', '.jsx', '.tsx'):
+            chunks = self._split_javascript(content, lines)
+        else:
+            # Chunk por líneas para otros lenguajes
+            chunks = self._split_by_lines(content, lines)
+
+        return chunks
+
+    def _split_python(self, content: str, lines: List[str]) -> List[Dict]:
+        """Divide código Python en chunks."""
+        chunks = []
+
+        try:
+            import ast
+            tree = ast.parse(content)
+
+            for node in ast.iter_child_nodes(tree):
+                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                    chunks.append({
+                        "content": "\n".join(lines[node.lineno - 1:node.end_lineno]),
+                        "start_line": node.lineno,
+                        "end_line": node.end_lineno,
+                        "type": "function",
+                        "name": node.name,
+                    })
+                elif isinstance(node, ast.ClassDef):
+                    # Clase completa
+                    class_content = "\n".join(lines[node.lineno - 1:node.end_lineno])
+                    chunks.append({
+                        "content": class_content,
+                        "start_line": node.lineno,
+                        "end_line": node.end_lineno,
+                        "type": "class",
+                        "name": node.name,
+                    })
+
+                    # Métodos individuales
+                    for item in node.body:
+                        if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                            chunks.append({
+                                "content": "\n".join(lines[item.lineno - 1:item.end_lineno]),
+                                "start_line": item.lineno,
+                                "end_line": item.end_lineno,
+                                "type": "method",
+                                "name": f"{node.name}.{item.name}",
+                            })
+
+        except SyntaxError:
+            # Fallback a división por líneas
+            return self._split_by_lines(content, lines)
+
+        # Si no hay chunks, crear uno para el archivo completo
+        if not chunks:
+            chunks.append({
+                "content": content,
+                "start_line": 1,
+                "end_line": len(lines),
+                "type": "module",
+            })
+
+        return chunks
+
+    def _split_javascript(self, content: str, lines: List[str]) -> List[Dict]:
+        """Divide código JavaScript en chunks."""
+        chunks = []
+
+        # Funciones
+        func_pattern = r'(?:export\s+)?(?:async\s+)?function\s+(\w+)\s*\([^)]*\)\s*\{'
+        for match in re.finditer(func_pattern, content):
+            name = match.group(1)
+            start = content[:match.start()].count('\n') + 1
+            # Simplificado: tomar ~50 líneas
+            end = min(start + 50, len(lines))
+            chunks.append({
+                "content": "\n".join(lines[start - 1:end]),
+                "start_line": start,
+                "end_line": end,
+                "type": "function",
+                "name": name,
+            })
+
+        # Clases
+        class_pattern = r'(?:export\s+)?class\s+(\w+)'
+        for match in re.finditer(class_pattern, content):
+            name = match.group(1)
+            start = content[:match.start()].count('\n') + 1
+            end = min(start + 100, len(lines))
+            chunks.append({
+                "content": "\n".join(lines[start - 1:end]),
+                "start_line": start,
+                "end_line": end,
+                "type": "class",
+                "name": name,
+            })
+
+        if not chunks:
+            return self._split_by_lines(content, lines)
+
+        return chunks
+
+    def _split_by_lines(self, content: str, lines: List[str], chunk_size: int = 50) -> List[Dict]:
+        """Divide por líneas fijas."""
+        chunks = []
+        for i in range(0, len(lines), chunk_size):
+            chunk_lines = lines[i:i + chunk_size]
+            chunks.append({
+                "content": "\n".join(chunk_lines),
+                "start_line": i + 1,
+                "end_line": min(i + chunk_size, len(lines)),
+                "type": "block",
+            })
+        return chunks
+
+    def _generate_chunk_id(self, file_path: Path, chunk_data: Dict) -> str:
+        """Genera ID único para un chunk."""
+        content_hash = hashlib.md5(chunk_data["content"].encode()).hexdigest()[:8]
+        name = chunk_data.get("name", f"L{chunk_data['start_line']}")
+        return f"{file_path.stem}:{name}:{content_hash}"
+
+    def _get_embedding(self, text: str) -> Optional[List[float]]:
+        """Obtiene embedding para un texto."""
+        if not self._model:
+            return None
+
+        # Verificar cache
+        text_hash = hashlib.md5(text.encode()).hexdigest()
+        if text_hash in self._embedding_cache:
+            return self._embedding_cache[text_hash]
+
+        try:
+            embedding = self._model.encode(text[:500]).tolist()  # Limitar texto
+            self._embedding_cache[text_hash] = embedding
+            return embedding
+        except Exception:
+            return None
+
+    def _update_keyword_index(self, chunk: CodeChunk):
+        """Actualiza índice invertido de keywords."""
+        # Tokenizar contenido
+        words = re.findall(r'\b\w+\b', chunk.content.lower())
+        words = [w for w in words if len(w) > 2]  # Ignorar palabras muy cortas
+
+        for word in set(words):
+            if chunk.id not in self._keyword_index[word]:
+                self._keyword_index[word].append(chunk.id)
+
+    def _semantic_search(self, query: str, limit: int) -> List[SearchResult]:
+        """Búsqueda semántica usando embeddings."""
+        if not self._model:
+            return []
+
+        query_embedding = self._get_embedding(query)
+        if not query_embedding:
+            return []
+
+        results = []
+        for chunk_id, chunk in self._chunks.items():
+            if chunk.embedding:
+                similarity = self._cosine_similarity(query_embedding, chunk.embedding)
+                if similarity > 0.2:
+                    results.append(SearchResult(
+                        chunk=chunk,
+                        score=similarity,
+                        match_type="semantic",
+                    ))
+
+        results.sort(key=lambda x: -x.score)
+        return results[:limit]
+
+    def _keyword_search(self, query: str, limit: int) -> List[SearchResult]:
+        """Búsqueda por keywords con fuzzy matching."""
+        query_words = set(re.findall(r'\b\w+\b', query.lower()))
+        query_words = {w for w in query_words if len(w) > 2}
+
+        if not query_words:
+            return []
+
+        # Contar matches por chunk
+        chunk_scores: Dict[str, float] = defaultdict(float)
+
+        for word in query_words:
+            # Exact match
+            if word in self._keyword_index:
+                for chunk_id in self._keyword_index[word]:
+                    chunk_scores[chunk_id] += 1.0
+
+            # Fuzzy match (prefijo)
+            for indexed_word in self._keyword_index:
+                if indexed_word.startswith(word) or word.startswith(indexed_word):
+                    for chunk_id in self._keyword_index[indexed_word]:
+                        chunk_scores[chunk_id] += 0.5
+
+        # Normalizar scores
+        max_score = max(chunk_scores.values()) if chunk_scores else 1
+        chunk_scores = {k: v / max_score for k, v in chunk_scores.items()}
+
+        # Crear resultados
+        results = []
+        for chunk_id, score in chunk_scores.items():
+            if chunk_id in self._chunks and score > 0.1:
+                results.append(SearchResult(
+                    chunk=self._chunks[chunk_id],
+                    score=score,
+                    match_type="keyword",
+                ))
+
+        results.sort(key=lambda x: -x.score)
+        return results[:limit]
+
+    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
+        """Calcula similitud de coseno."""
+        if len(a) != len(b):
+            return 0.0
+
+        dot_product = sum(x * y for x, y in zip(a, b))
+        norm_a = sum(x * x for x in a) ** 0.5
+        norm_b = sum(x * x for x in b) ** 0.5
+
+        if norm_a == 0 or norm_b == 0:
+            return 0.0
+
+        return dot_product / (norm_a * norm_b)
+
+    def _matches_filter(self, file_path: str, pattern: str) -> bool:
+        """Verifica si un archivo coincide con un patrón."""
+        from fnmatch import fnmatch
+        return fnmatch(file_path, pattern)
+
+    def _load_index(self):
+        """Carga índice desde disco."""
+        index_file = self.index_dir / "semantic_index.json"
+
+        if not index_file.exists():
+            return
+
+        try:
+            data = json.loads(index_file.read_text(encoding="utf-8"))
+
+            for chunk_data in data.get("chunks", []):
+                chunk = CodeChunk(
+                    id=chunk_data["id"],
+                    file_path=chunk_data["file_path"],
+                    content=chunk_data["content"],
+                    start_line=chunk_data["start_line"],
+                    end_line=chunk_data["end_line"],
+                    chunk_type=chunk_data["chunk_type"],
+                    name=chunk_data.get("name"),
+                    embedding=chunk_data.get("embedding"),
+                )
+                self._chunks[chunk.id] = chunk
+                self._file_chunks[chunk.file_path].append(chunk.id)
+
+        except Exception as e:
+            print(f"[SemanticSearch] Error loading index: {e}")
+
+    def _save_index(self):
+        """Guarda índice a disco."""
+        index_file = self.index_dir / "semantic_index.json"
+
+        try:
+            data = {
+                "chunks": [c.to_dict() for c in self._chunks.values()],
+                "project_root": str(self.project_root),
+                "timestamp": time.time(),
+            }
+
+            index_file.write_text(
+                json.dumps(data, indent=2, ensure_ascii=False),
+                encoding="utf-8"
+            )
+
+        except Exception as e:
+            print(f"[SemanticSearch] Error saving index: {e}")
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# INSTANCIA GLOBAL
+# ══════════════════════════════════════════════════════════════════════════════
+
+_search_instance: Optional[SemanticSearch] = None
+
+
+def get_semantic_search(project_root: Path = None) -> SemanticSearch:
+    """Obtiene la instancia global de búsqueda semántica."""
+    global _search_instance
+    if _search_instance is None or (project_root and _search_instance.project_root != project_root):
+        _search_instance = SemanticSearch(project_root or Path.cwd())
+    return _search_instance