contextual-engine 0.1.0__py3-none-any.whl

contextual/mcp/tools.py

@@ -0,0 +1,443 @@
+"""MCP tool implementations for search, indexing, and stats.
+
+This module provides the server-side implementations for the MCP tools
+that expose Contextual's core functionality to clients.
+"""
+from __future__ import annotations
+
+import json
+import os
+import re
+import threading
+import time
+import warnings
+from pathlib import Path
+from typing import Any, cast
+
+from contextual.core.errors import ErrorCode, MCPError
+from contextual.docs.retrieval import docs_search
+from contextual.embedding import embed_query_simple
+from contextual.embedding.embedder import CodeEmbedder, get_docs_embedder, get_embedder
+from contextual.indexing.processor import index_project
+from contextual.mcp.server import mcp
+from contextual.observability.logging import get_logger
+from contextual.retrieval.context_assembler import assemble_context
+from contextual.retrieval.search import hybrid_search
+from contextual.storage.sqlite_pool import SQLitePool
+
+
+# === CRITICAL: Suppress ALL stdout pollution BEFORE any imports ===
+os.environ["TQDM_DISABLE"] = "1"  # Disable tqdm progress bars
+os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"  # Disable HuggingFace progress
+os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"  # Disable telemetry
+os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Silence tokenizers warning
+
+warnings.filterwarnings("ignore")
+# ===================================================================
+
+logger = get_logger(__name__)
+
+# Module-level connection pool singleton
+_pool: SQLitePool | None = None
+_pool_lock = threading.Lock()
+
+
+def _get_pool() -> SQLitePool:
+    """Get or create global SQLite pool with thread-safe double-checked locking.
+
+    Looks for database in this order:
+    1. .contextual/contextual.db in current working directory (project-local)
+    2. ~/.contextual/contextual.db (user-global fallback)
+    """
+    global _pool
+    if _pool is None:
+        with _pool_lock:
+            if _pool is None:
+                # Try project-local database first
+                project_db = Path.cwd() / ".contextual" / "contextual.db"
+                if project_db.exists():
+                    db_path = project_db
+                    logger.info("Using project-local database", db_path=str(db_path))
+                else:
+                    # Fallback to home directory
+                    db_path = Path.home() / ".contextual" / "contextual.db"
+                    db_path.parent.mkdir(parents=True, exist_ok=True)
+                    logger.info("Using user-global database", db_path=str(db_path))
+
+                _pool = SQLitePool(db_path)
+    return _pool
+
+
+@mcp.tool()
+def search_codebase(
+    query: str,
+    top_k: int = 10,
+    language: str | None = None,
+    token_budget: int = 100_000,
+) -> str:
+    """Semantically search the indexed codebase.
+
+    Performs hybrid BM25 + vector search and returns formatted context
+    ready for LLM consumption.
+
+    Args:
+        query: Natural-language question or concept to search for.
+        top_k: Number of results to return (1-50).
+        language: Optional language filter (python, typescript, javascript, rust, go).
+        token_budget: Maximum tokens for assembled context (default 100k).
+
+    Returns:
+        Formatted context string with search results and metadata.
+
+    Raises:
+        MCPError: If search fails or database unavailable.
+    """
+    try:
+        # Embed query (no caching for queries)
+        query_embedding = embed_query_simple(query)
+
+        # Execute hybrid search with code-specific RRF weights
+        # Code benefits more from exact symbol matching (BM25) than semantic similarity
+        with _get_pool().reader() as conn:
+            results = hybrid_search(
+                conn,
+                query,
+                query_embedding,
+                top_k,
+                language,
+                bm25_weight=0.55,  # Exact term matching for code symbols
+                dense_weight=0.45,  # Less weight on semantic similarity for code
+            )
+
+        # Assemble context with token budget
+        context = assemble_context(results, token_budget)
+
+        logger.info(
+            "search_codebase_success",
+            query=query,
+            top_k=top_k,
+            language=language,
+            results_count=len(results),
+            total_tokens=context.total_tokens,
+        )
+
+        return context.content
+
+    except Exception as e:
+        logger.error("Tool failed", tool="search_codebase", error=str(e))
+        raise MCPError(
+            code=ErrorCode.MCP_TOOL_CALL_FAILED,
+            message=f"Tool 'search_codebase' failed: {e}",
+            context={"tool": "search_codebase"},
+        ) from e
+
+
+@mcp.tool()
+def index_codebase(
+    path: str,
+) -> dict[str, Any]:
+    """Index or re-index a code repository.
+
+    Discovers source files, parses with tree-sitter, chunks code,
+    generates embeddings, and stores in SQLite + vec0.
+
+    Args:
+        path: Absolute path to repository root directory.
+
+    Returns:
+        Dictionary with indexing statistics:
+        - files_indexed: Number of files processed
+        - chunks_created: Number of code chunks generated
+        - duration_ms: Indexing duration in milliseconds
+
+    Raises:
+        MCPError: If path invalid or indexing fails.
+    """
+    start = time.perf_counter()
+
+    try:
+        # Validate path
+        p = Path(path)
+        if not p.exists() or not p.is_dir():
+            raise ValueError(f"Invalid path: {path} is not an existing directory.")
+
+        # Resolve and security check
+        resolved_path = p.resolve()
+        if not resolved_path.is_absolute():
+            raise ValueError("Path must be absolute for security reasons.")
+
+        # Run indexing
+        chunks = index_project(resolved_path)
+        duration_ms = int((time.perf_counter() - start) * 1000)
+
+        result = {
+            "files_indexed": len({chunk.path for chunk in chunks}),
+            "chunks_created": len(chunks),
+            "duration_ms": duration_ms,
+        }
+
+        logger.info("index_codebase_success", path=path, stats=result)
+
+        return result
+
+    except Exception as e:
+        logger.error("Tool failed", tool="index_codebase", error=str(e))
+        raise MCPError(
+            code=ErrorCode.MCP_TOOL_CALL_FAILED,
+            message=f"Tool 'index_codebase' failed: {e}",
+            context={"tool": "index_codebase"},
+        ) from e
+
+
+@mcp.tool()
+def get_stats() -> dict[str, Any]:
+    """Get current index statistics.
+
+    Returns summary of indexed repository state including file counts,
+    chunk counts, languages, and last indexing timestamp.
+
+    Returns:
+        Dictionary with statistics:
+        - file_count: Total indexed files
+        - chunk_count: Total code chunks
+        - vector_count: Total vectors in vec_chunks
+        - last_indexed_at: Unix timestamp (ms) of last index, or None
+        - languages: List of detected languages
+    """
+    try:
+        with _get_pool().reader() as conn:
+            # Get counts
+            file_count = conn.execute("SELECT COUNT(*) FROM files").fetchone()[0]
+            chunk_count = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
+
+            # Vec0 vector count
+            vector_count_row = conn.execute("SELECT COUNT(*) FROM vec_chunks").fetchone()
+            vector_count = vector_count_row[0] if vector_count_row else 0
+
+            # Last indexed timestamp
+            last_indexed_row = conn.execute("SELECT MAX(indexed_at) FROM files").fetchone()
+            last_indexed_at = last_indexed_row[0] if last_indexed_row else None
+
+            # Languages
+            lang_rows = conn.execute(
+                "SELECT DISTINCT language FROM files ORDER BY language"
+            ).fetchall()
+            languages = [row[0] for row in lang_rows if row[0]]
+
+        stats = {
+            "file_count": file_count or 0,
+            "chunk_count": chunk_count or 0,
+            "vector_count": vector_count or 0,
+            "last_indexed_at": last_indexed_at,
+            "languages": languages,
+        }
+
+        logger.info("get_stats_success", stats=stats)
+
+        return stats
+
+    except Exception as e:
+        logger.error("Tool failed", tool="get_stats", error=str(e))
+
+        # Handle fresh database (no tables yet)
+        if "no such table" in str(e):
+            empty_stats: dict[str, Any] = {
+                "file_count": 0,
+                "chunk_count": 0,
+                "vector_count": 0,
+                "last_indexed_at": None,
+                "languages": [],
+            }
+            logger.warning("get_stats_empty_db", reason=str(e))
+            return empty_stats
+
+        raise MCPError(
+            code=ErrorCode.MCP_TOOL_CALL_FAILED,
+            message=f"Tool 'get_stats' failed: {e}",
+            context={"tool": "get_stats"},
+        ) from e
+
+
+# ---------------------------------------------------------------------------
+# Query heuristics — detect code vs prose intent
+# ---------------------------------------------------------------------------
+
+_CODE_SIGNAL_RE = re.compile(
+    r"[a-z_]{2,}\.[a-z_]{2,}\("      # method call: pool.reader(
+    r"|[A-Z][a-zA-Z0-9]{2,}\("        # class call: SQLitePool(
+    r"|`[^`]+`"                        # backtick symbol: `hybrid_search`
+    r"|\b[a-z_]{2,}_[a-z_]{2,}\b"     # snake_case: get_pool, chunk_id
+    r"|\bdef\b|\bclass\b|\bimport\b"   # Python keywords
+    r"|\b[A-Z][a-zA-Z0-9]{3,}(?:Pool|Manager|Writer|Handler|Client|Server|Cache|Store|Engine|Pipeline|Chunker|Embedder|Watcher|Builder|Parser|Runner|Iterator|Provider|Retrieval|Indexer)\b"
+    r"|\bthreading\b|\bsqlite3?\b|\bfastembed\b|\bwatchdog\b"
+    r"|\b(sqlite|vec0|fts5|onnx|pytest|async|await|bool|None|str|int)\b"
+    r"|\.(py|ts|js|rs|go|toml|yaml|json)\b"
+    r"|\b[A-Z][a-zA-Z0-9]{3,}(?:"
+    r"Pool|Manager|Writer|Handler|Client|Server|Cache|Store|Engine|Pipeline|"
+    r"Chunker|Embedder|Watcher|Builder|Parser|Runner|Iterator|Provider|Retrieval|"
+    r"Indexer"
+    r")\b"  # PascalCase class names with technical suffixes
+    r"|\bthreading\b|\bsqlite3?\b|\bfastembed\b|\bwatchdog\b",  # code-domain signals
+    re.IGNORECASE,
+)
+
+
+def _query_source_weights(query: str) -> tuple[float, float]:
+    """Return (code_weight, docs_weight) from heuristic signal detection.
+
+    Strong code signals → (0.70, 0.30)
+    Strong prose signals → (0.30, 0.70)
+    Ambiguous           → (0.50, 0.50)
+    """
+    words = query.split()
+    hits = len(_CODE_SIGNAL_RE.findall(query))
+    ratio = hits / max(len(words), 1)
+
+    if ratio >= 0.25:
+        return 0.70, 0.30
+    if ratio <= 0.05:
+        return 0.30, 0.70
+    return 0.50, 0.50
+
+
+@mcp.tool()
+def contextual_search(
+    query: str,
+    k: int = 10,
+    source: str = "auto",
+) -> str:
+    """Unified semantic search across codebase AND documentation.
+
+    This is the recommended default search tool. It automatically decides
+    whether your query needs code, docs, or both — and merges the results
+    into a single ranked list.
+
+    When to use this vs the specific tools:
+      contextual_search   — use by default; auto-routes intelligently
+      search_codebase     — use when you know you want code only
+      docs_search_tool    — use when you know you want docs only
+
+    Args:
+        query:  Natural language or keyword query.
+                Examples:
+                  "how is the db connection pool implemented"  → auto: both
+                  "SQLitePool.writer signature"                → auto: code
+                  "installation prerequisites macOS"           → auto: docs
+
+        k:      Total number of results to return (1-50, default 10).
+
+        source: Routing override.
+                "auto"  — detect from query content (default)
+                "code"  — search codebase index only
+                "docs"  — search docs index only (requires CONTEXTUAL_DOCS_ENABLED=1)
+                "both"  — always search both, merge results
+
+    Returns:
+        JSON array sorted by relevance score (highest first). Each result:
+          {
+            "source":       "code" | "docs",
+            "file_path":    "contextual/storage/sqlite_pool.py",
+            "content":      "..chunk text..",
+            "score":        0.0412,
+            "symbol_name":  "writer"          (code only, empty string for docs),
+            "language":     "python"          (code only, empty string for docs),
+            "heading_path": "Installation > macOS"  (docs only, empty for code),
+            "start_line":   172,
+            "end_line":     213
+          }
+    """
+    k = max(1, min(k, 50))
+    half_k = max(k, 10)   # over-fetch per corpus so merge has enough candidates
+
+    docs_enabled = os.getenv("CONTEXTUAL_DOCS_ENABLED", "0") == "1"
+
+    # ── Routing decision ────────────────────────────────────────────────────
+    if source == "auto":
+        run_code = True
+        run_docs = docs_enabled
+        code_w, docs_w = _query_source_weights(query)
+    elif source == "code":
+        run_code, run_docs = True, False
+        code_w, docs_w = 1.0, 0.0
+    elif source == "docs":
+        run_code, run_docs = False, docs_enabled
+        code_w, docs_w = 0.0, 1.0
+    else:  # "both"
+        run_code = True
+        run_docs = docs_enabled
+        code_w, docs_w = 0.50, 0.50
+
+    candidates: list[dict[str, Any]] = []
+
+    # ── Code path ───────────────────────────────────────────────────────────
+    if run_code:
+        try:
+            embedder = get_embedder()
+            query_embedding = embedder.embed([query])[0]
+
+            with _get_pool().reader() as conn:
+                hits = hybrid_search(
+                    conn=conn,
+                    query=query,
+                    query_embedding=query_embedding,
+                    top_k=half_k,
+                )
+
+            for hit in hits:
+                candidates.append({
+                    "source":       "code",
+                    "file_path":    hit.path,
+                    "content":      hit.snippet,
+                    "score":        round(hit.score * code_w, 6),
+                    "symbol_name":  hit.symbol_name or "",
+                    "language":     "",
+                    "heading_path": "",
+                    "start_line":   hit.start_line,
+                    "end_line":     hit.end_line,
+                })
+        except Exception as _exc:
+            logger.warning("contextual_search: code path failed", error=str(_exc))
+
+    # ── Docs path ───────────────────────────────────────────────────────────
+    if run_docs:
+        try:
+            with _get_pool().reader() as conn:
+                doc_hits = docs_search(
+                    query=query,
+                    conn=conn,
+                    embedder=cast(CodeEmbedder, get_docs_embedder()),
+                    k=half_k,
+                )
+
+            for doc_hit in doc_hits:
+                candidates.append({
+                    "source":       "docs",
+                    "file_path":    doc_hit.file_path,
+                    "content":      doc_hit.content,
+                    "score":        round(doc_hit.score * docs_w, 6),
+                    "symbol_name":  "",
+                    "language":     "",
+                    "heading_path": doc_hit.heading_path,
+                    "start_line":   doc_hit.start_line,
+                    "end_line":     doc_hit.end_line,
+                })
+        except Exception as _exc:
+            logger.warning("contextual_search: docs path failed", error=str(_exc))
+
+    # ── Merge + rank ────────────────────────────────────────────────────────
+    candidates.sort(key=lambda x: x["score"], reverse=True)
+    results = candidates[:k]
+
+    logger.info(
+        "contextual_search",
+        query=query[:80],
+        source=source,
+        run_code=run_code,
+        run_docs=run_docs,
+        code_w=code_w,
+        docs_w=docs_w,
+        candidates=len(candidates),
+        returned=len(results),
+    )
+
+    return json.dumps(results, ensure_ascii=False, indent=2)

