luckyd-code 1.2.2__py3-none-any.whl

luckyd_code/brain/graph.py

@@ -0,0 +1,311 @@
+"""Knowledge graph — stores and queries codebase structure across sessions."""
+
+import json
+import time
+from pathlib import Path
+from typing import Any
+
+from ..log import get_logger
+from .constants import BRAIN_DIR
+
+
+GRAPH_FILE = BRAIN_DIR / "graph.json"
+
+Node = dict[str, Any]
+Edge = dict[str, str]
+
+
+class KnowledgeGraph:
+    """Persistent knowledge graph of codebase structure.
+
+    Nodes: modules, classes, functions
+    Edges: imports, contains, calls, inherits
+    """
+
+    def __init__(self) -> None:
+        self.nodes: dict[str, Node] = {}
+        self.edges: list[Edge] = []
+        self.stats: dict[str, Any] = {
+            "node_count": 0,
+            "edge_count": 0,
+            "last_built": 0,
+            "files_parsed": 0,
+            "errors": 0,
+        }
+
+    def build(self, project_root: str, parsed_files: list[dict[str, Any]]) -> None:
+        self.nodes = {}
+        self.edges = []
+        self.stats["last_built"] = time.time()
+        self.stats["files_parsed"] = len(parsed_files)
+        self.stats["errors"] = 0
+
+        for pf in parsed_files:
+            if pf["errors"]:
+                self.stats["errors"] += len(pf["errors"])
+                continue
+
+            rel_path = pf["module"]
+            module_id = f"module:{rel_path}"
+
+            self.nodes[module_id] = {
+                "type": "module",
+                "name": Path(rel_path).name,
+                "file": rel_path,
+                "line": 1,
+                "doc": "",
+                "size": pf["size"],
+            }
+
+            for imp in pf["imports"]:
+                import_id = f"import:{imp['module']}:{imp['name']}"
+                if import_id not in self.nodes:
+                    self.nodes[import_id] = {
+                        "type": "import",
+                        "name": imp["name"],
+                        "module": imp["module"],
+                        "alias": imp.get("alias"),
+                        "file": rel_path,
+                        "line": 0,
+                        "doc": "",
+                    }
+                self.edges.append({"from": module_id, "to": import_id, "type": "imports"})
+
+            for cls in pf["classes"]:
+                cls_id = f"class:{rel_path}:{cls['name']}"
+                self.nodes[cls_id] = {
+                    "type": "class",
+                    "name": cls["name"],
+                    "file": rel_path,
+                    "line": cls["line"],
+                    "end_line": cls["end_line"],
+                    "bases": cls["base_names"],
+                    "decorators": cls["decorators"],
+                    "doc": cls["docstring"][:200],
+                }
+                self.edges.append({"from": module_id, "to": cls_id, "type": "contains"})
+
+                for base in cls["base_names"]:
+                    if base and base != "object":
+                        self.edges.append({
+                            "from": cls_id, "to": f"class:??:{base}", "type": "inherits"
+                        })
+
+                for method in cls["methods"]:
+                    method_id = f"method:{rel_path}:{cls['name']}.{method['name']}"
+                    self.nodes[method_id] = {
+                        "type": "method",
+                        "name": method["name"],
+                        "class": cls["name"],
+                        "file": rel_path,
+                        "line": method["line"],
+                        "end_line": method["end_line"],
+                        "decorators": method["decorators"],
+                        "doc": method["docstring"][:200],
+                    }
+                    self.edges.append({"from": cls_id, "to": method_id, "type": "contains"})
+                    for call in method["calls"]:
+                        self.edges.append({"from": method_id, "to": f"func:??:{call}", "type": "calls"})
+
+            for func in pf["functions"]:
+                func_id = f"func:{rel_path}:{func['name']}"
+                self.nodes[func_id] = {
+                    "type": "function",
+                    "name": func["name"],
+                    "file": rel_path,
+                    "line": func["line"],
+                    "end_line": func["end_line"],
+                    "decorators": func["decorators"],
+                    "doc": func["docstring"][:200],
+                }
+                self.edges.append({"from": module_id, "to": func_id, "type": "contains"})
+                for call in func["calls"]:
+                    self.edges.append({"from": func_id, "to": f"func:??:{call}", "type": "calls"})
+
+        self.stats["node_count"] = len(self.nodes)
+        self.stats["edge_count"] = len(self.edges)
+
+    # --- Persistence ---
+
+    def save(self) -> None:
+        BRAIN_DIR.mkdir(parents=True, exist_ok=True)
+        data = {"nodes": self.nodes, "edges": self.edges, "stats": self.stats}
+        GRAPH_FILE.write_text(json.dumps(data, indent=2), encoding="utf-8")
+
+    def load(self) -> bool:
+        if GRAPH_FILE.exists():
+            try:
+                data: Any = json.loads(GRAPH_FILE.read_text(encoding="utf-8"))
+                self.nodes = data.get("nodes", {})
+                self.edges = data.get("edges", {}) if isinstance(data.get("edges"), dict) else data.get("edges", [])
+                self.stats = data.get("stats", {})
+                return True
+            except (json.JSONDecodeError, OSError):
+                get_logger().warning("Could not load knowledge graph from %s", GRAPH_FILE, exc_info=True)
+        return False
+
+    def search(self, query: str, max_results: int = 15) -> list[Node]:
+        q = query.lower()
+        scored: list[tuple[int, str, Node]] = []
+
+        for nid, node in self.nodes.items():
+            score = 0
+            if q in node.get("name", "").lower():
+                score += 10
+            if q in node.get("file", "").lower():
+                score += 5
+            if q in node.get("doc", "").lower():
+                score += 3
+            if q in node.get("module", "").lower():
+                score += 2
+            if q in node.get("class", "").lower():
+                score += 2
+            if score > 0:
+                scored.append((score, nid, node))
+
+        scored.sort(key=lambda x: -x[0])
+        seen: set[str] = set()
+        top: list[Node] = []
+        for _score, nid, node in scored[:max_results]:
+            if nid not in seen:
+                top.append(node)
+                seen.add(nid)
+        return top
+
+    def get_related(self, node_id: str, max_depth: int = 1) -> list[Node]:
+        related: set[str] = set()
+        current = {node_id}
+
+        for _ in range(max_depth):
+            next_set: set[str] = set()
+            for edge in self.edges:
+                if isinstance(edge, dict):
+                    if edge["from"] in current:
+                        next_set.add(edge["to"])
+                    if edge["to"] in current:
+                        next_set.add(edge["from"])
+            current = next_set
+            related.update(current)
+
+        return [
+            self.nodes.get(nid, {"name": nid, "type": "unknown", "file": ""})
+            for nid in related if nid != node_id
+        ]
+
+    def get_by_file(self, filepath: str) -> list[Node]:
+        return [
+            node for node in self.nodes.values()
+            if node.get("file", "").endswith(filepath)
+        ]
+
+    def get_by_type(self, node_type: str) -> list[Node]:
+        return [
+            node for node in self.nodes.values()
+            if node.get("type") == node_type
+        ]
+
+    def find_dependents(self, symbol_name: str, max_results: int = 15) -> list[dict]:
+        """Find all nodes that depend on a symbol by traversing incoming edges."""
+        matches = self.search(symbol_name, max_results=5)
+        if not matches:
+            return []
+
+        # Build a lookup of known node IDs so we can match even with incomplete IDs
+        node_ids = set(self.nodes.keys())
+        dependents: list[dict] = []
+        seen: set[str] = set()
+
+        for match in matches:
+            # Generate the most likely node ID for this match
+            candidate_ids = [
+                f"{match['type']}:{match.get('file', '')}:{match['name']}",
+                f"func:??:{match['name']}",
+                f"method:??:{match['name']}",
+                f"class:??:{match['name']}",
+            ]
+            matched_id = None
+            for cid in candidate_ids:
+                if cid in node_ids:
+                    matched_id = cid
+                    break
+
+            if not matched_id:
+                continue
+
+            for edge in self.edges:
+                if isinstance(edge, dict) and edge.get("to") == matched_id:
+                    src = self.nodes.get(edge["from"])
+                    if src and edge["from"] not in seen:
+                        seen.add(edge["from"])
+                        dependents.append({
+                            "name": f"{src.get('type', '?')}:{src.get('name', '?')}",
+                            "file": src.get("file", ""),
+                            "relation": edge.get("type", ""),
+                            "line": src.get("line", 0),
+                        })
+
+        dependents.sort(key=lambda x: (x["file"], x["line"]))
+        return dependents[:max_results]
+
+    def summarize(self, max_modules: int = 20) -> str:
+        lines = ["<knowledge-graph>"]
+        lines.append(f"Graph: {self.stats.get('node_count', 0)} symbols, "
+                     f"{self.stats.get('edge_count', 0)} relationships, "
+                     f"{self.stats.get('files_parsed', 0)} files")
+
+        by_file: dict[str, list[Node]] = {}
+        for node in self.nodes.values():
+            f = node.get("file", "")
+            if f:
+                by_file.setdefault(f, []).append(node)
+
+        count = 0
+        for filepath, nodes in sorted(by_file.items()):
+            if count >= max_modules:
+                break
+            count += 1
+            classes = [n for n in nodes if n["type"] == "class"]
+            functions = [n for n in nodes if n["type"] == "function"]
+
+            short_path = Path(filepath).name
+            parts = [short_path]
+            if classes:
+                parts.append(f"classes={{{','.join(c['name'] for c in classes)}}}")
+            if functions:
+                parts.append(f"functions={{{','.join(f['name'] for f in functions)}}}")
+            lines.append(f"  {' | '.join(parts)}")
+
+        lines.append("</knowledge-graph>")
+        return "\n".join(lines)
+
+    def stats_text(self) -> str:
+        by_type: dict[str, int] = {}
+        for node in self.nodes.values():
+            t = node.get("type", "unknown")
+            by_type[t] = by_type.get(t, 0) + 1
+
+        by_file: dict[str, int] = {}
+        for node in self.nodes.values():
+            f = node.get("file", "")
+            if f:
+                by_file[f] = by_file.get(f, 0) + 1
+
+        lines = [
+            f"Nodes: {self.stats.get('node_count', 0)}",
+            f"Edges: {self.stats.get('edge_count', 0)}",
+            f"Files parsed: {self.stats.get('files_parsed', 0)}",
+            f"Parse errors: {self.stats.get('errors', 0)}",
+        ]
+        if by_type:
+            lines.append("\nBy type:")
+            for t, c in sorted(by_type.items(), key=lambda x: -x[1]):
+                lines.append(f"  {t}: {c}")
+        if by_file:
+            lines.append("\nBy file:")
+            for f, c in sorted(by_file.items(), key=lambda x: -x[1])[:20]:
+                lines.append(f"  {f}: {c} symbols")
+        if self.stats.get("last_built"):
+            last = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.stats["last_built"]))
+            lines.append(f"\nLast built: {last}")
+
+        return "\n".join(lines)

