sourcefire 0.2.0__py3-none-any.whl

sourcefire/indexer/metadata.py

@@ -0,0 +1,289 @@
+"""Language-agnostic code metadata extractor.
+
+Extracts structural metadata from source files using tree-sitter (when available)
+or regex fallback. Language-specific behavior is driven by LanguageProfile.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Any, Optional
+
+from sourcefire.indexer.language_profiles import LanguageProfile, get_profile_for_extension
+
+# ---------------------------------------------------------------------------
+# Optional tree-sitter import
+# ---------------------------------------------------------------------------
+try:
+    from tree_sitter_languages import get_language, get_parser  # type: ignore
+
+    _HAS_TREE_SITTER = True
+except Exception:
+    _HAS_TREE_SITTER = False
+
+# Cache loaded parsers by language name
+_PARSERS: dict[str, Any] = {}
+
+
+def _get_parser(language: str) -> Any | None:
+    """Get a tree-sitter parser for the given language, or None."""
+    if not _HAS_TREE_SITTER:
+        return None
+    if language not in _PARSERS:
+        try:
+            _PARSERS[language] = get_parser(language)
+        except Exception:
+            _PARSERS[language] = None
+    return _PARSERS[language]
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def extract_metadata(source: str, file_path: str, profile: Optional[LanguageProfile] = None) -> dict[str, Any]:
+    """Return a metadata dict for a source file.
+
+    Keys:
+        imports   : list[str] — imported URIs/modules
+        exports   : list[str] — top-level declaration names
+        layer     : str       — architecture layer inferred from path
+        feature   : str       — feature name inferred from path
+        file_type : str       — file role inferred from path
+    """
+    if profile is None:
+        ext = Path(file_path).suffix
+        profile = get_profile_for_extension(ext)
+
+    if profile is None:
+        # No profile — return path-only defaults
+        return {
+            "imports": [],
+            "exports": [],
+            "layer": "unknown",
+            "feature": "unknown",
+            "file_type": "unknown",
+        }
+
+    parser = _get_parser(profile.tree_sitter_language) if profile.tree_sitter_language else None
+
+    if parser is not None and source:
+        imports = _extract_imports_tree_sitter(source, parser, profile)
+        exports = _extract_exports_tree_sitter(source, parser, profile)
+    elif source:
+        imports = _extract_imports_regex(source, profile)
+        exports = _extract_exports_regex(source, profile)
+    else:
+        imports = []
+        exports = []
+
+    return {
+        "imports": imports,
+        "exports": exports,
+        "layer": _infer_layer(file_path, profile),
+        "feature": _infer_feature(file_path, profile),
+        "file_type": _infer_file_type(file_path, profile),
+    }
+
+
+def chunk_source_file(
+    source: str,
+    file_path: str,
+    profile: Optional[LanguageProfile] = None,
+    chunk_size: int = 1000,
+) -> list[dict[str, Any]]:
+    """Split a source file into chunks and attach metadata to each chunk.
+
+    Splitting strategy:
+    1. If tree-sitter is available and a profile exists, split at declaration boundaries.
+    2. Otherwise, use regex boundary splitting.
+    3. If the source is shorter than chunk_size, return a single chunk.
+
+    Each chunk dict has keys:
+        text     : str  — the chunk text
+        metadata : dict — output of extract_metadata for the whole file
+    """
+    if profile is None:
+        ext = Path(file_path).suffix
+        profile = get_profile_for_extension(ext)
+
+    metadata = extract_metadata(source, file_path, profile)
+
+    if profile is None:
+        # No profile — return source as a single chunk
+        return [{"text": source, "metadata": metadata}]
+
+    parser = _get_parser(profile.tree_sitter_language) if profile.tree_sitter_language else None
+
+    if parser is not None:
+        raw_chunks = _chunk_tree_sitter(source, parser, profile, chunk_size)
+    else:
+        raw_chunks = _chunk_regex(source, profile, chunk_size)
+
+    return [{"text": text, "metadata": metadata} for text in raw_chunks]
+
+
+# ---------------------------------------------------------------------------
+# Tree-sitter implementations
+# ---------------------------------------------------------------------------
+
+
+def _extract_imports_tree_sitter(source: str, parser: Any, profile: LanguageProfile) -> list[str]:
+    """Extract import URIs/modules using tree-sitter."""
+    tree = parser.parse(source.encode())
+    imports: list[str] = []
+    _walk_for_imports(tree.root_node, imports, profile)
+    return imports
+
+
+def _walk_for_imports(node: Any, imports: list[str], profile: LanguageProfile) -> None:
+    if node.type in profile.import_node_types:
+        for child in node.children:
+            if child.type == profile.string_literal_type:
+                uri = child.text.decode().strip("'\"")
+                imports.append(uri)
+    for child in node.children:
+        _walk_for_imports(child, imports, profile)
+
+
+def _extract_exports_tree_sitter(source: str, parser: Any, profile: LanguageProfile) -> list[str]:
+    """Extract top-level declaration names using tree-sitter."""
+    tree = parser.parse(source.encode())
+    exports: list[str] = []
+    _walk_for_exports(tree.root_node, exports, profile)
+    return exports
+
+
+def _walk_for_exports(node: Any, exports: list[str], profile: LanguageProfile) -> None:
+    if node.type in profile.export_node_types:
+        for child in node.children:
+            if child.type == "identifier":
+                exports.append(child.text.decode())
+                break
+    for child in node.children:
+        _walk_for_exports(child, exports, profile)
+
+
+def _chunk_tree_sitter(source: str, parser: Any, profile: LanguageProfile, chunk_size: int) -> list[str]:
+    """Split source at top-level declaration boundaries using tree-sitter."""
+    tree = parser.parse(source.encode())
+    chunks: list[str] = []
+    current: list[str] = []
+    current_len = 0
+
+    for node in tree.root_node.children:
+        if node.type in profile.boundary_node_types:
+            text = node.text.decode()
+            if current_len + len(text) > chunk_size and current:
+                chunks.append("\n".join(current).strip())
+                current = []
+                current_len = 0
+            current.append(text)
+            current_len += len(text)
+
+    if current:
+        chunks.append("\n".join(current).strip())
+
+    return chunks or [source]
+
+
+# ---------------------------------------------------------------------------
+# Regex implementations
+# ---------------------------------------------------------------------------
+
+
+def _extract_imports_regex(source: str, profile: LanguageProfile) -> list[str]:
+    if not profile.import_pattern:
+        return []
+    regex = re.compile(profile.import_pattern, re.MULTILINE)
+    results = []
+    for m in regex.finditer(source):
+        # Take the first non-None group (different patterns use different groups)
+        for g in m.groups():
+            if g:
+                results.append(g)
+                break
+    return results
+
+
+def _extract_exports_regex(source: str, profile: LanguageProfile) -> list[str]:
+    if not profile.export_pattern:
+        return []
+    return re.compile(profile.export_pattern, re.MULTILINE).findall(source)
+
+
+def _chunk_regex(source: str, profile: LanguageProfile, chunk_size: int) -> list[str]:
+    """Split source at declaration boundaries using regex."""
+    if len(source) <= chunk_size:
+        return [source]
+
+    if not profile.boundary_pattern:
+        # No boundary pattern — fall back to size-based splitting
+        return [source[i : i + chunk_size] for i in range(0, len(source), chunk_size)]
+
+    boundary_re = re.compile(profile.boundary_pattern, re.MULTILINE)
+    boundaries = [m.start() for m in boundary_re.finditer(source)]
+
+    if not boundaries:
+        return [source[i : i + chunk_size] for i in range(0, len(source), chunk_size)]
+
+    # Include the preamble (imports, top-level comments) before first boundary.
+    segments: list[str] = []
+    starts = boundaries + [len(source)]
+    if boundaries[0] > 0:
+        preamble = source[: boundaries[0]].strip()
+        if preamble:
+            segments.append(preamble)
+
+    current_text = ""
+    for idx, start in enumerate(boundaries):
+        end = starts[idx + 1]
+        segment = source[start:end].strip()
+        if len(current_text) + len(segment) > chunk_size and current_text:
+            segments.append(current_text.strip())
+            current_text = segment
+        else:
+            current_text = (current_text + "\n\n" + segment).strip() if current_text else segment
+
+    if current_text:
+        segments.append(current_text.strip())
+
+    return segments or [source]
+
+
+# ---------------------------------------------------------------------------
+# Path-based inference helpers
+# ---------------------------------------------------------------------------
+
+
+def _infer_layer(file_path: str, profile: LanguageProfile) -> str:
+    for part in profile.layer_parts:
+        if f"/{part}/" in file_path:
+            return part
+    return "unknown"
+
+
+def _infer_feature(file_path: str, profile: LanguageProfile) -> str:
+    if profile.feature_regex:
+        match = re.search(profile.feature_regex, file_path)
+        if match:
+            return match.group(1)
+    if "/core/" in file_path:
+        return "core"
+    return "unknown"
+
+
+def _infer_file_type(file_path: str, profile: LanguageProfile) -> str:
+    stem = Path(file_path).stem.lower()
+
+    for suffix, file_type in profile.file_type_suffixes:
+        if stem.endswith(suffix.lower()):
+            return file_type
+
+    for pattern, file_type in profile.directory_type_patterns.items():
+        if pattern in file_path:
+            return file_type
+
+    return "unknown"

sourcefire/indexer/pipeline.py

@@ -0,0 +1,406 @@
+"""Full indexing pipeline for Sourcefire.
+
+Scans a codebase, chunks files (AST-aware when a language profile exists,
+simple split otherwise), embeds all chunks, and inserts them into ChromaDB.
+Also builds the import graph for graph-augmented retrieval.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+from pathlib import Path
+from typing import Any
+
+import chromadb
+
+from sourcefire.config import SourcefireConfig
+from sourcefire.db import add_chunks, reset_collection, delete_file_chunks, get_indexed_files, get_stored_mtimes
+from sourcefire.indexer.embeddings import embed_batch
+from sourcefire.indexer.language_profiles import LanguageProfile, get_profile, get_profile_for_extension
+from sourcefire.indexer.metadata import chunk_source_file, extract_metadata
+
+
+# ---------------------------------------------------------------------------
+# .gitignore parsing
+# ---------------------------------------------------------------------------
+
+
+def _parse_gitignore(codebase_path: Path) -> list[str]:
+    """Parse .gitignore and return a list of glob patterns to exclude."""
+    gitignore_path = codebase_path / ".gitignore"
+    if not gitignore_path.is_file():
+        return []
+
+    patterns: list[str] = []
+    try:
+        for line in gitignore_path.read_text(encoding="utf-8", errors="replace").splitlines():
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            if line.startswith("/"):
+                line = line[1:]
+            if line.endswith("/"):
+                line = line + "**"
+            patterns.append(line)
+    except OSError:
+        pass
+    return patterns
+
+
+# ---------------------------------------------------------------------------
+# Pattern matching helpers
+# ---------------------------------------------------------------------------
+
+
+def _match_patterns(rel_path: str, patterns: list[str]) -> bool:
+    """Return True if *rel_path* matches any of *patterns* using fnmatch."""
+    for pattern in patterns:
+        if fnmatch.fnmatch(rel_path, pattern):
+            return True
+    return False
+
+
+# ---------------------------------------------------------------------------
+# Non-AST chunker (simple recursive text split)
+# ---------------------------------------------------------------------------
+
+
+def _chunk_plain_text(
+    text: str,
+    chunk_size: int = 1000,
+    chunk_overlap: int = 300,
+) -> list[str]:
+    """Split *text* into overlapping chunks of at most *chunk_size* characters."""
+    if len(text) <= chunk_size:
+        return [text]
+
+    chunks: list[str] = []
+    start = 0
+    while start < len(text):
+        end = start + chunk_size
+        chunks.append(text[start:end])
+        if end >= len(text):
+            break
+        start = end - chunk_overlap
+
+    return chunks
+
+
+# ---------------------------------------------------------------------------
+# File scanning
+# ---------------------------------------------------------------------------
+
+
+def _collect_files(
+    codebase_path: Path,
+    config: SourcefireConfig,
+    profile: LanguageProfile | None,
+) -> list[Path]:
+    """Return all files under *codebase_path* that pass include/exclude filters.
+
+    Config patterns are authoritative. If config has include patterns, those
+    are used. Otherwise falls back to language profile patterns.
+    """
+    include_patterns: list[str] = list(config.include) if config.include else []
+    exclude_patterns: list[str] = list(config.exclude) if config.exclude else []
+
+    # If no config patterns, fall back to profile
+    if not include_patterns and profile:
+        include_patterns = list(profile.include_patterns)
+    if not exclude_patterns and profile:
+        exclude_patterns = list(profile.exclude_patterns)
+
+    # Always exclude .gitignore patterns
+    exclude_patterns.extend(_parse_gitignore(codebase_path))
+
+    if not include_patterns:
+        include_patterns = ["**/*"]
+
+    matched: list[Path] = []
+    for pattern in include_patterns:
+        for file_path in codebase_path.glob(pattern):
+            if not file_path.is_file():
+                continue
+            rel = file_path.relative_to(codebase_path).as_posix()
+            if not _match_patterns(rel, exclude_patterns):
+                matched.append(file_path)
+
+    seen: set[Path] = set()
+    unique: list[Path] = []
+    for p in matched:
+        if p not in seen:
+            seen.add(p)
+            unique.append(p)
+
+    return unique
+
+
+# ---------------------------------------------------------------------------
+# Chunk production
+# ---------------------------------------------------------------------------
+
+
+def _chunks_for_file(
+    file_path: Path,
+    codebase_path: Path,
+    profile: LanguageProfile | None,
+    chunk_size: int = 1000,
+) -> list[dict[str, Any]]:
+    """Return a list of chunk dicts for *file_path*."""
+    rel = file_path.relative_to(codebase_path).as_posix()
+    source = file_path.read_text(encoding="utf-8", errors="replace")
+
+    file_profile = get_profile_for_extension(file_path.suffix) or profile
+
+    if file_profile and file_path.suffix in [e for e in file_profile.file_extensions]:
+        raw_chunks = chunk_source_file(source, rel, file_profile, chunk_size=chunk_size)
+        chunks_out: list[dict[str, Any]] = []
+        for idx, chunk in enumerate(raw_chunks):
+            meta = chunk["metadata"]
+            chunks_out.append({
+                "filename": rel,
+                "location": f"{rel}:{idx}",
+                "code": chunk["text"],
+                "feature": meta.get("feature", ""),
+                "layer": meta.get("layer", ""),
+                "file_type": meta.get("file_type", ""),
+            })
+        return chunks_out
+    else:
+        meta = extract_metadata("", rel, file_profile)
+        raw_texts = _chunk_plain_text(source, chunk_size=chunk_size)
+        return [
+            {
+                "filename": rel,
+                "location": f"{rel}:{idx}",
+                "code": text,
+                "feature": meta.get("feature", ""),
+                "layer": meta.get("layer", ""),
+                "file_type": meta.get("file_type", ""),
+            }
+            for idx, text in enumerate(raw_texts)
+        ]
+
+
+# ---------------------------------------------------------------------------
+# Public entry points
+# ---------------------------------------------------------------------------
+
+
+def run_indexing(
+    collection: chromadb.Collection,
+    config: SourcefireConfig,
+    client: chromadb.ClientAPI | None = None,
+    full: bool = True,
+) -> dict[str, Any]:
+    """Run the indexing pipeline.
+
+    Args:
+        collection: ChromaDB collection to write to.
+        config: Sourcefire configuration.
+        client: ChromaDB client (needed for full reset).
+        full: If True, reset collection and re-index everything.
+              If False, incremental — compare file mtimes and only
+              re-index changed/new files, delete removed files.
+
+    Returns:
+        A stats dict with keys: files, chunks, edges, language, import_edges.
+    """
+    codebase_path = config.project_dir
+    print(f"[pipeline] Scanning codebase at: {codebase_path}")
+
+    language_override = config.language if config.language != "auto" else None
+    profile = get_profile(codebase_path, language_override)
+    lang_name = profile.language if profile else "generic"
+    print(f"[pipeline] Detected language: {lang_name}")
+
+    # Collect all files on disk
+    all_disk_files = _collect_files(codebase_path, config, profile)
+    print(f"[pipeline] Found {len(all_disk_files)} files to index.")
+
+    if not all_disk_files:
+        print("[pipeline] Error: No source files found matching the configured patterns.")
+        print("Run `sourcefire --reinit` to regenerate patterns, or edit .sourcefire/config.toml manually.")
+        return {
+            "files": 0, "chunks": 0, "edges": 0,
+            "language": lang_name, "import_edges": {},
+        }
+
+    # Determine which files to process
+    if full and client:
+        collection = reset_collection(client)
+        print("[pipeline] Collection reset for full re-index.")
+        files_to_index = all_disk_files
+    elif not full:
+        # Incremental: compare mtimes
+        indexed_files = get_indexed_files(collection)
+        stored_mtimes = get_stored_mtimes(collection)
+
+        current_files: dict[str, Path] = {}
+        for f in all_disk_files:
+            rel = f.relative_to(codebase_path).as_posix()
+            current_files[rel] = f
+
+        # Find changed/new files
+        changed: list[Path] = []
+        for rel, f in current_files.items():
+            stored_mtime = stored_mtimes.get(rel, 0.0)
+            if rel not in indexed_files or f.stat().st_mtime > stored_mtime:
+                changed.append(f)
+
+        # Find deleted files
+        deleted = indexed_files - set(current_files.keys())
+        for rel in deleted:
+            delete_file_chunks(collection, rel)
+
+        if not changed and not deleted:
+            print("[pipeline] Index is up to date.")
+            return {
+                "files": len(all_disk_files), "chunks": collection.count(), "edges": 0,
+                "language": lang_name, "import_edges": {},
+            }
+
+        print(f"[pipeline] {len(changed)} changed, {len(deleted)} deleted files.")
+        files_to_index = changed
+
+        # Delete old chunks for changed files before re-inserting
+        for f in changed:
+            rel = f.relative_to(codebase_path).as_posix()
+            delete_file_chunks(collection, rel)
+    else:
+        files_to_index = all_disk_files
+
+    # Produce chunks and collect imports
+    all_chunks: list[dict[str, Any]] = []
+    file_imports: dict[str, list[str]] = {}
+
+    for file_path in files_to_index:
+        rel = file_path.relative_to(codebase_path).as_posix()
+        chunks = _chunks_for_file(file_path, codebase_path, profile, chunk_size=config.chunk_size)
+        all_chunks.extend(chunks)
+
+        file_profile = get_profile_for_extension(file_path.suffix) or profile
+        if file_profile and file_path.suffix in file_profile.file_extensions:
+            source = file_path.read_text(encoding="utf-8", errors="replace")
+            meta = extract_metadata(source, rel, file_profile)
+            if meta.get("imports"):
+                file_imports[rel] = meta["imports"]
+
+    print(f"[pipeline] Produced {len(all_chunks)} chunks.")
+
+    if not all_chunks:
+        return {
+            "files": len(all_disk_files), "chunks": 0, "edges": 0,
+            "language": lang_name, "import_edges": file_imports,
+        }
+
+    # Embed
+    print("[pipeline] Embedding chunks...")
+    texts = [c["code"] for c in all_chunks]
+    embeddings = embed_batch(texts)
+    print("[pipeline] Embeddings done.")
+
+    # Build mtime lookup
+    file_mtimes: dict[str, str] = {}
+    for file_path in files_to_index:
+        rel = file_path.relative_to(codebase_path).as_posix()
+        file_mtimes[rel] = str(file_path.stat().st_mtime)
+
+    # Insert into ChromaDB in batches
+    BATCH_SIZE = 5000
+    print("[pipeline] Inserting into ChromaDB...")
+    for i in range(0, len(all_chunks), BATCH_SIZE):
+        batch = all_chunks[i:i + BATCH_SIZE]
+        batch_emb = embeddings[i:i + BATCH_SIZE]
+        add_chunks(
+            collection,
+            ids=[c["location"] for c in batch],
+            documents=[c["code"] for c in batch],
+            embeddings=batch_emb,
+            metadatas=[
+                {
+                    "filename": c["filename"],
+                    "location": c["location"],
+                    "feature": c["feature"],
+                    "layer": c["layer"],
+                    "file_type": c["file_type"],
+                    "mtime": file_mtimes.get(c["filename"], "0"),
+                }
+                for c in batch
+            ],
+        )
+    print(f"[pipeline] Inserted {len(all_chunks)} chunks.")
+
+    edge_count = sum(len(v) for v in file_imports.values())
+    print(f"[pipeline] Import edges: {edge_count}")
+
+    return {
+        "files": len(all_disk_files),
+        "chunks": len(all_chunks),
+        "edges": edge_count,
+        "language": lang_name,
+        "import_edges": file_imports,
+    }
+
+
+def index_files(
+    collection: chromadb.Collection,
+    file_paths: list[Path],
+    config: SourcefireConfig,
+    profile: LanguageProfile | None,
+) -> dict[str, list[str]]:
+    """Index specific files (for incremental re-indexing by the watcher).
+
+    Deletes existing chunks for each file, then re-chunks, embeds, and inserts.
+    Returns the import map for updated files.
+    """
+    codebase_path = config.project_dir
+    file_imports: dict[str, list[str]] = {}
+    all_chunks: list[dict[str, Any]] = []
+
+    for file_path in file_paths:
+        rel = file_path.relative_to(codebase_path).as_posix()
+        delete_file_chunks(collection, rel)
+
+        chunks = _chunks_for_file(file_path, codebase_path, profile, chunk_size=config.chunk_size)
+        all_chunks.extend(chunks)
+
+        file_profile = get_profile_for_extension(file_path.suffix) or profile
+        if file_profile and file_path.suffix in file_profile.file_extensions:
+            source = file_path.read_text(encoding="utf-8", errors="replace")
+            meta = extract_metadata(source, rel, file_profile)
+            if meta.get("imports"):
+                file_imports[rel] = meta["imports"]
+
+    if all_chunks:
+        texts = [c["code"] for c in all_chunks]
+        embeddings = embed_batch(texts)
+
+        file_mtimes: dict[str, str] = {}
+        for file_path in file_paths:
+            rel = file_path.relative_to(codebase_path).as_posix()
+            file_mtimes[rel] = str(file_path.stat().st_mtime)
+
+        BATCH_SIZE = 5000
+        for i in range(0, len(all_chunks), BATCH_SIZE):
+            batch = all_chunks[i:i + BATCH_SIZE]
+            batch_emb = embeddings[i:i + BATCH_SIZE]
+            add_chunks(
+                collection,
+                ids=[c["location"] for c in batch],
+                documents=[c["code"] for c in batch],
+                embeddings=batch_emb,
+                metadatas=[
+                    {
+                        "filename": c["filename"],
+                        "location": c["location"],
+                        "feature": c["feature"],
+                        "layer": c["layer"],
+                        "file_type": c["file_type"],
+                        "mtime": file_mtimes.get(c["filename"], "0"),
+                    }
+                    for c in batch
+                ],
+            )
+
+    return file_imports