contextual/observability/__init__.py

@@ -0,0 +1,21 @@
+"""Observability: structured logging, metrics collection, and graceful shutdown.
+
+Provides structlog-based logging pipeline with JSON output, local metrics collection,
+and graceful shutdown handlers for WAL checkpointing and async task draining.
+"""
+from __future__ import annotations
+
+from contextual.observability.logging import (
+    add_context,
+    clear_context,
+    configure_logging,
+    get_logger,
+)
+
+
+__all__ = [
+    "add_context",
+    "clear_context",
+    "configure_logging",
+    "get_logger",
+]

contextual/observability/logging.py

@@ -0,0 +1,115 @@
+"""Structured logging pipeline using structlog."""
+
+from __future__ import annotations
+
+import logging
+import sys
+from logging.handlers import RotatingFileHandler
+from typing import Any
+
+import structlog
+from structlog.types import Processor
+
+from contextual.config import get_log_dir
+
+
+def configure_logging(log_level: str = "INFO", enable_file_logging: bool = True) -> None:
+    """Configure the structlog pipeline.
+    
+    Sets up JSON file logging and colored console logging.
+    
+    Args:
+        log_level: The minimum log level to capture.
+        enable_file_logging: Whether to write logs to a file.
+    """
+    processors: list[Processor] = [
+        structlog.contextvars.merge_contextvars,
+        structlog.processors.add_log_level,
+        structlog.processors.TimeStamper(fmt="iso", utc=True),
+        structlog.processors.StackInfoRenderer(),
+        structlog.processors.format_exc_info,
+        structlog.processors.CallsiteParameterAdder(),
+    ]
+
+    # Console renderer (when sys.stderr.isatty())
+    if sys.stderr.isatty():
+        console_processor = structlog.dev.ConsoleRenderer()
+    else:
+        console_processor = structlog.processors.JSONRenderer()
+
+    # Configure stdlib logging handlers
+    handlers: list[logging.Handler] = []
+
+    # Console handler
+    console_handler = logging.StreamHandler(sys.stderr)
+    console_handler.setFormatter(structlog.stdlib.ProcessorFormatter(
+        processor=console_processor,
+    ))
+    handlers.append(console_handler)
+
+    if enable_file_logging:
+        try:
+            log_dir = get_log_dir()
+            log_dir.mkdir(parents=True, exist_ok=True)
+            log_file = log_dir / "contextual.log"
+
+            file_handler = RotatingFileHandler(
+                log_file,
+                maxBytes=10 * 1024 * 1024,  # 10 MB
+                backupCount=5,
+                encoding="utf-8"
+            )
+            # Ensure log file has proper permissions (0o644)
+            if log_file.exists():
+                log_file.chmod(0o644)
+
+            file_handler.setFormatter(structlog.stdlib.ProcessorFormatter(
+                processor=structlog.processors.JSONRenderer(),
+            ))
+            handlers.append(file_handler)
+        except Exception as e:
+            # Fallback to console-only if file logging setup fails
+            print(f"Warning: Failed to setup file logging: {e}", file=sys.stderr)
+
+    # Route stdlib logs through structlog
+    logging.basicConfig(
+        format="%(message)s",
+        level=log_level,
+        handlers=handlers,
+    )
+
+    structlog.configure(
+        processors=processors + [structlog.stdlib.ProcessorFormatter.wrap_for_formatter],
+        wrapper_class=structlog.make_filtering_bound_logger(getattr(logging, log_level.upper())),
+        context_class=dict,
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        cache_logger_on_first_use=True,
+    )
+
+
+def get_logger(name: str) -> structlog.BoundLogger:
+    """Get a structured logger.
+    
+    Args:
+        name: Logger name (usually __name__).
+        
+    Returns:
+        Configured structlog logger.
+    """
+    return structlog.get_logger(name)
+
+
+def add_context(**kwargs: Any) -> None:
+    """Add context variables to all logs in current context.
+    
+    Uses structlog.contextvars to bind variables.
+    
+    Args:
+        **kwargs: Key-value pairs to add to log context.
+    """
+    structlog.contextvars.bind_contextvars(**kwargs)
+
+
+def clear_context() -> None:
+    """Clear all context variables."""
+    structlog.contextvars.clear_contextvars()

contextual/retrieval/__init__.py

@@ -0,0 +1,24 @@
+"""Retrieval module: hybrid search, ranking, and context assembly.
+
+Public API:
+- hybrid_search: BM25 + vector search with RRF fusion
+- assemble_context: Token-budgeted context formatting
+- deduplicate_results: Content deduplication
+- apply_mmr: Diversity filtering
+"""
+from __future__ import annotations
+
+from contextual.retrieval.context_assembler import assemble_context
+from contextual.retrieval.search import SearchResult, hybrid_search
+from contextual.retrieval.ranker import apply_mmr, deduplicate_results
+
+
+# ranker imports will be added after delegation
+
+__all__ = [
+    "SearchResult",
+    "assemble_context",
+    "hybrid_search",
+    "apply_mmr",
+    "deduplicate_results"
+]