luckyd-code 1.2.2__py3-none-any.whl

luckyd_code/brain/chunker.py

@@ -0,0 +1,345 @@
+"""Code chunker — splits source files into overlapping chunks for embedding."""
+
+import ast
+import os
+import re
+from pathlib import Path
+
+from ..log import get_logger
+from .constants import LANGUAGE_MAP, Chunk, should_skip
+
+# Regex patterns for non-Python language structure detection
+STRUCTURE_PATTERNS: dict[str, list[tuple[str, str, str]]] = {
+    "javascript": [
+        (r"(?:^|\n)\s*function\s+\*?\s*(\w+)\s*\(", "function", "function"),
+        (r"(?:^|\n)\s*(?:async\s+)?function\s+\*?\s*(\w+)\s*\(", "function", "function"),
+        (r"(?:^|\n)\s*class\s+(\w+)", "class", "class"),
+        (r"(?:^|\n)\s*(?:export\s+)?(?:default\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\(?.*\)?\s*=>", "function", "arrow_function"),
+        (r"(?:^|\n)\s*(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*function", "function", "function_expression"),
+    ],
+    "typescript": [
+        (r"(?:^|\n)\s*function\s+\*?\s*(\w+)\s*\(", "function", "function"),
+        (r"(?:^|\n)\s*(?:async\s+)?function\s+\*?\s*(\w+)\s*\(", "function", "function"),
+        (r"(?:^|\n)\s*class\s+(\w+)", "class", "class"),
+        (r"(?:^|\n)\s*interface\s+(\w+)", "class", "interface"),
+        (r"(?:^|\n)\s*type\s+(\w+)\s*=", "class", "type_alias"),
+        (r"(?:^|\n)\s*(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\(?.*\)?\s*=>", "function", "arrow_function"),
+        (r"(?:^|\n)\s*(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*function", "function", "function_expression"),
+    ],
+    "go": [
+        (r"(?:^|\n)\s*func\s+(\w+)", "function", "function"),
+        (r"(?:^|\n)\s*type\s+(\w+)\s+struct", "class", "struct"),
+        (r"(?:^|\n)\s*type\s+(\w+)\s+interface", "class", "interface"),
+    ],
+    "rust": [
+        (r"(?:^|\n)\s*fn\s+(\w+)", "function", "function"),
+        (r"(?:^|\n)\s*struct\s+(\w+)", "class", "struct"),
+        (r"(?:^|\n)\s*enum\s+(\w+)", "class", "enum"),
+        (r"(?:^|\n)\s*trait\s+(\w+)", "class", "trait"),
+        (r"(?:^|\n)\s*impl\s+(\w+)", "class", "impl"),
+    ],
+}
+
+
+def _chunk_python(filepath: Path, content: str) -> list[Chunk]:
+    chunks: list[Chunk] = []
+    lines = content.split("\n")
+    rel_path = str(filepath)
+
+    try:
+        tree = ast.parse(content, filename=str(filepath))
+    except SyntaxError:
+        return _chunk_by_lines(filepath, content, "python")
+
+    header_end = 1
+    module_doc = ast.get_docstring(tree) or ""
+    for node in ast.iter_child_nodes(tree):
+        if isinstance(node, (ast.Import, ast.ImportFrom)):
+            header_end = node.end_lineno or node.lineno
+        elif isinstance(node, ast.Expr) and module_doc:
+            header_end = node.end_lineno or node.lineno
+        else:
+            break
+
+    header_content = "\n".join(lines[:header_end])
+    if header_content.strip():
+        chunks.append(Chunk(
+            file_path=rel_path,
+            chunk_id=f"{rel_path}:module",
+            start_line=1,
+            end_line=header_end,
+            type="module",
+            name=Path(rel_path).name,
+            language="python",
+            content=header_content,
+        ))
+
+    for node in ast.iter_child_nodes(tree):
+        if isinstance(node, ast.ClassDef):
+            end = node.end_lineno or node.lineno
+            cls_lines = lines[node.lineno - 1:end]
+            if end < len(lines):
+                cls_lines.append(lines[end])
+
+            content = "\n".join(cls_lines)
+            chunks.append(Chunk(
+                file_path=rel_path,
+                chunk_id=f"{rel_path}:class:{node.name}",
+                start_line=node.lineno,
+                end_line=end,
+                type="class",
+                name=node.name,
+                language="python",
+                content=content,
+            ))
+
+            for child in ast.iter_child_nodes(node):
+                if isinstance(child, ast.FunctionDef):
+                    m_end = child.end_lineno or child.lineno
+                    method_lines = lines[child.lineno - 1:m_end]
+                    if m_end < len(lines):
+                        method_lines.append(lines[m_end])
+                    chunks.append(Chunk(
+                        file_path=rel_path,
+                        chunk_id=f"{rel_path}:method:{child.name}",
+                        start_line=child.lineno,
+                        end_line=m_end,
+                        type="method",
+                        name=child.name,
+                        language="python",
+                        content="\n".join(method_lines),
+                    ))
+
+        elif isinstance(node, ast.FunctionDef):
+            end = node.end_lineno or node.lineno
+            func_lines = lines[node.lineno - 1:end]
+            if end < len(lines):
+                func_lines.append(lines[end])
+
+            chunks.append(Chunk(
+                file_path=rel_path,
+                chunk_id=f"{rel_path}:function:{node.name}",
+                start_line=node.lineno,
+                end_line=end,
+                type="function",
+                name=node.name,
+                language="python",
+                content="\n".join(func_lines),
+            ))
+
+    return chunks
+
+
+def _chunk_by_lines(filepath: Path, content: str, language: str) -> list[Chunk]:
+    chunks: list[Chunk] = []
+    lines = content.split("\n")
+    rel_path = str(filepath)
+
+    header_lines: list[str] = []
+    i = 0
+    while i < len(lines):
+        if lines[i].strip() == "" and any(l.strip() for l in lines[i + 1:i + 3]):
+            break
+        header_lines.append(lines[i])
+        i += 1
+    if header_lines:
+        chunks.append(Chunk(
+            file_path=rel_path,
+            chunk_id=f"{rel_path}:module",
+            start_line=1,
+            end_line=i,
+            type="module",
+            name=Path(rel_path).name,
+            language=language,
+            content="\n".join(header_lines).strip(),
+        ))
+
+    block_start = i + 1
+    for match in re.finditer(r"\n\s*\n", content):
+        block_end = match.start()
+        if block_end > block_start:
+            block_content = content[block_start:block_end].strip()
+            if block_content:
+                start_line = content[:block_start].count("\n") + 1
+                end_line = content[:block_end].count("\n") + 1
+                chunks.append(Chunk(
+                    file_path=rel_path,
+                    chunk_id=f"{rel_path}:block:{start_line}",
+                    start_line=start_line,
+                    end_line=end_line,
+                    type="block",
+                    name="",
+                    language=language,
+                    content=block_content,
+                ))
+        block_start = match.end()
+
+    if block_start < len(content):
+        last_block = content[block_start:].strip()
+        if last_block:
+            start_line = content[:block_start].count("\n") + 1
+            end_line = content.count("\n") + 1
+            chunks.append(Chunk(
+                file_path=rel_path,
+                chunk_id=f"{rel_path}:block:{start_line}",
+                start_line=start_line,
+                end_line=end_line,
+                type="block",
+                name="",
+                language=language,
+                content=last_block,
+            ))
+
+    return chunks
+
+
+def _chunk_with_regex(filepath: Path, content: str, language: str) -> list[Chunk]:
+    chunks: list[Chunk] = []
+    lines = content.split("\n")
+    rel_path = str(filepath)
+    patterns = STRUCTURE_PATTERNS.get(language, [])
+
+    header_end = min(20, len(lines))
+    for pattern, _type, _subtype in patterns:
+        for m in re.finditer(pattern, content):
+            line_num = content[:m.start()].count("\n") + 1
+            header_end = min(header_end, line_num - 1)
+            break
+
+    header_content = "\n".join(lines[:header_end]).strip()
+    if header_content:
+        chunks.append(Chunk(
+            file_path=rel_path,
+            chunk_id=f"{rel_path}:module",
+            start_line=1,
+            end_line=header_end,
+            type="module",
+            name=Path(rel_path).name,
+            language=language,
+            content=header_content,
+        ))
+
+    finds: list[tuple[int, str, str, str]] = []
+    for pattern, chunk_type, subtype in patterns:
+        for m in re.finditer(pattern, content):
+            line_num = content[:m.start()].count("\n") + 1
+            name = m.group(1)
+            finds.append((line_num, name, chunk_type, subtype))
+
+    finds.sort(key=lambda x: x[0])
+
+    for idx, (start_line, name, chunk_type, subtype) in enumerate(finds):
+        byte_pos = 0
+        for _ in range(start_line - 1):
+            byte_pos = content.index("\n", byte_pos) + 1
+
+        end_byte = _find_block_end(content, byte_pos)
+        if end_byte <= byte_pos:
+            end_byte = len(content)
+
+        end_line = content[:end_byte].count("\n") + 1
+
+        chunk_content = content[byte_pos:end_byte].rstrip()
+        if not chunk_content:
+            continue
+
+        if idx + 1 < len(finds):
+            next_start = finds[idx + 1][0]
+            next_byte = 0
+            for _ in range(next_start - 1):
+                next_byte = content.index("\n", next_byte) + 1
+            overlap_end = min(end_byte + 200, next_byte)
+            overlap_content = content[end_byte:overlap_end]
+            chunk_content += overlap_content
+
+        chunks.append(Chunk(
+            file_path=rel_path,
+            chunk_id=f"{rel_path}:{chunk_type}:{name}",
+            start_line=start_line,
+            end_line=end_line,
+            type=chunk_type,
+            name=name,
+            language=language,
+            content=chunk_content,
+        ))
+
+    return chunks
+
+
+def _find_block_end(content: str, start_byte: int) -> int:
+    """Find the end of a brace-delimited block starting at start_byte."""
+    # Find the first '{' after start_byte
+    brace_start = content.find("{", start_byte)
+    if brace_start == -1:
+        return len(content)
+
+    depth = 0
+    in_string = False
+    string_char = None
+    i = brace_start
+
+    while i < len(content):
+        ch = content[i]
+        if not in_string:
+            if ch == '"' or ch == "'":
+                in_string = True
+                string_char = ch
+            elif ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    return i + 1
+        else:
+            if ch == "\\":
+                i += 1  # skip escaped char
+            elif ch == string_char:
+                in_string = False
+        i += 1
+
+    return len(content)
+
+
+def chunk_file(filepath: Path) -> list[Chunk]:
+    suffix = filepath.suffix.lower()
+    language = LANGUAGE_MAP.get(suffix)
+    if not language:
+        return []
+
+    try:
+        content = filepath.read_text(encoding="utf-8", errors="replace")
+    except OSError:
+        get_logger().warning("Could not read file %s", filepath, exc_info=True)
+        return []
+
+    if not content.strip():
+        return []
+
+    if language == "python":
+        return _chunk_python(filepath, content)
+    elif language in STRUCTURE_PATTERNS:
+        return _chunk_with_regex(filepath, content, language)
+    else:
+        return _chunk_by_lines(filepath, content, language)
+
+
+def chunk_project(project_root: str) -> list[Chunk]:
+    root = Path(project_root).resolve()
+    all_chunks: list[Chunk] = []
+
+    for dirpath, dirnames, filenames in os.walk(root):
+        dirnames[:] = [d for d in dirnames if not should_skip(d)]
+
+        for fname in filenames:
+            suffix = Path(fname).suffix.lower()
+            if suffix not in LANGUAGE_MAP:
+                continue
+
+            fpath = Path(dirpath) / fname
+            try:
+                chunks = chunk_file(fpath)
+                all_chunks.extend(chunks)
+            except Exception:
+                get_logger().warning("Error chunking %s", fpath, exc_info=True)
+
+    return all_chunks

