sourcefire 0.2.0__py3-none-any.whl

sourcefire/init.py

@@ -0,0 +1,189 @@
+"""Auto-initialization for Sourcefire — creates .sourcefire/ with LLM-generated config."""
+
+from __future__ import annotations
+
+import os
+import re
+import tomllib
+from pathlib import Path
+
+from sourcefire.config import SourcefireConfig, default_config, save_config
+from sourcefire.indexer.language_profiles import get_profile
+
+
+def scan_file_tree(project_dir: Path, max_files: int = 5000) -> str:
+    """Scan the project directory and return a text representation of the file tree."""
+    skip_dirs = {
+        ".git", "node_modules", "__pycache__", "build", "dist", "target",
+        ".dart_tool", ".next", "venv", ".venv", ".idea", ".vs", ".sourcefire",
+        ".tox", "eggs",
+    }
+
+    lines: list[str] = []
+    file_count = 0
+
+    for root, dirs, files in os.walk(project_dir):
+        dirs[:] = [d for d in dirs if d not in skip_dirs and not d.startswith(".")]
+
+        rel_root = Path(root).relative_to(project_dir).as_posix()
+        if rel_root == ".":
+            rel_root = ""
+
+        for f in sorted(files):
+            if file_count >= max_files:
+                break
+            rel_path = f"{rel_root}/{f}" if rel_root else f
+            lines.append(rel_path)
+            file_count += 1
+
+        if file_count >= max_files:
+            break
+
+    return "\n".join(lines)
+
+
+def _generate_patterns_via_llm(file_tree: str, api_key: str) -> dict[str, list[str]] | None:
+    """Ask the LLM to generate include/exclude patterns from the file tree."""
+    try:
+        from langchain_google_genai import ChatGoogleGenerativeAI
+        from langchain_core.messages import HumanMessage
+
+        llm = ChatGoogleGenerativeAI(
+            model="gemini-2.5-flash",
+            google_api_key=api_key,
+        )
+
+        prompt = (
+            "Given this project file tree, determine which files are source code worth "
+            "indexing for a code RAG system. Respond with ONLY a TOML code block containing "
+            "two arrays: `include` (glob patterns for source files, configs, and docs) and "
+            "`exclude` (glob patterns for build artifacts, dependencies, generated files, "
+            "and non-code assets). Be comprehensive but conservative.\n\n"
+            "Always include these in exclude: .git/**, .sourcefire/**\n\n"
+            f"```\n{file_tree}\n```"
+        )
+
+        response = llm.invoke([HumanMessage(content=prompt)])
+        text = response.content if hasattr(response, "content") else str(response)
+
+        match = re.search(r"```(?:toml)?\s*\n(.*?)\n```", text, re.DOTALL)
+        if not match:
+            return None
+
+        toml_str = match.group(1)
+        data = tomllib.loads(toml_str)
+
+        include = data.get("include", [])
+        exclude = data.get("exclude", [])
+
+        if isinstance(include, list) and isinstance(exclude, list):
+            if ".sourcefire/**" not in exclude:
+                exclude.append(".sourcefire/**")
+            return {"include": include, "exclude": exclude}
+
+        return None
+
+    except Exception as exc:
+        print(f"[init] LLM config generation failed: {exc}")
+        return None
+
+
+def _fallback_patterns(project_dir: Path, language_override: str | None = None) -> dict[str, list[str]]:
+    """Generate patterns from language profile when LLM is unavailable."""
+    profile = get_profile(project_dir, language_override)
+
+    if profile:
+        include = list(profile.include_patterns)
+        exclude = list(profile.exclude_patterns)
+    else:
+        include = ["**/*"]
+        exclude = []
+
+    for pat in [".git/**", ".sourcefire/**", "node_modules/**", "__pycache__/**",
+                "*.pyc", ".venv/**", "venv/**", "dist/**", "build/**"]:
+        if pat not in exclude:
+            exclude.append(pat)
+
+    for pat in ["README.md", "CLAUDE.md"]:
+        if pat not in include:
+            include.append(pat)
+
+    return {"include": include, "exclude": exclude}
+
+
+def auto_init(
+    project_dir: Path,
+    sourcefire_dir: Path | None = None,
+    api_key: str = "",
+    language_override: str | None = None,
+) -> SourcefireConfig:
+    """Initialize .sourcefire/ directory with LLM-generated config."""
+    if sourcefire_dir is None:
+        sourcefire_dir = project_dir / ".sourcefire"
+
+    print(f"[init] Initializing Sourcefire for: {project_dir.name}")
+
+    sourcefire_dir.mkdir(parents=True, exist_ok=True)
+
+    print("[init] Scanning project structure...")
+    file_tree = scan_file_tree(project_dir)
+
+    patterns: dict[str, list[str]] | None = None
+    if api_key:
+        print("[init] Generating config via LLM...")
+        patterns = _generate_patterns_via_llm(file_tree, api_key)
+
+    if patterns:
+        print(f"[init] LLM generated {len(patterns['include'])} include, {len(patterns['exclude'])} exclude patterns.")
+    else:
+        print("[init] Using language-profile defaults for patterns.")
+        patterns = _fallback_patterns(project_dir, language_override)
+
+    profile = get_profile(project_dir, language_override)
+    language = profile.language if profile else "auto"
+
+    config = default_config(project_dir)
+    config.sourcefire_dir = sourcefire_dir
+    config.include = patterns["include"]
+    config.exclude = patterns["exclude"]
+    config.language = language
+
+    save_config(config)
+    print(f"[init] Config written to: {config.config_path}")
+
+    print("\nTip: Add to your .gitignore:")
+    print("  .sourcefire/chroma/")
+    print("  .sourcefire/graph.json")
+    print("  .sourcefire/.lock\n")
+
+    return config
+
+
+def reinit_patterns(
+    config: SourcefireConfig,
+    api_key: str = "",
+) -> SourcefireConfig:
+    """Regenerate only the [indexer] include/exclude patterns, preserving other config."""
+    print(f"[init] Regenerating patterns for: {config.project_dir.name}")
+
+    file_tree = scan_file_tree(config.project_dir)
+
+    patterns: dict[str, list[str]] | None = None
+    if api_key:
+        print("[init] Generating patterns via LLM...")
+        patterns = _generate_patterns_via_llm(file_tree, api_key)
+
+    if patterns:
+        print(f"[init] LLM generated {len(patterns['include'])} include, {len(patterns['exclude'])} exclude patterns.")
+    else:
+        print("[init] Using language-profile defaults.")
+        language_override = config.language if config.language != "auto" else None
+        patterns = _fallback_patterns(config.project_dir, language_override)
+
+    config.include = patterns["include"]
+    config.exclude = patterns["exclude"]
+
+    save_config(config)
+    print(f"[init] Updated patterns in: {config.config_path}")
+
+    return config

