sourcefire 0.2.0__py3-none-any.whl

sourcefire/api/models.py

@@ -0,0 +1,24 @@
+"""Pydantic request/response models for the Sourcefire API."""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class QueryRequest(BaseModel):
+    query: str
+    mode: Literal["debug", "feature", "explain"] = "debug"
+    model: Literal["gemini-2.5-flash", "gemini-2.5-pro"] = "gemini-2.5-flash"
+    history: list[dict] = []
+
+
+class StatusResponse(BaseModel):
+    files_indexed: int
+    last_indexed: str
+    index_status: str
+    language: str = "generic"
+
+
+class SourceResponse(BaseModel):
+    content: str
+    language: str

sourcefire/api/routes.py

@@ -0,0 +1,166 @@
+"""FastAPI router for the Sourcefire API."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, AsyncGenerator
+
+from fastapi import APIRouter, HTTPException, Query
+from sse_starlette.sse import EventSourceResponse
+
+from sourcefire.api.models import QueryRequest, SourceResponse, StatusResponse
+
+router = APIRouter(prefix="/api")
+
+# ---------------------------------------------------------------------------
+# Module-level dependency state — set once at startup via init_dependencies()
+# ---------------------------------------------------------------------------
+
+_collection: Any = None
+_graph: Any = None
+_profile: Any = None
+_project_dir: Path | None = None
+_gemini_api_key: str = ""
+_index_status: dict[str, Any] = {
+    "files_indexed": 0,
+    "last_indexed": "never",
+    "index_status": "not_ready",
+    "language": "generic",
+}
+
+
+def init_dependencies(
+    collection: Any,
+    graph: Any,
+    index_status: dict[str, Any],
+    profile: Any = None,
+    project_dir: Path | None = None,
+    gemini_api_key: str = "",
+) -> None:
+    """Inject shared dependencies from the application lifespan."""
+    global _collection, _graph, _index_status, _profile, _project_dir, _gemini_api_key
+    _collection = collection
+    _graph = graph
+    _index_status = index_status
+    _profile = profile
+    _project_dir = project_dir
+    _gemini_api_key = gemini_api_key
+
+
+# ---------------------------------------------------------------------------
+# Language detection helper
+# ---------------------------------------------------------------------------
+
+_EXTENSION_TO_LANGUAGE: dict[str, str] = {
+    ".dart": "dart",
+    ".py": "python",
+    ".md": "markdown",
+    ".yaml": "yaml",
+    ".yml": "yaml",
+    ".json": "json",
+    ".ts": "typescript",
+    ".tsx": "typescript",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".html": "html",
+    ".css": "css",
+    ".sh": "bash",
+    ".go": "go",
+    ".rs": "rust",
+    ".java": "java",
+    ".kt": "kotlin",
+    ".swift": "swift",
+    ".rb": "ruby",
+    ".php": "php",
+    ".c": "c",
+    ".cpp": "cpp",
+    ".h": "c",
+    ".hpp": "cpp",
+    ".toml": "toml",
+    ".xml": "xml",
+    ".sql": "sql",
+    ".graphql": "graphql",
+    ".proto": "protobuf",
+    ".tf": "hcl",
+    ".dockerfile": "dockerfile",
+}
+
+
+def _detect_language(file_path: Path) -> str:
+    name = file_path.name.lower()
+    if name == "dockerfile":
+        return "dockerfile"
+    if name == "makefile":
+        return "makefile"
+    return _EXTENSION_TO_LANGUAGE.get(file_path.suffix.lower(), "plaintext")
+
+
+# ---------------------------------------------------------------------------
+# Routes
+# ---------------------------------------------------------------------------
+
+
+@router.post("/query")
+async def query(request: QueryRequest) -> EventSourceResponse:
+    """Stream a RAG response for the given query via Server-Sent Events."""
+    if not _gemini_api_key:
+        raise HTTPException(
+            status_code=503,
+            detail="GEMINI_API_KEY is not configured.",
+        )
+
+    from sourcefire.chain.rag_chain import stream_rag_response
+
+    async def _event_generator() -> AsyncGenerator[dict[str, str], None]:
+        async for chunk in stream_rag_response(
+            collection=_collection,
+            graph=_graph,
+            query=request.query,
+            mode=request.mode,
+            model=request.model,
+            history=request.history,
+            profile=_profile,
+            project_dir=_project_dir,
+            gemini_api_key=_gemini_api_key,
+        ):
+            yield {"data": json.dumps(chunk)}
+
+    return EventSourceResponse(_event_generator())
+
+
+@router.get("/sources", response_model=SourceResponse)
+async def sources(path: str = Query(..., description="Relative path within the codebase")) -> SourceResponse:
+    """Return the content and detected language of a source file."""
+    if _project_dir is None:
+        raise HTTPException(status_code=503, detail="Project directory not initialized.")
+
+    codebase_resolved = _project_dir.resolve()
+    full_path = (_project_dir / path).resolve()
+
+    if not str(full_path).startswith(str(codebase_resolved)):
+        raise HTTPException(status_code=400, detail="Path traversal detected.")
+
+    if not full_path.is_file():
+        raise HTTPException(status_code=404, detail=f"File not found: {path}")
+
+    try:
+        content = full_path.read_text(encoding="utf-8", errors="replace")
+    except OSError as exc:
+        raise HTTPException(status_code=500, detail=f"Could not read file: {exc}") from exc
+
+    return SourceResponse(
+        content=content,
+        language=_detect_language(full_path),
+    )
+
+
+@router.get("/status", response_model=StatusResponse)
+async def status() -> StatusResponse:
+    """Return current index status."""
+    return StatusResponse(
+        files_indexed=_index_status.get("files_indexed", 0),
+        last_indexed=str(_index_status.get("last_indexed", "never")),
+        index_status=str(_index_status.get("index_status", "not_ready")),
+        language=str(_index_status.get("language", "generic")),
+    )

sourcefire/chain/prompts.py

@@ -0,0 +1,195 @@
+"""Prompt assembly and token budget management for Sourcefire."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from sourcefire.config import MAX_HISTORY_PAIRS, MAX_TOKEN_BUDGET, RESPONSE_HEADROOM
+
+# ---------------------------------------------------------------------------
+# System template
+# ---------------------------------------------------------------------------
+
+from importlib.resources import files as _resource_files
+
+_SYSTEM_MD_PATH = Path(str(_resource_files("sourcefire") / "prompts" / "system.md"))
+_SYSTEM_TEMPLATE: str = _SYSTEM_MD_PATH.read_text(encoding="utf-8")
+
+# ---------------------------------------------------------------------------
+# Priority ordering (lower = higher priority)
+# ---------------------------------------------------------------------------
+
+_PRIORITY_ORDER: dict[str, int] = {"direct": 0, "semantic": 1, "graph": 2}
+
+# ---------------------------------------------------------------------------
+# Per-mode suffixes appended to the system prompt
+# ---------------------------------------------------------------------------
+
+_MODE_SUFFIXES: dict[str, str] = {
+    "debug": (
+        "\n\n## Mode: Debug\n"
+        "Diagnose the root cause from the retrieved context first. "
+        "If the stack trace references files not in the context, use tools to read them. "
+        "Trace the call chain. Show exact files and lines. Suggest a minimal fix."
+    ),
+    "feature": (
+        "\n\n## Mode: Feature\n"
+        "Use the retrieved context to identify the project structure and patterns. "
+        "Show where new code should live based on existing conventions you can see. "
+        "If you need to see similar features for reference, use semantic_code_search or find_files_by_name."
+    ),
+    "explain": (
+        "\n\n## Mode: Explain\n"
+        "Explain from the retrieved context first — it already contains the most relevant code. "
+        "Walk through files in dependency order. "
+        "Only use tools if you need to trace deeper connections not visible in the context."
+    ),
+}
+
+# ---------------------------------------------------------------------------
+# Token utilities
+# ---------------------------------------------------------------------------
+
+
+def estimate_tokens(text: str) -> int:
+    """Rough token estimate: 1 token ~ 4 characters."""
+    return len(text) // 4
+
+
+# ---------------------------------------------------------------------------
+# Chunk truncation with token budget enforcement
+# ---------------------------------------------------------------------------
+
+_MAX_CHUNK_CHARS = 6_000  # ~1500 tokens per chunk
+
+
+def truncate_chunks(
+    chunks: list[dict[str, Any]],
+    max_tokens: int,
+) -> list[dict[str, Any]]:
+    """Return a subset of *chunks* that fits within *max_tokens*."""
+    sorted_chunks = sorted(
+        chunks,
+        key=lambda c: (
+            _PRIORITY_ORDER.get(c.get("priority", "graph"), 2),
+            -float(c.get("relevance", 0.0)),
+        ),
+    )
+
+    capped: list[dict[str, Any]] = []
+    for chunk in sorted_chunks:
+        c = dict(chunk)
+        if len(c.get("code", "")) > _MAX_CHUNK_CHARS:
+            c["code"] = c["code"][:_MAX_CHUNK_CHARS]
+        capped.append(c)
+
+    result: list[dict[str, Any]] = []
+    used_tokens = 0
+    for chunk in capped:
+        chunk_tokens = estimate_tokens(chunk.get("code", ""))
+        if used_tokens + chunk_tokens <= max_tokens:
+            result.append(chunk)
+            used_tokens += chunk_tokens
+
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Prompt assembly
+# ---------------------------------------------------------------------------
+
+
+def assemble_prompt(
+    mode: str,
+    query: str,
+    chunks: list[dict[str, Any]],
+    claude_md: str,
+    memory_content: str,
+    history: list[dict[str, str]],
+    model: str,
+    highlight_language: str = "text",
+) -> dict[str, Any]:
+    """Assemble the full prompt dict for the LLM call.
+
+    Returns a dict with keys:
+    - ``system``   — full system prompt
+    - ``context``  — formatted retrieved code chunks
+    - ``query``    — the user's question (unchanged)
+    - ``history``  — trimmed conversation history
+    - ``stats``    — token usage summary dict
+    """
+    # 1. Compute available context budget
+    token_budget = MAX_TOKEN_BUDGET.get(model, 100_000)
+    system_tokens = estimate_tokens(_SYSTEM_TEMPLATE + claude_md + memory_content)
+    query_tokens = estimate_tokens(query)
+    history_tokens = sum(estimate_tokens(m.get("content", "")) for m in history)
+    overhead = system_tokens + query_tokens + history_tokens + RESPONSE_HEADROOM
+    context_budget = max(0, token_budget - overhead)
+
+    # 2. Truncate chunks to fit context budget
+    kept_chunks = truncate_chunks(chunks, max_tokens=context_budget)
+
+    # 3. Build system prompt
+    mode_suffix = _MODE_SUFFIXES.get(mode, "")
+    system_parts = [
+        _SYSTEM_TEMPLATE,
+        "\n\n---\n\n## Project Rules (CLAUDE.md)\n\n",
+        claude_md,
+        "\n\n---\n\n## Developer Memory\n\n",
+        memory_content,
+        mode_suffix,
+    ]
+    system_prompt = "".join(system_parts)
+
+    # 4. Build context block — use the detected language for syntax highlighting
+    context_parts: list[str] = []
+    for chunk in kept_chunks:
+        filename = chunk.get("filename", "unknown")
+        location = chunk.get("location", "")
+        relevance = chunk.get("relevance", 0.0)
+        code = chunk.get("code", "")
+
+        # Detect per-file highlight language from extension
+        ext = Path(filename).suffix.lower()
+        _ext_map = {
+            ".py": "python", ".js": "javascript", ".ts": "typescript",
+            ".tsx": "typescript", ".jsx": "javascript", ".dart": "dart",
+            ".go": "go", ".rs": "rust", ".java": "java", ".md": "markdown",
+            ".yaml": "yaml", ".yml": "yaml", ".json": "json", ".html": "html",
+            ".css": "css", ".sh": "bash", ".toml": "toml",
+        }
+        lang = _ext_map.get(ext, highlight_language)
+
+        header = f"### {filename}"
+        if location:
+            header += f" ({location})"
+        header += f"  [relevance: {relevance:.2f}]"
+        context_parts.append(f"{header}\n```{lang}\n{code}\n```")
+
+    context_block = "\n\n".join(context_parts)
+
+    # 5. Trim history to MAX_HISTORY_PAIRS
+    trimmed_history = history[-(MAX_HISTORY_PAIRS * 2):]
+
+    # 6. Stats
+    context_tokens = estimate_tokens(context_block)
+    stats = {
+        "model": model,
+        "token_budget": token_budget,
+        "system_tokens": system_tokens,
+        "context_tokens": context_tokens,
+        "query_tokens": query_tokens,
+        "history_tokens": history_tokens,
+        "total_estimated": system_tokens + context_tokens + query_tokens + history_tokens,
+        "chunks_used": len(kept_chunks),
+        "chunks_dropped": len(chunks) - len(kept_chunks),
+    }
+
+    return {
+        "system": system_prompt,
+        "context": context_block,
+        "query": query,
+        "history": trimmed_history,
+        "stats": stats,
+    }