luckyd_code/brain/constants.py

@@ -0,0 +1,73 @@
+"""Shared constants, types, and helpers for the brain module."""
+
+from typing import TypedDict
+
+from .._data_dir import data_path
+
+# Directories to skip during project scanning
+SKIP_DIRS = {
+    ".git", "__pycache__", "node_modules", ".venv", "venv", "env",
+    ".tox", ".eggs", "dist", "build", ".next", ".nuxt",
+    "target", "vendor", ".bundle", ".claude", ".deepseek-code", ".vscode", ".idea",
+    ".mypy_cache", ".pytest_cache", ".ruff_cache", ".ruff",
+    ".svn", ".hg", "egg-info",
+}
+
+# Supported file extensions and their languages
+LANGUAGE_MAP = {
+    ".py": "python",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".ts": "typescript",
+    ".tsx": "typescript",
+    ".go": "go",
+    ".rs": "rust",
+    ".java": "java",
+    ".cpp": "cpp",
+    ".cc": "cpp",
+    ".cxx": "cpp",
+    ".c": "c",
+    ".h": "c",
+    ".hpp": "cpp",
+    ".cs": "csharp",
+    ".rb": "ruby",
+    ".php": "php",
+    ".swift": "swift",
+    ".kt": "kotlin",
+    ".kts": "kotlin",
+    ".sh": "bash",
+    ".bash": "bash",
+    ".zsh": "bash",
+    ".lua": "lua",
+    ".r": "r",
+    ".R": "r",
+    ".scala": "scala",
+    ".ex": "elixir",
+    ".exs": "elixir",
+    ".md": "markdown",
+    ".toml": "toml",
+    ".yaml": "yaml",
+    ".yml": "yaml",
+    ".json": "json",
+    ".sql": "sql",
+}
+
+# Storage directory for index and cache files
+BRAIN_DIR = data_path("brain")
+
+
+class Chunk(TypedDict):
+    """A chunk of source code ready for embedding and search."""
+    file_path: str
+    chunk_id: str
+    start_line: int
+    end_line: int
+    type: str      # module, class, function, method, block
+    name: str
+    language: str
+    content: str
+    # score is added by the retriever at search time
+
+
+def should_skip(name: str) -> bool:
+    return name in SKIP_DIRS or name.startswith(".")

