sourcefire 0.2.0__py3-none-any.whl

sourcefire/cli.py

@@ -0,0 +1,293 @@
+"""Sourcefire CLI — single command entry point."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import fcntl
+import os
+import sys
+import webbrowser
+from contextlib import asynccontextmanager
+from datetime import datetime, timezone
+from pathlib import Path
+
+import uvicorn
+from dotenv import load_dotenv
+from fastapi import FastAPI
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+
+load_dotenv()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="sourcefire",
+        description="Sourcefire — AI-powered codebase RAG from your terminal",
+    )
+    parser.add_argument("--port", type=int, default=None, help="Server port (default: from config or 8000)")
+    parser.add_argument("--no-open", action="store_true", help="Don't auto-open browser")
+    parser.add_argument("--reinit", action="store_true", help="Regenerate .sourcefire/config.toml via LLM")
+    parser.add_argument("--verbose", action="store_true", help="Verbose logging")
+    return parser.parse_args()
+
+
+def discover_project() -> tuple[Path, Path]:
+    """Walk up from cwd to find .sourcefire/, like git finds .git/.
+
+    Returns (project_dir, sourcefire_dir).
+    If not found, returns (cwd, cwd/.sourcefire).
+    """
+    current = Path.cwd().resolve()
+    while True:
+        candidate = current / ".sourcefire"
+        if candidate.is_dir():
+            return current, candidate
+        parent = current.parent
+        if parent == current:
+            break
+        current = parent
+
+    cwd = Path.cwd().resolve()
+    return cwd, cwd / ".sourcefire"
+
+
+def acquire_lock(lock_path: Path) -> int | None:
+    """Acquire an exclusive file lock. Returns fd on success, None on failure."""
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR)
+        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        return fd
+    except (OSError, BlockingIOError):
+        return None
+
+
+def release_lock(fd: int, lock_path: Path) -> None:
+    """Release the file lock."""
+    try:
+        fcntl.flock(fd, fcntl.LOCK_UN)
+        os.close(fd)
+    except OSError:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# App state (shared between main() and lifespan)
+# ---------------------------------------------------------------------------
+
+_app_state: dict = {}
+
+
+# ---------------------------------------------------------------------------
+# Lifespan
+# ---------------------------------------------------------------------------
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Start-up: index codebase, build graph, start watcher."""
+    config = _app_state["config"]
+    project_dir = _app_state["project_dir"]
+    api_key = _app_state["api_key"]
+    args = _app_state["args"]
+
+    from sourcefire.db import create_client, get_collection
+    from sourcefire.indexer.language_profiles import get_profile
+    from sourcefire.indexer.pipeline import run_indexing
+    from sourcefire.retriever.graph import ImportGraph
+    from sourcefire.api.routes import init_dependencies
+    from sourcefire.watcher import watch_and_reindex
+
+    print(f"[sourcefire] Project: {project_dir.name}")
+    print(f"[sourcefire] Config: {config.config_path}")
+
+    # Detect language
+    language_override = config.language if config.language != "auto" else None
+    profile = get_profile(project_dir, language_override)
+    lang_name = profile.language if profile else "generic"
+    print(f"[sourcefire] Language: {lang_name}")
+
+    # Create ChromaDB client
+    client = create_client(config.chroma_dir)
+    collection = get_collection(client)
+
+    # Determine if this is a first run (empty collection)
+    existing_count = collection.count()
+    is_first_run = existing_count == 0
+
+    # Run indexing
+    if is_first_run:
+        print("[sourcefire] First run — full index...")
+        stats = run_indexing(collection, config, client=client, full=True)
+    else:
+        print("[sourcefire] Checking for changes...")
+        stats = run_indexing(collection, config, client=client, full=False)
+
+    print(f"[sourcefire] Indexed: {stats['files']} files, {stats['chunks']} chunks")
+
+    # Build import graph
+    external_prefixes = profile.external_import_prefixes if profile else ()
+    graph = ImportGraph(external_prefixes=external_prefixes)
+
+    import_edges = stats.get("import_edges", {})
+    if import_edges:
+        for source_file, imports in import_edges.items():
+            for imp in imports:
+                resolved = ImportGraph._resolve_import(source_file, imp)
+                graph.add_edge(source_file, resolved)
+    elif config.graph_path.is_file():
+        graph = ImportGraph.load(config.graph_path, external_prefixes=external_prefixes)
+
+    print(f"[sourcefire] Import graph: {graph.node_count} nodes")
+
+    # Build index status
+    index_status = {
+        "files_indexed": stats.get("files", 0),
+        "last_indexed": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "index_status": "ready",
+        "language": lang_name,
+    }
+
+    # Inject dependencies into routes
+    init_dependencies(collection, graph, index_status, profile, project_dir, api_key)
+
+    # Start file watcher
+    watcher_task = asyncio.create_task(
+        watch_and_reindex(config, collection, graph, profile)
+    )
+
+    # Open browser
+    url = f"http://{config.host}:{config.port}"
+    print(f"[sourcefire] Ready — {url}")
+    if not args.no_open:
+        webbrowser.open(url)
+
+    yield
+
+    # Shutdown
+    print("[sourcefire] Shutting down...")
+    watcher_task.cancel()
+    try:
+        await watcher_task
+    except asyncio.CancelledError:
+        pass
+
+    # Save graph
+    graph.save(config.graph_path)
+    print("[sourcefire] Graph saved.")
+
+
+# ---------------------------------------------------------------------------
+# App
+# ---------------------------------------------------------------------------
+
+from importlib.resources import files as _resource_files
+
+_static_dir = str(Path(_resource_files("sourcefire")) / "static")
+
+app = FastAPI(
+    title="Sourcefire",
+    description="AI-powered codebase RAG. Created by Athar Wani.",
+    version="0.2.0",
+    lifespan=lifespan,
+)
+
+from sourcefire.api.routes import router  # noqa: E402
+
+app.include_router(router)
+app.mount("/static", StaticFiles(directory=_static_dir), name="static")
+
+
+@app.get("/", include_in_schema=False)
+async def root() -> FileResponse:
+    return FileResponse(os.path.join(_static_dir, "index.html"))
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    """Sourcefire CLI entry point."""
+    args = parse_args()
+
+    project_dir, sourcefire_dir = discover_project()
+
+    # Acquire lock
+    lock_fd = acquire_lock(sourcefire_dir / ".lock")
+    if lock_fd is None:
+        print("Error: Another sourcefire instance is already running for this project.")
+        sys.exit(1)
+
+    # Check for API key — prompt interactively if missing
+    api_key = os.getenv("GEMINI_API_KEY", "")
+    if not api_key:
+        print("No GEMINI_API_KEY found in environment.")
+        try:
+            api_key = input("Enter your Gemini API key: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nAborted.")
+            release_lock(lock_fd, sourcefire_dir / ".lock")
+            sys.exit(1)
+
+        if not api_key:
+            print("Error: A Gemini API key is required to run Sourcefire.")
+            release_lock(lock_fd, sourcefire_dir / ".lock")
+            sys.exit(1)
+
+        # Persist to .env in project root
+        env_path = project_dir / ".env"
+        with open(env_path, "a") as f:
+            f.write(f"\nGEMINI_API_KEY={api_key}\n")
+        os.environ["GEMINI_API_KEY"] = api_key
+        print(f"API key saved to {env_path}")
+
+    # Auto-init or reinit
+    needs_init = not sourcefire_dir.exists() or not (sourcefire_dir / "config.toml").exists()
+
+    if needs_init:
+        from sourcefire.init import auto_init
+        config = auto_init(
+            project_dir=project_dir,
+            sourcefire_dir=sourcefire_dir,
+            api_key=api_key,
+        )
+    elif args.reinit:
+        from sourcefire.config import load_config
+        from sourcefire.init import reinit_patterns
+        config = load_config(project_dir, sourcefire_dir)
+        config = reinit_patterns(config, api_key=api_key)
+    else:
+        from sourcefire.config import load_config
+        config = load_config(project_dir, sourcefire_dir)
+
+    # Override port from CLI
+    if args.port:
+        config.port = args.port
+
+    # Store state for lifespan access
+    _app_state["config"] = config
+    _app_state["project_dir"] = project_dir
+    _app_state["sourcefire_dir"] = sourcefire_dir
+    _app_state["api_key"] = api_key
+    _app_state["args"] = args
+    _app_state["lock_fd"] = lock_fd
+
+    # Run server
+    try:
+        uvicorn.run(
+            "sourcefire.cli:app",
+            host=config.host,
+            port=config.port,
+            reload=False,
+            log_level="info" if args.verbose else "warning",
+        )
+    finally:
+        release_lock(lock_fd, sourcefire_dir / ".lock")
+
+
+if __name__ == "__main__":
+    main()

