okb 1.0.0__py3-none-any.whl

okb/migrate.py

@@ -0,0 +1,53 @@
+"""Migration runner for okb database schema."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from yoyo import get_backend, read_migrations
+
+
+def get_migrations_path() -> str:
+    """Get path to migrations directory."""
+    return str(Path(__file__).parent / "migrations")
+
+
+def _convert_db_url(db_url: str) -> str:
+    """Convert psycopg3 URL to yoyo-compatible format.
+
+    yoyo uses psycopg2 by default. We convert:
+    postgresql://... -> postgresql+psycopg://...
+    to use psycopg v3.
+    """
+    if db_url.startswith("postgresql://") and "+psycopg" not in db_url:
+        return db_url.replace("postgresql://", "postgresql+psycopg://", 1)
+    return db_url
+
+
+def run_migrations(db_url: str) -> list[str]:
+    """Apply pending migrations, return list of applied migration IDs."""
+    backend = get_backend(_convert_db_url(db_url))
+    migrations = read_migrations(get_migrations_path())
+
+    with backend.lock():
+        to_apply = backend.to_apply(migrations)
+        if to_apply:
+            backend.apply_migrations(to_apply)
+
+    return [m.id for m in to_apply]
+
+
+def get_pending(db_url: str) -> list[str]:
+    """Get list of pending migration IDs."""
+    backend = get_backend(_convert_db_url(db_url))
+    migrations = read_migrations(get_migrations_path())
+    return [m.id for m in backend.to_apply(migrations)]
+
+
+def get_applied(db_url: str) -> list[str]:
+    """Get list of applied migration IDs."""
+    backend = get_backend(_convert_db_url(db_url))
+    migrations = read_migrations(get_migrations_path())
+    to_apply = backend.to_apply(migrations)
+    to_apply_ids = {m.id for m in to_apply}
+    return [m.id for m in migrations if m.id not in to_apply_ids]

okb/migrations/0001.initial-schema.sql

@@ -0,0 +1,91 @@
+-- Initial schema - documents, chunks, indexes
+-- depends:
+
+CREATE EXTENSION IF NOT EXISTS vector;
+
+-- Main documents table
+CREATE TABLE IF NOT EXISTS documents (
+    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    source_path TEXT NOT NULL,
+    source_type TEXT NOT NULL,
+    title TEXT,
+    content TEXT NOT NULL,
+    metadata JSONB DEFAULT '{}',
+    created_at TIMESTAMPTZ DEFAULT NOW(),
+    updated_at TIMESTAMPTZ DEFAULT NOW(),
+    content_hash TEXT NOT NULL,
+    UNIQUE(content_hash)
+);
+
+-- Chunks for semantic search
+CREATE TABLE IF NOT EXISTS chunks (
+    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    document_id UUID REFERENCES documents(id) ON DELETE CASCADE,
+    chunk_index INTEGER NOT NULL,
+    content TEXT NOT NULL,
+    embedding_text TEXT NOT NULL,
+    embedding vector(768),
+    token_count INTEGER,
+    metadata JSONB DEFAULT '{}',
+    created_at TIMESTAMPTZ DEFAULT NOW()
+);
+
+-- HNSW index for similarity search
+CREATE INDEX IF NOT EXISTS chunks_embedding_idx ON chunks
+USING hnsw (embedding vector_cosine_ops)
+WITH (m = 16, ef_construction = 64);
+
+-- Full-text search indexes
+CREATE INDEX IF NOT EXISTS documents_content_fts ON documents
+USING gin(to_tsvector('english', content));
+
+CREATE INDEX IF NOT EXISTS chunks_content_fts ON chunks
+USING gin(to_tsvector('english', content));
+
+-- Lookup indexes
+CREATE INDEX IF NOT EXISTS documents_source_path_idx ON documents(source_path);
+CREATE INDEX IF NOT EXISTS documents_source_type_idx ON documents(source_type);
+CREATE INDEX IF NOT EXISTS documents_metadata_idx ON documents USING gin(metadata);
+
+-- Function to update timestamp
+CREATE OR REPLACE FUNCTION update_updated_at()
+RETURNS TRIGGER AS $$
+BEGIN
+    NEW.updated_at = NOW();
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Trigger (drop first to make idempotent)
+DROP TRIGGER IF EXISTS documents_updated_at ON documents;
+CREATE TRIGGER documents_updated_at
+    BEFORE UPDATE ON documents
+    FOR EACH ROW
+    EXECUTE FUNCTION update_updated_at();
+
+-- Helper view for search results
+CREATE OR REPLACE VIEW search_results AS
+SELECT
+    c.id as chunk_id,
+    c.content,
+    c.chunk_index,
+    c.token_count,
+    c.embedding,
+    d.id as document_id,
+    d.source_path,
+    d.source_type,
+    d.title,
+    d.metadata
+FROM chunks c
+JOIN documents d ON c.document_id = d.id;
+
+-- Stats view
+CREATE OR REPLACE VIEW index_stats AS
+SELECT
+    source_type,
+    COUNT(DISTINCT d.id) as document_count,
+    COUNT(c.id) as chunk_count,
+    SUM(c.token_count) as total_tokens
+FROM documents d
+LEFT JOIN chunks c ON c.document_id = d.id
+GROUP BY source_type;

okb/migrations/0002.sync-state.sql

@@ -0,0 +1,22 @@
+-- Sync state for API sources (plugin system)
+-- depends: 0001.initial-schema
+
+CREATE TABLE IF NOT EXISTS sync_state (
+    source_name TEXT NOT NULL,
+    database_name TEXT NOT NULL DEFAULT 'default',
+    last_sync TIMESTAMPTZ,
+    cursor TEXT,
+    extra JSONB DEFAULT '{}',
+    updated_at TIMESTAMPTZ DEFAULT NOW(),
+    PRIMARY KEY (source_name, database_name)
+);
+
+-- Index for listing sync states
+CREATE INDEX IF NOT EXISTS sync_state_updated_idx ON sync_state(updated_at DESC);
+
+-- Trigger for updated_at
+DROP TRIGGER IF EXISTS sync_state_updated_at ON sync_state;
+CREATE TRIGGER sync_state_updated_at
+    BEFORE UPDATE ON sync_state
+    FOR EACH ROW
+    EXECUTE FUNCTION update_updated_at();

okb/migrations/0003.structured-fields.sql

@@ -0,0 +1,22 @@
+-- Structured fields for actionable items (tasks, events, emails)
+-- depends: 0002.sync-state
+
+-- Add structured fields to documents table for temporal/status queries
+ALTER TABLE documents ADD COLUMN IF NOT EXISTS due_date TIMESTAMPTZ;
+ALTER TABLE documents ADD COLUMN IF NOT EXISTS event_start TIMESTAMPTZ;
+ALTER TABLE documents ADD COLUMN IF NOT EXISTS event_end TIMESTAMPTZ;
+ALTER TABLE documents ADD COLUMN IF NOT EXISTS status TEXT;
+ALTER TABLE documents ADD COLUMN IF NOT EXISTS priority INTEGER;
+
+-- Indexes for common query patterns
+CREATE INDEX IF NOT EXISTS documents_due_date_idx ON documents(due_date) WHERE due_date IS NOT NULL;
+CREATE INDEX IF NOT EXISTS documents_event_start_idx ON documents(event_start) WHERE event_start IS NOT NULL;
+CREATE INDEX IF NOT EXISTS documents_status_idx ON documents(status) WHERE status IS NOT NULL;
+
+-- Composite index for "incomplete tasks due soon" queries
+CREATE INDEX IF NOT EXISTS documents_actionable_idx ON documents(due_date, status)
+    WHERE due_date IS NOT NULL AND status IS NOT NULL;
+
+-- Composite index for "events in date range" queries
+CREATE INDEX IF NOT EXISTS documents_event_range_idx ON documents(event_start, event_end)
+    WHERE event_start IS NOT NULL;

okb/migrations/0004.tokens.sql

@@ -0,0 +1,13 @@
+-- API tokens for HTTP authentication
+-- depends: 0003.structured-fields
+
+CREATE TABLE IF NOT EXISTS tokens (
+    token_hash TEXT PRIMARY KEY,  -- SHA256(full_token), hex encoded
+    permissions TEXT NOT NULL CHECK (permissions IN ('ro', 'rw')),
+    description TEXT,
+    created_at TIMESTAMPTZ DEFAULT NOW(),
+    last_used_at TIMESTAMPTZ
+);
+
+-- Index for listing tokens
+CREATE INDEX IF NOT EXISTS tokens_created_at_idx ON tokens(created_at DESC);

okb/migrations/0005.database-metadata.sql

@@ -0,0 +1,19 @@
+-- Database metadata table for storing LLM-enhanced descriptions
+-- depends: 0004.tokens
+
+CREATE TABLE IF NOT EXISTS database_metadata (
+    id SERIAL PRIMARY KEY,
+    key TEXT NOT NULL UNIQUE,  -- 'description', 'topics', etc.
+    value JSONB NOT NULL,
+    source TEXT NOT NULL DEFAULT 'llm',  -- 'config' or 'llm'
+    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+);
+
+-- Index for quick key lookups
+CREATE INDEX IF NOT EXISTS idx_database_metadata_key ON database_metadata(key);
+
+-- Insert default entries that can be updated by LLM
+INSERT INTO database_metadata (key, value, source) VALUES
+    ('llm_description', 'null'::jsonb, 'llm'),
+    ('llm_topics', '[]'::jsonb, 'llm')
+ON CONFLICT (key) DO NOTHING;

okb/migrations/0006.llm-cache.sql

@@ -0,0 +1,13 @@
+-- LLM response cache for avoiding redundant API calls
+-- depends: 0005.database-metadata
+
+CREATE TABLE IF NOT EXISTS llm_cache (
+    content_hash TEXT PRIMARY KEY,  -- SHA256 of (prompt + system + model)
+    provider TEXT NOT NULL,
+    model TEXT NOT NULL,
+    response TEXT NOT NULL,  -- JSON: {content, input_tokens, output_tokens}
+    created_at TIMESTAMPTZ DEFAULT NOW()
+);
+
+-- Index for cache cleanup queries
+CREATE INDEX IF NOT EXISTS llm_cache_created_at_idx ON llm_cache(created_at);

okb/modal_embedder.py

@@ -0,0 +1,120 @@
+"""
+Modal-based GPU embedding service.
+
+Provides on-demand GPU access for batch embedding generation.
+Costs approximately $0.02 per 1000 chunks on T4 GPU.
+
+Usage:
+    modal deploy modal_embedder.py
+
+Then call from Python:
+    embedder = modal.Cls.from_name("knowledge-embedder", "Embedder")()
+    embeddings = embedder.embed_batch.remote(texts)
+"""
+
+import modal
+
+app = modal.App("knowledge-embedder")
+
+# Container image with all dependencies
+embedder_image = modal.Image.debian_slim(python_version="3.11").pip_install(
+    "sentence-transformers>=2.2.0",
+    "torch>=2.0.0",
+    "numpy>=1.24.0",
+    "einops>=0.7.0",  # Required by nomic
+)
+
+
+@app.cls(
+    image=embedder_image,
+    gpu="T4",  # Cheapest option, sufficient for embedding
+    timeout=600,
+    scaledown_window=300,  # Keep warm for 5 min
+    retries=2,
+)
+class Embedder:
+    """GPU-accelerated embedding generator using nomic-embed-text."""
+
+    @modal.enter()
+    def load_model(self):
+        """Load model once when container starts."""
+        from sentence_transformers import SentenceTransformer
+        import torch
+
+        self.model = SentenceTransformer(
+            "nomic-ai/nomic-embed-text-v1.5",
+            trust_remote_code=True,
+        )
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(device)
+        print(f"Model loaded on {device}")
+
+    @modal.method()
+    def embed_batch(
+        self,
+        texts: list[str],
+        is_query: bool = False,
+        batch_size: int = 16,
+    ) -> list[list[float]]:
+        """
+        Generate embeddings for a batch of texts.
+
+        Args:
+            texts: List of text chunks to embed
+            is_query: If True, use query prefix; otherwise document prefix
+            batch_size: Processing batch size
+
+        Returns:
+            List of embedding vectors (768 dimensions)
+        """
+        # Nomic model requires task-specific prefixes
+        prefix = "search_query: " if is_query else "search_document: "
+        prefixed = [f"{prefix}{t}" for t in texts]
+
+        embeddings = self.model.encode(
+            prefixed,
+            batch_size=batch_size,
+            show_progress_bar=len(texts) > 100,
+            convert_to_numpy=True,
+            normalize_embeddings=True,  # For cosine similarity
+        )
+
+        return embeddings.tolist()
+
+    @modal.method()
+    def embed_single(self, text: str, is_query: bool = False) -> list[float]:
+        """Embed a single text (convenience method)."""
+        prefix = "search_query: " if is_query else "search_document: "
+        embedding = self.model.encode(
+            f"{prefix}{text}",
+            convert_to_numpy=True,
+            normalize_embeddings=True,
+        )
+        return embedding.tolist()
+
+
+# For testing
+@app.local_entrypoint()
+def test():
+    """Test the embedder."""
+    embedder = Embedder()
+
+    test_texts = [
+        "Django ORM query optimization using select_related and prefetch_related",
+        "PostgreSQL VACUUM and ANALYZE for table maintenance",
+        "Kubernetes pod scheduling with node affinity rules",
+    ]
+
+    print(f"Embedding {len(test_texts)} texts...")
+    embeddings = embedder.embed_batch.remote(test_texts)
+
+    print(f"Generated {len(embeddings)} embeddings")
+    print(f"Embedding dimension: {len(embeddings[0])}")
+
+    # Test similarity
+    import numpy as np
+
+    emb = np.array(embeddings)
+    similarity = np.dot(emb, emb.T)
+    print(f"\nSimilarity matrix:\n{similarity}")

okb/modal_llm.py

@@ -0,0 +1,178 @@
+"""
+Modal-based GPU LLM service for document classification.
+
+Provides on-demand GPU access for LLM inference using open models.
+Uses Llama 3.2 3B by default - fast and efficient for classification tasks.
+
+Usage:
+    modal deploy modal_llm.py
+
+Then call from Python:
+    llm = modal.Cls.from_name("knowledge-llm", "LLM")()
+    response = llm.complete.remote("Classify this document", system="You are a classifier")
+"""
+
+import modal
+
+app = modal.App("knowledge-llm")
+
+# Container image with transformers and torch
+llm_image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .pip_install(
+        "transformers>=4.40.0",
+        "torch>=2.0.0",
+        "accelerate>=0.27.0",
+        "bitsandbytes>=0.42.0",  # For quantization
+    )
+    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
+)
+
+# Default model - Llama 3.2 3B is fast and good for classification
+DEFAULT_MODEL = "meta-llama/Llama-3.2-3B-Instruct"
+
+
+@app.cls(
+    image=llm_image,
+    gpu="T4",  # T4 is sufficient for 3B model with quantization
+    timeout=300,
+    scaledown_window=300,  # Keep warm for 5 min
+    retries=1,
+)
+class LLM:
+    """GPU-accelerated LLM for document classification."""
+
+    model_id: str = DEFAULT_MODEL
+
+    @modal.enter()
+    def load_model(self):
+        """Load model once when container starts."""
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+        print(f"Loading model: {self.model_id}")
+
+        # Use 4-bit quantization for memory efficiency
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_id,
+            quantization_config=quantization_config,
+            device_map="auto",
+            torch_dtype=torch.float16,
+        )
+
+        # Set pad token if not set
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        print(f"Model loaded on {self.model.device}")
+
+    @modal.method()
+    def complete(
+        self,
+        prompt: str,
+        system: str | None = None,
+        max_tokens: int = 256,
+        temperature: float = 0.1,
+    ) -> dict:
+        """Generate a completion for the given prompt.
+
+        Args:
+            prompt: User prompt
+            system: Optional system prompt
+            max_tokens: Maximum tokens to generate
+            temperature: Sampling temperature (lower = more deterministic)
+
+        Returns:
+            Dict with 'content', 'model', 'input_tokens', 'output_tokens'
+        """
+        import torch
+
+        # Build messages in chat format
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+
+        # Apply chat template
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
+        input_tokens = inputs.input_ids.shape[1]
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=temperature > 0,
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+
+        # Decode only the new tokens
+        output_tokens = outputs.shape[1] - input_tokens
+        response_text = self.tokenizer.decode(
+            outputs[0][input_tokens:],
+            skip_special_tokens=True,
+        )
+
+        return {
+            "content": response_text.strip(),
+            "model": self.model_id,
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+        }
+
+    @modal.method()
+    def complete_batch(
+        self,
+        prompts: list[str],
+        system: str | None = None,
+        max_tokens: int = 256,
+    ) -> list[dict]:
+        """Generate completions for multiple prompts.
+
+        Args:
+            prompts: List of user prompts
+            system: Optional system prompt (same for all)
+            max_tokens: Maximum tokens per response
+
+        Returns:
+            List of response dicts
+        """
+        # For now, process sequentially (batched inference is more complex)
+        return [self.complete(prompt, system=system, max_tokens=max_tokens) for prompt in prompts]
+
+
+# For testing
+@app.local_entrypoint()
+def test():
+    """Test the LLM."""
+    llm = LLM()
+
+    system = "You are a document classifier. Respond with JSON: {action, reason}"
+    prompt = """Classify this email:
+
+Subject: 50% OFF - Limited Time Offer!
+From: deals@marketing-spam.com
+
+Don't miss out on our biggest sale of the year!
+Click here to claim your discount before it expires!"""
+
+    print("Testing LLM classification...")
+    response = llm.complete.remote(prompt, system=system, max_tokens=100)
+
+    print(f"Model: {response['model']}")
+    print(f"Tokens: {response['input_tokens']} in, {response['output_tokens']} out")
+    print(f"Response:\n{response['content']}")

okb/plugins/__init__.py

@@ -0,0 +1,8 @@
+"""Plugin system for LKB - extensible file parsers and API sources."""
+
+# Re-export Document from ingest for plugin authors
+from ..ingest import Document
+from .base import APISource, FileParser, SyncState
+from .registry import PluginRegistry
+
+__all__ = ["FileParser", "APISource", "SyncState", "Document", "PluginRegistry"]

okb/plugins/base.py

@@ -0,0 +1,110 @@
+"""Protocol definitions for LKB plugins."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from ..ingest import Document
+
+
+@dataclass
+class SyncState:
+    """Tracks sync progress for incremental updates."""
+
+    last_sync: datetime | None = None
+    cursor: str | None = None
+    extra: dict = field(default_factory=dict)
+
+
+@runtime_checkable
+class FileParser(Protocol):
+    """Protocol for file format parsers.
+
+    Plugins implement this to add support for new file types.
+
+    Example:
+        class MyParser:
+            extensions = ['.xyz']
+            source_type = 'xyz'
+
+            def can_parse(self, path: Path) -> bool:
+                return path.suffix.lower() == '.xyz'
+
+            def parse(self, path: Path, extra_metadata: dict | None = None) -> Document:
+                ...
+    """
+
+    extensions: list[str]  # e.g., ['.pdf', '.PDF'] - for fast pre-filtering
+    source_type: str  # e.g., 'pdf'
+
+    def can_parse(self, path: Path) -> bool:
+        """Check if this parser can handle the file (beyond just extension).
+
+        Called after extension match. Can inspect file content, magic bytes, etc.
+        Return False to let other parsers try.
+        """
+        ...
+
+    def parse(self, path: Path, extra_metadata: dict | None = None) -> Document:
+        """Parse the file and return a Document.
+
+        Args:
+            path: Path to the file to parse
+            extra_metadata: Optional metadata to merge into the document
+
+        Returns:
+            Document instance ready for ingestion
+        """
+        ...
+
+
+@runtime_checkable
+class APISource(Protocol):
+    """Protocol for API-based data sources.
+
+    Plugins implement this to sync data from external services.
+
+    Example:
+        class GitHubSource:
+            name = 'github'
+            source_type = 'github-source'
+
+            def configure(self, config: dict) -> None:
+                self._token = config['token']
+                self._repos = config.get('repos', [])  # From CLI --repo flags
+
+            def fetch(self, state: SyncState | None = None) -> tuple[list[Document], SyncState]:
+                # Use state.last_sync for incremental fetching
+                # Return (documents, new_state)
+                ...
+    """
+
+    name: str  # e.g., 'github'
+    source_type: str  # e.g., 'github-issue'
+
+    def configure(self, config: dict) -> None:
+        """Configure the source with settings from config file.
+
+        Config values may include resolved environment variables.
+
+        Args:
+            config: Source-specific configuration dict
+        """
+        ...
+
+    def fetch(self, state: SyncState | None = None) -> tuple[list[Document], SyncState]:
+        """Fetch documents from the external source.
+
+        Should support incremental fetching using the state object.
+
+        Args:
+            state: Previous sync state for incremental updates, or None for full sync
+
+        Returns:
+            Tuple of (list of documents, new sync state)
+        """
+        ...