sourcefire/prompts/system.md

@@ -0,0 +1,28 @@
+You are Sourcefire — an expert codebase guide. You answer questions about the codebase accurately and concisely.
+
+## How you work
+
+You receive two types of information:
+
+1. **Retrieved Context** (in the user message under "Retrieved Code Context") — these are semantically relevant code chunks pulled from a vector database. They are your PRIMARY source of truth. Trust them. Base your answers on what you see in these chunks first.
+
+2. **Tools** — you have tools to read files, search code, query the vector DB, trace imports, and check git history. Use them ONLY when the retrieved context is insufficient — for example when you need to see a full file, trace a call chain deeper, or verify something not covered by the chunks.
+
+## Rules
+
+- NEVER invent file paths, function names, or code that isn't in the retrieved context or tool results. If you don't know, say so and use a tool to find out.
+- NEVER mix in knowledge from other projects. Only reference code you can see.
+- When referencing a file, format it as `[filename](file://path)` so the UI can make it clickable.
+- Show full file paths so the developer can navigate to the source.
+- Trace connections between files — show the chain, not just the endpoint.
+- Be concise. Lead with the answer, then support with code references.
+
+## Tool strategy
+
+1. First, answer from the retrieved context — it's fast, semantic, and already relevant.
+2. If the context is partial or you need more detail, use `semantic_code_search` to find related code by meaning.
+3. If you need exact code, use `read_local_file` or `read_lines`.
+4. If you need to find where something is defined or used, use `find_definition` or `find_references`.
+5. If you need project structure, use `get_file_structure` or `find_files_by_name`.
+6. For history/blame, use git tools.
+7. Do NOT call tools for information already present in the retrieved context.

sourcefire/retriever/graph.py