sourcefire/config.py

@@ -0,0 +1,148 @@
+"""Configuration for Sourcefire — loaded from .sourcefire/config.toml."""
+
+from __future__ import annotations
+
+import os
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import tomli_w
+
+
+@dataclass
+class SourcefireConfig:
+    """All Sourcefire configuration for a project."""
+
+    # Resolved at runtime, not stored in TOML
+    project_dir: Path = field(default_factory=Path.cwd)
+    sourcefire_dir: Path = field(default_factory=lambda: Path.cwd() / ".sourcefire")
+
+    # [project]
+    project_name: str = ""
+    language: str = "auto"
+
+    # [indexer]
+    include: list[str] = field(default_factory=list)
+    exclude: list[str] = field(default_factory=list)
+    chunk_size: int = 1000
+    chunk_overlap: int = 300
+
+    # [llm]
+    provider: str = "gemini"
+    model: str = "gemini-2.5-flash"
+    api_key_env: str = "GEMINI_API_KEY"
+
+    # [server]
+    host: str = "127.0.0.1"
+    port: int = 8000
+
+    # [retrieval]
+    top_k: int = 8
+    relevance_threshold: float = 0.3
+
+    # Versioning
+    config_version: int = 1
+
+    @property
+    def gemini_api_key(self) -> str:
+        return os.getenv(self.api_key_env, "")
+
+    @property
+    def chroma_dir(self) -> Path:
+        return self.sourcefire_dir / "chroma"
+
+    @property
+    def graph_path(self) -> Path:
+        return self.sourcefire_dir / "graph.json"
+
+    @property
+    def config_path(self) -> Path:
+        return self.sourcefire_dir / "config.toml"
+
+    @property
+    def lock_path(self) -> Path:
+        return self.sourcefire_dir / ".lock"
+
+
+# Constants used by other modules
+EMBEDDING_MODEL: str = "sentence-transformers/all-MiniLM-L6-v2"
+MAX_TOKEN_BUDGET: dict[str, int] = {
+    "gemini-2.5-flash": 100_000,
+    "gemini-2.5-pro": 200_000,
+}
+MAX_HISTORY_PAIRS: int = 5
+RESPONSE_HEADROOM: int = 8_000
+
+
+def default_config(project_dir: Path) -> SourcefireConfig:
+    """Return a SourcefireConfig with sensible defaults for the given project."""
+    return SourcefireConfig(
+        project_dir=project_dir,
+        sourcefire_dir=project_dir / ".sourcefire",
+        project_name=project_dir.name,
+    )
+
+
+def load_config(project_dir: Path, sourcefire_dir: Path) -> SourcefireConfig:
+    """Load config from .sourcefire/config.toml."""
+    config_path = sourcefire_dir / "config.toml"
+    raw = config_path.read_text(encoding="utf-8")
+    data = tomllib.loads(raw)
+
+    project = data.get("project", {})
+    indexer = data.get("indexer", {})
+    llm = data.get("llm", {})
+    server = data.get("server", {})
+    retrieval = data.get("retrieval", {})
+
+    return SourcefireConfig(
+        project_dir=project_dir,
+        sourcefire_dir=sourcefire_dir,
+        config_version=data.get("config_version", 1),
+        project_name=project.get("name", project_dir.name),
+        language=project.get("language", "auto"),
+        include=indexer.get("include", []),
+        exclude=indexer.get("exclude", []),
+        chunk_size=indexer.get("chunk_size", 1000),
+        chunk_overlap=indexer.get("chunk_overlap", 300),
+        provider=llm.get("provider", "gemini"),
+        model=llm.get("model", "gemini-2.5-flash"),
+        api_key_env=llm.get("api_key_env", "GEMINI_API_KEY"),
+        host=server.get("host", "127.0.0.1"),
+        port=server.get("port", 8000),
+        top_k=retrieval.get("top_k", 8),
+        relevance_threshold=retrieval.get("relevance_threshold", 0.3),
+    )
+
+
+def save_config(config: SourcefireConfig) -> None:
+    """Write config to .sourcefire/config.toml."""
+    data = {
+        "config_version": config.config_version,
+        "project": {
+            "name": config.project_name,
+            "language": config.language,
+        },
+        "indexer": {
+            "include": config.include,
+            "exclude": config.exclude,
+            "chunk_size": config.chunk_size,
+            "chunk_overlap": config.chunk_overlap,
+        },
+        "llm": {
+            "provider": config.provider,
+            "model": config.model,
+            "api_key_env": config.api_key_env,
+        },
+        "server": {
+            "host": config.host,
+            "port": config.port,
+        },
+        "retrieval": {
+            "top_k": config.top_k,
+            "relevance_threshold": config.relevance_threshold,
+        },
+    }
+    config.config_path.parent.mkdir(parents=True, exist_ok=True)
+    config.config_path.write_text(tomli_w.dumps(data), encoding="utf-8")

