candlekeep 0.1.0__py3-none-any.whl

candlekeep/__init__.py

@@ -0,0 +1,36 @@
+"""Candlekeep: RAG knowledge base server."""
+import sys
+
+# Check Python version
+if sys.version_info < (3, 10):
+    raise RuntimeError(
+        f"Candlekeep requires Python 3.10+, found {sys.version_info.major}.{sys.version_info.minor}"
+    )
+
+# Check SQLite support (ChromaDB dependency)
+try:
+    import sqlite3
+except ImportError:
+    # Try pysqlite3-binary as fallback
+    try:
+        import pysqlite3 as sqlite3
+        sys.modules['sqlite3'] = sqlite3
+    except ImportError:
+        raise RuntimeError(
+            "Candlekeep requires Python with SQLite support (ChromaDB dependency)."
+        )
+
+# ChromaDB requires SQLite >= 3.35.0
+if sqlite3.sqlite_version_info < (3, 35, 0):
+    try:
+        __import__("pysqlite3")
+        sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
+    except ImportError:
+        raise RuntimeError(
+            f"ChromaDB requires SQLite >= 3.35.0, found {sqlite3.sqlite_version}\n"
+            "Fix: pip install pysqlite3-binary"
+        )
+
+from candlekeep.mcp.server import mcp, main
+
+__all__ = ["mcp", "main"]

candlekeep/__main__.py

@@ -0,0 +1,5 @@
+"""Entry point for candlekeep command."""
+from candlekeep.mcp.server import main
+
+if __name__ == "__main__":
+    main()

candlekeep/config.py

@@ -0,0 +1,150 @@
+"""Configuration for candlekeep with ChromaDB authentication."""
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+from urllib.parse import urlparse
+
+EmbeddingModel = Literal["minilm", "bge-small", "nomic"]
+DeviceType = Literal["auto", "cpu", "mps", "cuda"]
+
+EMBEDDING_MODELS = {
+    "minilm": "sentence-transformers/all-MiniLM-L6-v2",
+    "bge-small": "BAAI/bge-small-en-v1.5",
+    "nomic": "nomic-ai/nomic-embed-text-v1.5",
+}
+
+
+def detect_device() -> str:
+    """Detect best available compute device."""
+    import torch
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def load_dotenv(env_file: Path | None = None) -> None:
+    """Load .env file into environment variables."""
+    if env_file is None:
+        for candidate in [Path(".env"), Path(__file__).parent.parent.parent / ".env"]:
+            if candidate.exists():
+                env_file = candidate
+                break
+    if env_file is None or not env_file.exists():
+        return
+    
+    for line in env_file.read_text().splitlines():
+        line = line.strip()
+        if line and not line.startswith("#") and "=" in line:
+            if line.startswith("export "):
+                line = line[7:]
+            key, _, value = line.partition("=")
+            key, value = key.strip(), value.strip().strip('"').strip("'")
+            if key and key not in os.environ:
+                os.environ[key] = value
+
+
+load_dotenv()
+
+
+def get_data_dir() -> Path:
+    """Get cross-platform data directory."""
+    if os.name == "nt":
+        base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local"))
+    elif os.name == "posix" and "darwin" in os.uname().sysname.lower():
+        base = Path.home() / "Library" / "Application Support"
+    else:
+        base = Path(os.environ.get("XDG_DATA_HOME", Path.home() / ".local" / "share"))
+    return base / "candlekeep"
+
+
+@dataclass
+class Settings:
+    """Candlekeep configuration settings."""
+    
+    # ChromaDB connection
+    chroma_url: str = field(default_factory=lambda: os.getenv("CHROMA_URL", "http://localhost:8000"))
+    chroma_auth_token: str = field(default_factory=lambda: os.getenv("CHROMA_AUTH_TOKEN", ""))
+    
+    # Embedding settings
+    embedding_model: EmbeddingModel = field(default_factory=lambda: os.getenv("CANDLEKEEP_EMBEDDING", "bge-small"))
+    embedding_cache_size: int = field(default_factory=lambda: int(os.getenv("CANDLEKEEP_EMBEDDING_CACHE_SIZE", "500")))
+    
+    # Inference device
+    device: str = field(default_factory=lambda: os.getenv("CANDLEKEEP_DEVICE", "auto"))
+    
+    # Document processing
+    chunk_size: int = field(default_factory=lambda: int(os.getenv("CANDLEKEEP_CHUNK_SIZE", "512")))
+    chunk_overlap: int = field(default_factory=lambda: int(os.getenv("CANDLEKEEP_CHUNK_OVERLAP", "50")))
+    
+    # Personality
+    spice: bool = field(default_factory=lambda: os.getenv("CANDLEKEEP_SPICE", "false").lower() == "true")
+    
+    # Structural Integrity (Bardic Knowledge)
+    bardic_knowledge: bool = field(default_factory=lambda: os.getenv("CANDLEKEEP_BARDIC_KNOWLEDGE", "true").lower() == "true")
+
+    # Sparse backend for hybrid path (bm25 or colbert)
+    sparse_backend: str = field(default_factory=lambda: os.getenv("CANDLEKEEP_SPARSE_BACKEND", "bm25"))
+
+    # Transport settings (multi-agent HTTP mode)
+    transport: str = field(default_factory=lambda: os.getenv("CANDLEKEEP_TRANSPORT", "stdio"))
+    http_host: str = field(default_factory=lambda: os.getenv("CANDLEKEEP_HTTP_HOST", "127.0.0.1"))
+    http_port: int = field(default_factory=lambda: int(os.getenv("CANDLEKEEP_HTTP_PORT", "8111")))
+    mcp_token: str = field(default_factory=lambda: os.getenv("CANDLEKEEP_MCP_TOKEN", ""))
+
+    # Rate limiting (HTTP mode only, per MCP session)
+    rate_limit_search: int = field(default_factory=lambda: int(os.getenv("CANDLEKEEP_RATE_LIMIT_SEARCH", "30")))
+    rate_limit_write: int = field(default_factory=lambda: int(os.getenv("CANDLEKEEP_RATE_LIMIT_WRITE", "5")))
+    rate_limit_window: int = field(default_factory=lambda: int(os.getenv("CANDLEKEEP_RATE_LIMIT_WINDOW", "60")))
+
+    # LLM / Vision provider selection (anthropic, openai, bedrock, openai_compat)
+    llm_provider: str = field(default_factory=lambda: os.getenv("CANDLEKEEP_LLM_PROVIDER", ""))
+    vlm_provider: str = field(default_factory=lambda: os.getenv("CANDLEKEEP_VLM_PROVIDER", ""))
+
+    # Data directory
+    data_dir: Path = field(default_factory=get_data_dir)
+
+    def __post_init__(self):
+        if self.device == "auto":
+            self.device = detect_device()
+        self.data_dir = Path(self.data_dir)
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+        (self.data_dir / "models").mkdir(exist_ok=True)
+        (self.data_dir / "chroma").mkdir(exist_ok=True)
+    
+    @classmethod
+    def from_env(cls) -> "Settings":
+        """Create settings from environment variables."""
+        return cls()
+    
+    @property
+    def chroma_host(self) -> str:
+        """Extract host from CHROMA_URL."""
+        parsed = urlparse(self.chroma_url)
+        return parsed.hostname or "localhost"
+    
+    @property
+    def chroma_port(self) -> int:
+        """Extract port from CHROMA_URL."""
+        parsed = urlparse(self.chroma_url)
+        return parsed.port or 8000
+    
+    @property
+    def chroma_ssl(self) -> bool:
+        """Check if HTTPS is used."""
+        parsed = urlparse(self.chroma_url)
+        return parsed.scheme == "https"
+    
+    @property
+    def models_dir(self) -> Path:
+        return self.data_dir / "models"
+    
+    @property
+    def chroma_dir(self) -> Path:
+        return self.data_dir / "chroma"
+    
+    @property
+    def embedding_model_name(self) -> str:
+        return EMBEDDING_MODELS[self.embedding_model]