luckyd_code/brain/embedder.py

@@ -0,0 +1,163 @@
+"""Embedder — generates embeddings for code chunks.
+
+Supports multiple backends:
+  - sentence-transformers (local, no API key)
+  - OpenAI text-embedding-3-small (requires API key)
+
+All dependencies are optional — embedder returns available=False gracefully.
+"""
+
+import os
+from typing import Optional
+
+from ..log import get_logger
+
+# Module-level singleton
+_embedder: Optional["Embedder"] = None
+
+
+def get_embedder() -> "Embedder":
+    """Get or create the shared embedder singleton."""
+    global _embedder
+    if _embedder is None:
+        _embedder = Embedder()
+        _embedder.load()
+    return _embedder
+
+
+class Embedder:
+    """Generates embeddings for text/code using available backends."""
+
+    def __init__(self):
+        self.available = False
+        self.dimension = 0
+        self.model_name = "none"
+        self._model = None
+        self._openai_client = None
+
+    def load(self, model_type: Optional[str] = None) -> bool:
+        """Load the embedding model.
+
+        Args:
+            model_type: "local" for sentence-transformers, "openai" for OpenAI API.
+                        If None, reads from settings or tries local first.
+
+        Returns:
+            True if a backend was successfully loaded.
+        """
+        if model_type is None:
+            try:
+                from .. import settings as cfg
+                model_type = cfg.load_settings().get("embedding_model", "local")
+            except Exception:
+                model_type = "local"
+
+        if model_type == "openai":
+            return self._load_openai()
+        return self._load_local()
+
+    def _load_local(self) -> bool:
+        """Load sentence-transformers for local embeddings."""
+        try:
+            import sentence_transformers
+
+            self._model = sentence_transformers.SentenceTransformer(
+                "all-MiniLM-L6-v2"
+            )
+            self.dimension = 384
+            self.model_name = "sentence-transformers/all-MiniLM-L6-v2"
+            self.available = True
+            get_logger().info(
+                "Loaded local embedding model: %s (dim=%d)",
+                self.model_name, self.dimension,
+            )
+            return True
+        except ImportError:
+            get_logger().info(
+                "sentence-transformers not installed. "
+                "Install with: pip install sentence-transformers"
+            )
+        except Exception as exc:
+            get_logger().warning("Failed to load sentence-transformers: %s", exc)
+
+        self.available = False
+        return False
+
+    def _load_openai(self) -> bool:
+        """Load OpenAI embedding API client."""
+        api_key = os.environ.get("OPENAI_API_KEY")
+        if not api_key:
+            get_logger().warning(
+                "OPENAI_API_KEY not set. Cannot use OpenAI embeddings."
+            )
+            self.available = False
+            return False
+
+        try:
+            from openai import OpenAI
+
+            self._openai_client = OpenAI(api_key=api_key)
+            self.dimension = 1536
+            self.model_name = "text-embedding-3-small"
+            self.available = True
+            get_logger().info("Loaded OpenAI embedding model: text-embedding-3-small")
+            return True
+        except ImportError:
+            get_logger().warning("openai package not installed.")
+        except Exception as exc:
+            get_logger().warning("Failed to load OpenAI embeddings: %s", exc)
+
+        self.available = False
+        return False
+
+    def embed(self, texts: list[str]) -> Optional[list[list[float]]]:
+        """Embed a list of texts into vectors.
+
+        Args:
+            texts: List of text strings to embed.
+
+        Returns:
+            List of embedding vectors (list of floats), or None if unavailable.
+        """
+        if not self.available or not texts:
+            return None
+
+        # Filter out empty texts
+        valid_texts = [t for t in texts if t and t.strip()]
+        if not valid_texts:
+            return None
+
+        # Truncate very long texts (sentence-transformers has 256/512 token limit)
+        truncated = [t[:8192] for t in valid_texts]
+
+        try:
+            if self._model is not None:
+                # sentence-transformers
+                import numpy as np
+
+                embeddings = self._model.encode(truncated, show_progress_bar=False)
+                return embeddings.tolist() if isinstance(embeddings, np.ndarray) else embeddings  # type: ignore[return-value]
+            elif self._openai_client is not None:
+                resp = self._openai_client.embeddings.create(
+                    model="text-embedding-3-small",
+                    input=truncated,
+                )
+                return [item.embedding for item in resp.data]
+        except Exception as exc:
+            get_logger().warning("Embedding failed: %s", exc)
+
+        return None
+
+    def embed_query(self, query: str) -> Optional[list[float]]:
+        """Embed a single query string.
+
+        Args:
+            query: The search query.
+
+        Returns:
+            Single embedding vector, or None if unavailable.
+        """
+        result = self.embed([query])
+        if result:
+            return result[0]
+        return None