luckyd_code/brain/indexer.py

@@ -0,0 +1,316 @@
+"""Vector indexer — builds and queries FAISS vector index for code chunks."""
+
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+from ..log import get_logger
+from .constants import BRAIN_DIR, LANGUAGE_MAP, SKIP_DIRS
+
+INDEX_FILE = BRAIN_DIR / "index.faiss"
+CHUNKS_FILE = BRAIN_DIR / "chunks.json"
+MTIMES_FILE = BRAIN_DIR / "mtimes.json"
+STATS_FILE = BRAIN_DIR / "stats.json"
+
+CHUNK_SIZE = 384  # all-MiniLM-L6-v2 dimension
+
+
+class VectorIndexer:
+    """Manages the FAISS vector index with mtime tracking."""
+
+    def __init__(self):
+        self.index: Any = None  # FAISS index object
+        self.chunks: list[dict[str, Any]] = []
+        self.file_mtimes: dict[str, tuple[float, int]] = {}
+        self.stats: dict[str, Any] = {
+            "chunks": 0,
+            "files": 0,
+            "languages": {},
+            "last_indexed": 0,
+            "dimension": 0,
+            "index_size_bytes": 0,
+        }
+        self._faiss_available = False
+
+    def _check_deps(self) -> bool:
+        """Check if FAISS and numpy are available."""
+        if not self._faiss_available:
+            try:
+                import faiss
+                import numpy as np
+
+                self._faiss = faiss
+                self._np = np
+                self._faiss_available = True
+                return True
+            except ImportError:
+                get_logger().info(
+                    "faiss-cpu not available. Vector search disabled. "
+                    "Install with: pip install faiss-cpu"
+                )
+                return False
+        return True
+
+    def build(self, chunks: list[dict[str, Any]]) -> dict[str, Any]:
+        """Build the FAISS index from chunks.
+
+        Args:
+            chunks: List of chunk dicts from chunker.
+
+        Returns:
+            Stats dict.
+        """
+        from .embedder import get_embedder
+
+        if not self._check_deps():
+            self.stats["chunks"] = len(chunks)
+            self.stats["files"] = len(set(c["file_path"] for c in chunks))
+            self.stats["last_indexed"] = time.time()
+            return self.stats
+
+        embedder = get_embedder()
+        if not embedder.available:
+            self.stats["chunks"] = len(chunks)
+            self.stats["files"] = len(set(c["file_path"] for c in chunks))
+            self.stats["last_indexed"] = time.time()
+            return self.stats
+
+        if not chunks:
+            self.chunks = []
+            self.index = None
+            self.stats["chunks"] = 0
+            self.stats["files"] = 0
+            self.stats["languages"] = {}
+            self.stats["last_indexed"] = time.time()
+            return self.stats
+
+        # Track languages
+        languages: dict[str, int] = {}
+        for c in chunks:
+            lang = c.get("language", "unknown")
+            languages[lang] = languages.get(lang, 0) + 1
+
+        # Sort chunks by file_path then start_line for stable ordering
+        chunks.sort(key=lambda c: (c["file_path"], c.get("start_line", 0)))
+        self.chunks = chunks
+
+        # Embed all chunk contents
+        texts = [c.get("content", "") for c in chunks]
+        embeddings = embedder.embed(texts)
+
+        if embeddings is None:
+            self.stats["chunks"] = len(chunks)
+            self.stats["files"] = len(set(c["file_path"] for c in chunks))
+            self.stats["last_indexed"] = time.time()
+            return self.stats
+
+        # Build FAISS index
+        dim = len(embeddings[0])
+        idx = self._faiss.IndexFlatIP(dim)  # Inner product = cosine sim for normalized vectors
+        vectors = self._np.array(embeddings, dtype=self._np.float32)
+
+        # Normalize vectors for cosine similarity
+        self._faiss.normalize_L2(vectors)
+        idx.add(vectors)
+
+        self.index = idx
+        self.stats = {
+            "chunks": len(chunks),
+            "files": len(set(c["file_path"] for c in chunks)),
+            "languages": languages,
+            "last_indexed": time.time(),
+            "dimension": dim,
+            "index_size_bytes": 0,
+        }
+
+        return self.stats
+
+    def search(
+        self, query: str, k: int = 10
+    ) -> list[dict[str, Any]]:
+        """Search the index by embedding the query.
+
+        Args:
+            query: Natural language search query.
+            k: Number of results to return.
+
+        Returns:
+            List of chunk dicts with a 'score' key added.
+        """
+        from .embedder import get_embedder
+
+        if not self._check_deps() or self.index is None or self.index.ntotal == 0:
+            return []
+
+        embedder = get_embedder()
+        if not embedder.available:
+            return []
+
+        query_vec = embedder.embed_query(query)
+        if query_vec is None:
+            return []
+
+        # Normalize query vector
+        q = self._np.array([query_vec], dtype=self._np.float32)
+        self._faiss.normalize_L2(q)
+
+        k_actual = min(k, self.index.ntotal)
+        if k_actual == 0:
+            return []
+
+        scores, indices = self.index.search(q, k_actual)
+
+        results = []
+        for score, idx in zip(scores[0], indices[0]):
+            if idx < 0 or idx >= len(self.chunks):
+                continue
+            chunk = dict(self.chunks[idx])
+            chunk["score"] = float(score)
+            results.append(chunk)
+
+        return results
+
+    def save(self) -> bool:
+        """Save the index, chunks, and mtimes to disk."""
+        BRAIN_DIR.mkdir(parents=True, exist_ok=True)
+
+        try:
+            # Save FAISS index
+            if self._faiss_available and self.index is not None:
+                self._faiss.write_index(self.index, str(INDEX_FILE))
+                self.stats["index_size_bytes"] = INDEX_FILE.stat().st_size
+
+            # Save chunks with content
+            CHUNKS_FILE.write_text(
+                json.dumps(self.chunks, indent=2), encoding="utf-8"
+            )
+
+            # Save mtimes
+            MTIMES_FILE.write_text(
+                json.dumps(self.file_mtimes), encoding="utf-8"
+            )
+
+            # Save stats
+            STATS_FILE.write_text(
+                json.dumps(self.stats), encoding="utf-8"
+            )
+
+            return True
+        except Exception as exc:
+            get_logger().warning("Failed to save vector index: %s", exc)
+            return False
+
+    def load(self) -> bool:
+        """Load the index and metadata from disk.
+
+        Returns:
+            True if index was loaded successfully.
+        """
+        if not INDEX_FILE.exists() or not CHUNKS_FILE.exists():
+            return False
+
+        try:
+            self._check_deps()
+
+            # Load chunks
+            self.chunks = json.loads(CHUNKS_FILE.read_text(encoding="utf-8"))
+            if not self.chunks:
+                return False
+
+            # Load FAISS index
+            if self._faiss_available and INDEX_FILE.exists():
+                self.index = self._faiss.read_index(str(INDEX_FILE))
+
+            # Load mtimes
+            if MTIMES_FILE.exists():
+                self.file_mtimes = json.loads(MTIMES_FILE.read_text(encoding="utf-8")) or {}
+
+            # Load stats
+            if STATS_FILE.exists():
+                self.stats = json.loads(STATS_FILE.read_text(encoding="utf-8")) or {}
+
+            return True
+
+        except Exception as exc:
+            get_logger().warning("Failed to load vector index: %s", exc)
+            return False
+
+    def stats_text(self) -> str:
+        """Return human-readable statistics."""
+        lines = [
+            f"Chunks indexed: {self.stats.get('chunks', 0)}",
+            f"Files: {self.stats.get('files', 0)}",
+        ]
+
+        languages = self.stats.get("languages", {})
+        if languages:
+            lines.append(f"Languages: {', '.join(f'{k}={v}' for k, v in sorted(languages.items()))}")
+
+        dim = self.stats.get("dimension", 0)
+        if dim:
+            lines.append(f"Vector dimension: {dim}")
+
+        size = self.stats.get("index_size_bytes", 0)
+        if size:
+            if size < 1024:
+                size_str = f"{size} B"
+            elif size < 1024 * 1024:
+                size_str = f"{size / 1024:.1f} KB"
+            else:
+                size_str = f"{size / 1024 / 1024:.1f} MB"
+            lines.append(f"Index size: {size_str}")
+
+        last = self.stats.get("last_indexed", 0)
+        if last:
+            last_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(last))
+            lines.append(f"Last indexed: {last_str}")
+
+        if not self._faiss_available:
+            lines.append("FAISS not available (install faiss-cpu for vector search)")
+
+        return "\n".join(lines)
+
+    def get_changed_files(self, project_root: str) -> list[str]:
+        """Check which files have changed since last index.
+
+        Args:
+            project_root: Root directory to scan.
+
+        Returns:
+            List of file paths that have changed or are new.
+        """
+        changed: list[str] = []
+        root = Path(project_root).resolve()
+
+        for dirpath, dirnames, filenames in os.walk(root):
+            dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")]
+
+            for fname in filenames:
+                suffix = Path(fname).suffix.lower()
+                if suffix not in LANGUAGE_MAP:
+                    continue
+
+                fpath = Path(dirpath) / fname
+                try:
+                    st = fpath.stat()
+                    mtime = st.st_mtime
+                    size = st.st_size
+                except OSError:
+                    continue
+
+                fpath_str = str(fpath)
+                if fpath_str in self.file_mtimes:
+                    old_mtime, old_size = self.file_mtimes[fpath_str]
+                    if old_mtime == mtime and old_size == size:
+                        continue
+
+                changed.append(fpath_str)
+
+        return changed
+
+    @property
+    def is_available(self) -> bool:
+        """Whether the index is loaded and ready."""
+        return self._faiss_available and self.index is not None and self.index.ntotal > 0