candlekeep/database/__init__.py

@@ -0,0 +1 @@
+"""Database layer for candlekeep."""

candlekeep/database/embeddings.py

@@ -0,0 +1,115 @@
+"""Embedding manager with model caching and selection."""
+import threading
+from collections import OrderedDict
+from sentence_transformers import SentenceTransformer
+from candlekeep.config import Settings, EMBEDDING_MODELS, EmbeddingModel
+
+
+class EmbeddingManager:
+    _instance: "EmbeddingManager | None" = None
+    _model: SentenceTransformer | None = None
+    _current_model: str | None = None
+
+    def __init__(self, settings: Settings | None = None):
+        self.settings = settings or Settings.from_env()
+        self._cache = OrderedDict()
+        self._cache_lock = threading.Lock()
+        self._hits = 0
+        self._misses = 0
+        self._max_cache_size = self.settings.embedding_cache_size
+
+
+    @classmethod
+    def get_instance(cls, settings: Settings | None = None) -> "EmbeddingManager":
+        if cls._instance is None:
+            cls._instance = cls(settings)
+        return cls._instance
+
+    def get_model(self, model_name: EmbeddingModel | None = None) -> SentenceTransformer:
+        """Load model from local cache. Exits if not found locally."""
+        model_name = model_name or self.settings.embedding_model
+        model_id = EMBEDDING_MODELS[model_name]
+
+        if self._model is not None and self._current_model == model_id:
+            return self._model
+
+        cache_dir = self.settings.models_dir
+        # Check if model exists locally before loading
+        if not any(cache_dir.glob(f"models--{model_id.replace('/', '--')}*")):
+            import sys
+            print(f"[candlekeep] ❌ Model '{model_id}' not found locally in {cache_dir}. "
+                  f"Run ./scripts/setup.sh to download models.", file=sys.stderr)
+            sys.exit(1)
+
+        trust_remote_code = model_name == "nomic"
+        self._model = SentenceTransformer(
+            model_id,
+            cache_folder=str(cache_dir),
+            trust_remote_code=trust_remote_code,
+            local_files_only=True,
+            device=self.settings.device,
+        )
+        
+        # Promote to float64 on CPU for better determinism and to match reranker precision
+        if self.settings.device == "cpu":
+            try:
+                import torch
+                self._model.to(torch.float64)
+            except Exception as e:
+                # Fallback to float32 if conversion fails (e.g. meta tensors in CI)
+                import sys
+                print(f"[candlekeep] ⚠ Could not promote embedding model to float64: {e}. "
+                      f"Continuing with float32.", file=sys.stderr)
+            
+        self._current_model = model_id
+        return self._model
+
+    def embed(self, texts: list[str] | str, model_name: EmbeddingModel | None = None, is_query: bool = False) -> list[list[float]]:
+        """Generate embeddings for texts."""
+        model_name = model_name or self.settings.embedding_model
+        model = self.get_model(model_name)
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        # Nomic requires prefixes
+        if model_name == "nomic":
+            prefix = "search_query: " if is_query else "search_document: "
+            texts = [prefix + t for t in texts]
+        
+        embeddings = model.encode(texts, convert_to_numpy=True)
+        return embeddings.tolist()
+
+    def embed_query(self, query: str, model_name: EmbeddingModel | None = None) -> list[float]:
+        """Generate embedding for a single query with LRU cache."""
+        model_name = model_name or self.settings.embedding_model
+        
+        if self._max_cache_size == 0:
+            return self.embed([query], model_name, is_query=True)[0]
+
+        cache_key = (query, model_name)
+        with self._cache_lock:
+            if cache_key in self._cache:
+                self._hits += 1
+                self._cache.move_to_end(cache_key)
+                return self._cache[cache_key]
+        
+        # Cache miss - compute embedding
+        embedding = self.embed([query], model_name, is_query=True)[0]
+        
+        with self._cache_lock:
+            self._misses += 1
+            self._cache[cache_key] = embedding
+            self._cache.move_to_end(cache_key)
+            if len(self._cache) > self._max_cache_size:
+                self._cache.popitem(last=False)
+        
+        return embedding
+
+    def get_cache_stats(self) -> dict:
+        """Get embedding cache hit/miss statistics."""
+        with self._cache_lock:
+            return {
+                "embedding_cache_hits": self._hits,
+                "embedding_cache_misses": self._misses,
+                "embedding_cache_size": len(self._cache)
+            }