@@ -0,0 +1,162 @@
+"""Bidirectional import graph for source files."""
+
+from __future__ import annotations
+
+import json
+import posixpath
+from collections import defaultdict, deque
+from pathlib import Path, PurePosixPath
+from typing import ClassVar
+
+
+class ImportGraph:
+    """Directed graph that tracks import relationships between source files.
+
+    Edges are stored in both forward (imports) and reverse (importers)
+    directions so either direction can be queried in O(1).
+    """
+
+    # Default external schemes — overridden at construction time by the profile
+    _external_prefixes: tuple[str, ...] = ()
+
+    def __init__(self, external_prefixes: tuple[str, ...] = ()) -> None:
+        self._forward: defaultdict[str, set[str]] = defaultdict(set)
+        self._reverse: defaultdict[str, set[str]] = defaultdict(set)
+        self._external_prefixes = external_prefixes
+
+    # ------------------------------------------------------------------
+    # Mutation
+    # ------------------------------------------------------------------
+
+    def add_edge(self, source: str, target: str) -> None:
+        """Record that *source* imports *target*."""
+        self._forward[source].add(target)
+        self._reverse[target].add(source)
+
+    # ------------------------------------------------------------------
+    # Queries
+    # ------------------------------------------------------------------
+
+    def get_imports(self, file_path: str) -> list[str]:
+        """Return files that *file_path* imports (forward edges)."""
+        return list(self._forward.get(file_path, []))
+
+    def get_importers(self, file_path: str) -> list[str]:
+        """Return files that import *file_path* (reverse edges)."""
+        return list(self._reverse.get(file_path, []))
+
+    def get_neighbors(self, file_path: str, hops: int = 1) -> list[str]:
+        """Return all files reachable from *file_path* within *hops* steps.
+
+        Traversal goes in *both* directions (imports and importers).
+        The seed file itself is excluded from the result.
+        """
+        visited: set[str] = {file_path}
+        queue: deque[tuple[str, int]] = deque([(file_path, 0)])
+
+        while queue:
+            current, depth = queue.popleft()
+            if depth >= hops:
+                continue
+            for neighbor in (*self._forward.get(current, ()), *self._reverse.get(current, ())):
+                if neighbor not in visited:
+                    visited.add(neighbor)
+                    queue.append((neighbor, depth + 1))
+
+        visited.discard(file_path)
+        return list(visited)
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def node_count(self) -> int:
+        """Number of unique nodes (files) in the graph."""
+        nodes: set[str] = set(self._forward.keys()) | set(self._reverse.keys())
+        return len(nodes)
+
+    # ------------------------------------------------------------------
+    # Construction helpers
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_import_map(
+        cls,
+        file_imports: dict[str, list[str]],
+        base_dir: str = "",
+        external_prefixes: tuple[str, ...] = (),
+    ) -> "ImportGraph":
+        """Build an ImportGraph from a mapping of ``{file: [import_strings]}``.
+
+        - Imports starting with any prefix in *external_prefixes* are skipped.
+        - Relative imports (``../foo``, ``./bar``) are resolved relative to
+          the importing file's directory.
+
+        Args:
+            file_imports: Mapping from source file path to its raw import list.
+            base_dir: Optional prefix (unused, kept for forward-compatibility).
+            external_prefixes: Prefixes that indicate external packages.
+        """
+        graph = cls(external_prefixes=external_prefixes)
+        for source_file, imports in file_imports.items():
+            for raw_import in imports:
+                if external_prefixes and any(raw_import.startswith(scheme) for scheme in external_prefixes):
+                    continue
+                resolved = cls._resolve_import(source_file, raw_import)
+                graph.add_edge(source_file, resolved)
+        return graph
+
+    @staticmethod
+    def _resolve_import(source_file: str, relative_import: str) -> str:
+        """Resolve *relative_import* relative to *source_file*'s directory."""
+        source_dir = str(PurePosixPath(source_file).parent)
+        joined = posixpath.join(source_dir, relative_import)
+        return posixpath.normpath(joined)
+
+    # ------------------------------------------------------------------
+    # File removal (for incremental re-index)
+    # ------------------------------------------------------------------
+
+    def remove_file(self, file_path: str) -> None:
+        """Remove all edges involving *file_path*."""
+        if file_path in self._forward:
+            for target in self._forward[file_path]:
+                self._reverse[target].discard(file_path)
+            del self._forward[file_path]
+        if file_path in self._reverse:
+            for source in self._reverse[file_path]:
+                self._forward[source].discard(file_path)
+            del self._reverse[file_path]
+
+    # ------------------------------------------------------------------
+    # JSON persistence
+    # ------------------------------------------------------------------
+
+    def to_dict(self) -> dict:
+        """Serialize to a dict for JSON storage."""
+        edges = []
+        for source, targets in self._forward.items():
+            for target in targets:
+                edges.append({"source": source, "target": target})
+        return {"edges": edges}
+
+    @classmethod
+    def from_dict(cls, data: dict, external_prefixes: tuple[str, ...] = ()) -> "ImportGraph":
+        """Deserialize from a dict."""
+        graph = cls(external_prefixes=external_prefixes)
+        for edge in data.get("edges", []):
+            graph.add_edge(edge["source"], edge["target"])
+        return graph
+
+    def save(self, path: Path) -> None:
+        """Save graph to a JSON file."""
+        path.write_text(json.dumps(self.to_dict(), indent=2), encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: Path, external_prefixes: tuple[str, ...] = ()) -> "ImportGraph":
+        """Load graph from a JSON file."""
+        if not path.is_file():
+            return cls(external_prefixes=external_prefixes)
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return cls.from_dict(data, external_prefixes=external_prefixes)