sourcefire/db.py

@@ -0,0 +1,196 @@
+"""ChromaDB wrapper for Sourcefire — async-safe via run_in_executor."""
+
+from __future__ import annotations
+
+import asyncio
+from functools import partial
+from pathlib import Path
+from typing import Any
+
+import chromadb
+
+COLLECTION_NAME = "code_chunks"
+
+
+def create_client(chroma_dir: Path) -> chromadb.ClientAPI:
+    """Create a persistent ChromaDB client."""
+    chroma_dir.mkdir(parents=True, exist_ok=True)
+    return chromadb.PersistentClient(path=str(chroma_dir))
+
+
+def get_collection(client: chromadb.ClientAPI) -> chromadb.Collection:
+    """Get or create the code_chunks collection."""
+    return client.get_or_create_collection(
+        name=COLLECTION_NAME,
+        metadata={"hnsw:space": "cosine"},
+    )
+
+
+def reset_collection(client: chromadb.ClientAPI) -> chromadb.Collection:
+    """Delete and recreate the collection (for full re-index)."""
+    try:
+        client.delete_collection(COLLECTION_NAME)
+    except ValueError:
+        pass
+    return get_collection(client)
+
+
+# ---------------------------------------------------------------------------
+# Sync operations
+# ---------------------------------------------------------------------------
+
+
+def add_chunks(
+    collection: chromadb.Collection,
+    ids: list[str],
+    documents: list[str],
+    embeddings: list[list[float]],
+    metadatas: list[dict[str, str]],
+) -> None:
+    """Add chunks to the collection."""
+    collection.add(
+        ids=ids,
+        documents=documents,
+        embeddings=embeddings,
+        metadatas=metadatas,
+    )
+
+
+def delete_file_chunks(collection: chromadb.Collection, filename: str) -> None:
+    """Delete all chunks for a given filename."""
+    collection.delete(where={"filename": filename})
+
+
+def query_similar(
+    collection: chromadb.Collection,
+    query_embedding: list[float],
+    n_results: int = 8,
+    where: dict | None = None,
+) -> list[dict[str, Any]]:
+    """Query for similar chunks. Returns list of result dicts."""
+    kwargs: dict[str, Any] = {
+        "query_embeddings": [query_embedding],
+        "n_results": n_results,
+        "include": ["documents", "metadatas", "distances"],
+    }
+    if where:
+        kwargs["where"] = where
+
+    results = collection.query(**kwargs)
+
+    rows: list[dict[str, Any]] = []
+    if not results["ids"] or not results["ids"][0]:
+        return rows
+
+    for i, doc_id in enumerate(results["ids"][0]):
+        meta = results["metadatas"][0][i] if results["metadatas"] else {}
+        distance = results["distances"][0][i] if results["distances"] else 1.0
+        relevance = 1.0 - distance  # cosine distance -> similarity
+
+        rows.append({
+            "filename": meta.get("filename", ""),
+            "location": meta.get("location", ""),
+            "code": results["documents"][0][i] if results["documents"] else "",
+            "feature": meta.get("feature", ""),
+            "layer": meta.get("layer", ""),
+            "file_type": meta.get("file_type", ""),
+            "relevance": relevance,
+        })
+
+    return rows
+
+
+def get_chunks_by_files(
+    collection: chromadb.Collection,
+    filenames: list[str],
+) -> list[dict[str, Any]]:
+    """Retrieve all chunks for the given filenames."""
+    if not filenames:
+        return []
+
+    results = collection.get(
+        where={"filename": {"$in": filenames}},
+        include=["documents", "metadatas"],
+    )
+
+    rows: list[dict[str, Any]] = []
+    if not results["ids"]:
+        return rows
+
+    for i, doc_id in enumerate(results["ids"]):
+        meta = results["metadatas"][i] if results["metadatas"] else {}
+        rows.append({
+            "filename": meta.get("filename", ""),
+            "location": meta.get("location", ""),
+            "code": results["documents"][i] if results["documents"] else "",
+            "feature": meta.get("feature", ""),
+            "layer": meta.get("layer", ""),
+            "file_type": meta.get("file_type", ""),
+        })
+
+    return rows
+
+
+def get_indexed_files(collection: chromadb.Collection) -> set[str]:
+    """Return set of all filenames currently in the collection."""
+    results = collection.get(include=["metadatas"])
+    files: set[str] = set()
+    if results["metadatas"]:
+        for meta in results["metadatas"]:
+            if meta and "filename" in meta:
+                files.add(meta["filename"])
+    return files
+
+
+def get_stored_mtimes(collection: chromadb.Collection) -> dict[str, float]:
+    """Get stored mtimes for all indexed files from ChromaDB metadata."""
+    results = collection.get(include=["metadatas"])
+    mtimes: dict[str, float] = {}
+    if results["metadatas"]:
+        for meta in results["metadatas"]:
+            if meta and "filename" in meta and "mtime" in meta:
+                try:
+                    mtimes[meta["filename"]] = float(meta["mtime"])
+                except (ValueError, TypeError):
+                    pass
+    return mtimes
+
+
+# ---------------------------------------------------------------------------
+# Async wrappers (for use in FastAPI routes / RAG chain)
+# ---------------------------------------------------------------------------
+
+
+async def async_query_similar(
+    collection: chromadb.Collection,
+    query_embedding: list[float],
+    n_results: int = 8,
+    where: dict | None = None,
+) -> list[dict[str, Any]]:
+    """Async wrapper for query_similar."""
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(
+        None, partial(query_similar, collection, query_embedding, n_results, where)
+    )
+
+
+async def async_get_chunks_by_files(
+    collection: chromadb.Collection,
+    filenames: list[str],
+) -> list[dict[str, Any]]:
+    """Async wrapper for get_chunks_by_files."""
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(
+        None, partial(get_chunks_by_files, collection, filenames)
+    )
+
+
+async def async_delete_file_chunks(
+    collection: chromadb.Collection,
+    filename: str,
+) -> None:
+    """Async wrapper for delete_file_chunks."""
+    loop = asyncio.get_event_loop()
+    await loop.run_in_executor(
+        None, partial(delete_file_chunks, collection, filename)
+    )

sourcefire/indexer/embeddings.py

@@ -0,0 +1,27 @@
+"""Shared embedding module using sentence-transformers directly."""
+
+from sentence_transformers import SentenceTransformer
+from sourcefire.config import EMBEDDING_MODEL
+
+_model: SentenceTransformer | None = None
+
+
+def get_model() -> SentenceTransformer:
+    global _model
+    if _model is None:
+        print(f"Loading embedding model: {EMBEDDING_MODEL}...")
+        _model = SentenceTransformer(EMBEDDING_MODEL)
+        print("Embedding model loaded.")
+    return _model
+
+
+def embed_text(text: str) -> list[float]:
+    """Embed a single text string. Returns a list of floats."""
+    model = get_model()
+    return model.encode(text).tolist()
+
+
+def embed_batch(texts: list[str]) -> list[list[float]]:
+    """Embed a batch of texts. More efficient than calling embed_text in a loop."""
+    model = get_model()
+    return model.encode(texts).tolist()