candlekeep/database/interface.py

@@ -0,0 +1,86 @@
+"""Abstract database interface for candlekeep."""
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Dict, Any
+
+
+@dataclass
+class SearchResult:
+    """Search result from vector database."""
+    text: str
+    metadata: Dict[str, Any]
+    score: float
+    doc_id: str
+
+
+@dataclass
+class Chunk:
+    """Document chunk for indexing."""
+    text: str
+    metadata: Dict[str, Any]
+    chunk_index: int
+
+
+class VectorDatabase(ABC):
+    """Abstract interface for vector database operations."""
+    
+    @abstractmethod
+    def search(self, query: str, n_results: int = 5, category: str | None = None) -> List[SearchResult]:
+        """Search for similar documents."""
+        pass
+    
+    @abstractmethod
+    def add_documents(self, chunks: List[Chunk], collection: str = "default") -> int:
+        """Add document chunks to database."""
+        pass
+    
+    @abstractmethod
+    def delete_by_source(self, source: str) -> int:
+        """Delete all chunks from a source."""
+        pass
+    
+    @abstractmethod
+    def clear(self) -> None:
+        """Clear all documents from database."""
+        pass
+    
+    @abstractmethod
+    def get_stats(self) -> Dict[str, Any]:
+        """Get database statistics."""
+        pass
+    
+    @abstractmethod
+    def list_documents(self) -> List[Dict[str, Any]]:
+        """List all indexed documents."""
+        pass
+
+    @abstractmethod
+    def get_chunks_by_source(self, source: str) -> List[SearchResult]:
+        """Get all chunks from a specific source document."""
+        pass
+
+    @abstractmethod
+    def get_all_chunks(self) -> List[SearchResult]:
+        """Get all chunks from the database."""
+        pass
+
+    @abstractmethod
+    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Get embeddings for a list of texts."""
+        pass
+
+    @abstractmethod
+    def embed_query(self, query: str) -> List[float]:
+        """Get cached embedding for a single query."""
+        pass
+
+    @abstractmethod
+    def get_stored_embeddings_by_source(self, source: str) -> Dict[int, List[float]]:
+
+        """Get stored embeddings for all chunks from a source document.
+
+        Returns a dict mapping chunk_index to the embedding vector that was
+        computed at ingestion time. This avoids re-computing embeddings at
+        query time for similarity-weighted expansion.
+        """
+        pass