sourcefire/retriever/search.py

@@ -0,0 +1,86 @@
+"""Vector search retriever for Sourcefire.
+
+Provides:
+- parse_file_references: extract file references from stack traces and error
+  messages, driven by language profile patterns.
+- semantic_search: cosine similarity search against ChromaDB.
+- get_chunks_by_filenames: retrieve all chunks for a specific set of files.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from sourcefire.db import async_query_similar, async_get_chunks_by_files
+
+
+# ---------------------------------------------------------------------------
+# Stack trace / error message parsing
+# ---------------------------------------------------------------------------
+
+_COMPILED_FILE_REF_PATTERNS: dict[str, list[re.Pattern]] = {}
+
+
+def parse_file_references(text: str, file_ref_patterns: list[str] | None = None) -> list[dict[str, Any]]:
+    """Extract file references from stack traces and error messages."""
+    if not file_ref_patterns:
+        file_ref_patterns = [r"\b([\w./\\-]+\.\w+)(?::(\d+))?"]
+
+    cache_key = str(file_ref_patterns)
+    if cache_key not in _COMPILED_FILE_REF_PATTERNS:
+        _COMPILED_FILE_REF_PATTERNS[cache_key] = [re.compile(p) for p in file_ref_patterns]
+
+    compiled = _COMPILED_FILE_REF_PATTERNS[cache_key]
+
+    results: list[dict[str, Any]] = []
+    seen: set[tuple[str, int]] = set()
+
+    for regex in compiled:
+        for m in regex.finditer(text):
+            raw_path = m.group(1) if m.group(1) else ""
+            if not raw_path:
+                continue
+            line = int(m.group(2)) if m.lastindex and m.lastindex >= 2 and m.group(2) else 0
+            key = (raw_path, line)
+            if key not in seen:
+                seen.add(key)
+                results.append({"file": raw_path, "line": line})
+
+    return results
+
+
+# ---------------------------------------------------------------------------
+# ChromaDB search functions
+# ---------------------------------------------------------------------------
+
+
+async def semantic_search(
+    collection: Any,
+    query_vector: list[float],
+    top_k: int = 8,
+    threshold: float = 0.3,
+    feature: str | None = None,
+    filenames: list[str] | None = None,
+) -> list[dict[str, Any]]:
+    """Cosine similarity search against ChromaDB."""
+    where: dict | None = None
+
+    if feature and filenames:
+        where = {"$and": [{"feature": feature}, {"filename": {"$in": filenames}}]}
+    elif feature:
+        where = {"feature": feature}
+    elif filenames:
+        where = {"filename": {"$in": filenames}}
+
+    rows = await async_query_similar(collection, query_vector, n_results=top_k, where=where)
+
+    return [r for r in rows if r.get("relevance", 0) >= threshold]
+
+
+async def get_chunks_by_filenames(
+    collection: Any,
+    filenames: list[str],
+) -> list[dict[str, Any]]:
+    """Retrieve all chunks for the given file paths."""
+    return await async_get_chunks_by_files(collection, filenames)