okb 1.0.0__py3-none-any.whl

okb/llm/__init__.py

@@ -0,0 +1,86 @@
+"""LLM integration for document classification and enrichment.
+
+This package provides a provider-agnostic interface for LLM operations,
+with support for Claude API, AWS Bedrock, and response caching.
+
+Usage:
+    from okb.llm import get_llm, complete
+
+    # Get configured provider (returns None if disabled)
+    llm = get_llm()
+    if llm:
+        response = llm.complete("Summarize this document", system="Be concise")
+
+    # Or use convenience function with caching
+    response = complete("Classify this email")
+"""
+
+from .base import LLMProvider, LLMResponse
+from .filter import FilterAction, FilterResult, filter_document, filter_documents
+
+__all__ = [
+    "LLMProvider",
+    "LLMResponse",
+    "FilterAction",
+    "FilterResult",
+    "filter_document",
+    "filter_documents",
+    "get_llm",
+    "complete",
+]
+
+
+def get_llm() -> LLMProvider | None:
+    """Get the configured LLM provider, or None if disabled.
+
+    Reads configuration from the global config object.
+    Lazily initializes the provider on first call.
+
+    Returns:
+        Configured LLMProvider instance, or None if llm.provider is not set
+    """
+    from .providers import get_provider
+
+    return get_provider()
+
+
+def complete(
+    prompt: str,
+    system: str | None = None,
+    max_tokens: int = 1024,
+    use_cache: bool = True,
+) -> LLMResponse | None:
+    """Generate a completion using the configured LLM provider.
+
+    Convenience function that handles caching and provider initialization.
+
+    Args:
+        prompt: The user prompt to complete
+        system: Optional system prompt for context/instructions
+        max_tokens: Maximum tokens in the response
+        use_cache: Whether to use cached responses (default True)
+
+    Returns:
+        LLMResponse with the generated content, or None if LLM is disabled
+    """
+    from .cache import cache_response, get_cached
+    from .providers import get_provider
+
+    provider = get_provider()
+    if provider is None:
+        return None
+
+    # Check cache first
+    if use_cache:
+        cached = get_cached(prompt, system, provider.name)
+        if cached is not None:
+            return cached
+
+    # Generate new response
+    response = provider.complete(prompt, system=system, max_tokens=max_tokens)
+
+    # Cache the response
+    if use_cache:
+        cache_response(prompt, system, provider.name, response)
+
+    return response

okb/llm/base.py

@@ -0,0 +1,83 @@
+"""Protocol definitions and types for LLM providers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+
+@dataclass
+class LLMResponse:
+    """Standard response from LLM providers."""
+
+    content: str  # Raw text response
+    model: str  # Model that generated it
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+
+    @property
+    def total_tokens(self) -> int | None:
+        """Total tokens used (input + output)."""
+        if self.input_tokens is None or self.output_tokens is None:
+            return None
+        return self.input_tokens + self.output_tokens
+
+
+@runtime_checkable
+class LLMProvider(Protocol):
+    """Protocol for LLM backend providers.
+
+    Plugins implement this to add support for different LLM services.
+
+    Example:
+        class MyProvider:
+            name = 'my-llm'
+
+            def configure(self, config: dict) -> None:
+                self._api_key = config.get('api_key')
+
+            def complete(self, prompt: str, system: str | None = None) -> LLMResponse:
+                # Call the LLM API
+                ...
+
+            def is_available(self) -> bool:
+                return self._api_key is not None
+    """
+
+    name: str  # Provider identifier, e.g., "claude", "bedrock", "ollama"
+
+    def configure(self, config: dict) -> None:
+        """Configure the provider with settings from config.
+
+        Config values may include resolved environment variables.
+
+        Args:
+            config: Provider-specific configuration dict
+        """
+        ...
+
+    def complete(
+        self,
+        prompt: str,
+        system: str | None = None,
+        max_tokens: int = 1024,
+    ) -> LLMResponse:
+        """Generate a completion for the given prompt.
+
+        Args:
+            prompt: The user prompt to complete
+            system: Optional system prompt for context/instructions
+            max_tokens: Maximum tokens in the response
+
+        Returns:
+            LLMResponse with the generated content and metadata
+        """
+        ...
+
+    def is_available(self) -> bool:
+        """Check if the provider is configured and reachable.
+
+        Returns:
+            True if the provider can accept requests
+        """
+        ...

okb/llm/cache.py

@@ -0,0 +1,217 @@
+"""LLM response caching to avoid redundant API calls."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+import psycopg
+
+from .base import LLMResponse
+
+if TYPE_CHECKING:
+    pass
+
+
+def _compute_cache_key(prompt: str, system: str | None, model: str) -> str:
+    """Compute cache key from prompt, system, and model.
+
+    Args:
+        prompt: User prompt
+        system: System prompt (may be None)
+        model: Model name
+
+    Returns:
+        SHA256 hash of the combined inputs
+    """
+    content = f"{prompt}\n---\n{system or ''}\n---\n{model}"
+    return hashlib.sha256(content.encode()).hexdigest()
+
+
+def get_cached(
+    prompt: str,
+    system: str | None,
+    provider: str,
+    db_url: str | None = None,
+) -> LLMResponse | None:
+    """Retrieve a cached LLM response if available.
+
+    Args:
+        prompt: User prompt
+        system: System prompt
+        provider: Provider name (e.g., "claude")
+        db_url: Database URL (default: from config)
+
+    Returns:
+        Cached LLMResponse or None if not found
+    """
+    from ..config import config
+
+    if db_url is None:
+        db_url = config.db_url
+
+    # Get model from config for cache key
+    model = config.llm_model or "default"
+    content_hash = _compute_cache_key(prompt, system, model)
+
+    try:
+        with psycopg.connect(db_url) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    SELECT response FROM llm_cache
+                    WHERE content_hash = %s AND provider = %s AND model = %s
+                    """,
+                    (content_hash, provider, model),
+                )
+                row = cur.fetchone()
+                if row is None:
+                    return None
+
+                # Parse cached response
+                data = json.loads(row[0])
+                return LLMResponse(
+                    content=data["content"],
+                    model=data["model"],
+                    input_tokens=data.get("input_tokens"),
+                    output_tokens=data.get("output_tokens"),
+                )
+    except psycopg.Error:
+        # Cache miss on error - don't block on cache failures
+        return None
+
+
+def cache_response(
+    prompt: str,
+    system: str | None,
+    provider: str,
+    response: LLMResponse,
+    db_url: str | None = None,
+) -> None:
+    """Store an LLM response in the cache.
+
+    Args:
+        prompt: User prompt
+        system: System prompt
+        provider: Provider name
+        response: LLMResponse to cache
+        db_url: Database URL (default: from config)
+    """
+    from ..config import config
+
+    if db_url is None:
+        db_url = config.db_url
+
+    model = config.llm_model or "default"
+    content_hash = _compute_cache_key(prompt, system, model)
+
+    # Serialize response
+    data = {
+        "content": response.content,
+        "model": response.model,
+        "input_tokens": response.input_tokens,
+        "output_tokens": response.output_tokens,
+    }
+
+    try:
+        with psycopg.connect(db_url) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    INSERT INTO llm_cache (content_hash, provider, model, response)
+                    VALUES (%s, %s, %s, %s)
+                    ON CONFLICT (content_hash) DO UPDATE SET
+                        response = EXCLUDED.response,
+                        created_at = NOW()
+                    """,
+                    (content_hash, provider, model, json.dumps(data)),
+                )
+            conn.commit()
+    except psycopg.Error:
+        # Don't fail on cache write errors
+        pass
+
+
+def clear_cache(
+    older_than: datetime | None = None,
+    db_url: str | None = None,
+) -> int:
+    """Clear cached LLM responses.
+
+    Args:
+        older_than: Only clear entries older than this datetime.
+                   If None, clears all entries.
+        db_url: Database URL (default: from config)
+
+    Returns:
+        Number of entries deleted
+    """
+    from ..config import config
+
+    if db_url is None:
+        db_url = config.db_url
+
+    try:
+        with psycopg.connect(db_url) as conn:
+            with conn.cursor() as cur:
+                if older_than:
+                    cur.execute(
+                        "DELETE FROM llm_cache WHERE created_at < %s",
+                        (older_than,),
+                    )
+                else:
+                    cur.execute("DELETE FROM llm_cache")
+                deleted = cur.rowcount
+            conn.commit()
+            return deleted
+    except psycopg.Error:
+        return 0
+
+
+def get_cache_stats(db_url: str | None = None) -> dict:
+    """Get statistics about the LLM cache.
+
+    Args:
+        db_url: Database URL (default: from config)
+
+    Returns:
+        Dict with cache statistics
+    """
+    from ..config import config
+
+    if db_url is None:
+        db_url = config.db_url
+
+    try:
+        with psycopg.connect(db_url) as conn:
+            with conn.cursor() as cur:
+                # Total entries
+                cur.execute("SELECT COUNT(*) FROM llm_cache")
+                total = cur.fetchone()[0]
+
+                # Entries by provider/model
+                cur.execute(
+                    """
+                    SELECT provider, model, COUNT(*) as count
+                    FROM llm_cache
+                    GROUP BY provider, model
+                    ORDER BY count DESC
+                    """
+                )
+                by_provider = [
+                    {"provider": r[0], "model": r[1], "count": r[2]} for r in cur.fetchall()
+                ]
+
+                # Oldest entry
+                cur.execute("SELECT MIN(created_at) FROM llm_cache")
+                oldest = cur.fetchone()[0]
+
+                return {
+                    "total_entries": total,
+                    "by_provider": by_provider,
+                    "oldest_entry": oldest.isoformat() if oldest else None,
+                }
+    except psycopg.Error:
+        return {"total_entries": 0, "by_provider": [], "oldest_entry": None}

okb/llm/filter.py

@@ -0,0 +1,187 @@
+"""LLM-based document filtering for pre-ingest classification."""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..ingest import Document
+
+
+class FilterAction(Enum):
+    """Actions the filter can take on a document."""
+
+    INGEST = "ingest"  # Process normally
+    SKIP = "skip"  # Don't ingest
+    REVIEW = "review"  # Flag for manual review (still ingest)
+
+
+@dataclass
+class FilterResult:
+    """Result of LLM filtering on a document."""
+
+    action: FilterAction
+    reason: str
+    confidence: float | None = None  # Optional confidence score 0-1
+
+    @property
+    def should_ingest(self) -> bool:
+        """Whether the document should be ingested."""
+        return self.action in (FilterAction.INGEST, FilterAction.REVIEW)
+
+
+DEFAULT_SYSTEM_PROMPT = """\
+You are a document classifier. Analyze the document and decide whether it should
+be ingested into a knowledge base.
+
+Respond with a JSON object containing:
+- "action": one of "ingest", "skip", or "review"
+- "reason": brief explanation (1 sentence)
+
+Use these guidelines:
+- "ingest": valuable content worth indexing (notes, docs, important emails, etc.)
+- "skip": low-value content (spam, marketing, automated notifications, duplicates)
+- "review": uncertain cases that need human review
+
+Respond ONLY with the JSON object, no other text."""
+
+
+def _build_filter_prompt(document: Document, custom_prompt: str | None = None) -> str:
+    """Build the prompt for document filtering.
+
+    Args:
+        document: Document to classify
+        custom_prompt: Optional custom instructions to append
+
+    Returns:
+        Formatted prompt string
+    """
+    parts = [
+        f"Title: {document.title}",
+        f"Source: {document.source_type}",
+    ]
+
+    if document.metadata and document.metadata.tags:
+        parts.append(f"Tags: {', '.join(document.metadata.tags)}")
+
+    # Include content preview (truncate if too long)
+    content_preview = document.content[:2000]
+    if len(document.content) > 2000:
+        content_preview += "\n[... truncated ...]"
+
+    parts.append(f"\nContent:\n{content_preview}")
+
+    if custom_prompt:
+        parts.append(f"\nAdditional instructions: {custom_prompt}")
+
+    return "\n".join(parts)
+
+
+def _parse_filter_response(response: str) -> FilterResult:
+    """Parse the LLM response into a FilterResult.
+
+    Args:
+        response: Raw LLM response text
+
+    Returns:
+        Parsed FilterResult
+
+    Raises:
+        ValueError: If response cannot be parsed
+    """
+    # Try to extract JSON from response
+    # Handle cases where LLM wraps in markdown code blocks
+    json_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", response, re.DOTALL)
+    if json_match:
+        json_str = json_match.group(1)
+    else:
+        # Try to find raw JSON object
+        json_match = re.search(r"\{[^{}]*\}", response, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(0)
+        else:
+            # Fallback: assume entire response is JSON
+            json_str = response.strip()
+
+    try:
+        data = json.loads(json_str)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Failed to parse filter response as JSON: {e}")
+
+    action_str = data.get("action", "").lower()
+    try:
+        action = FilterAction(action_str)
+    except ValueError:
+        # Default to ingest if action is invalid
+        action = FilterAction.INGEST
+
+    reason = data.get("reason", "No reason provided")
+    confidence = data.get("confidence")
+
+    return FilterResult(action=action, reason=reason, confidence=confidence)
+
+
+def filter_document(
+    document: Document,
+    custom_prompt: str | None = None,
+    use_cache: bool = True,
+) -> FilterResult:
+    """Filter a single document using the configured LLM.
+
+    Args:
+        document: Document to filter
+        custom_prompt: Optional custom classification instructions
+        use_cache: Whether to use cached responses
+
+    Returns:
+        FilterResult with action and reason
+
+    Raises:
+        RuntimeError: If LLM is not configured
+    """
+    from . import complete
+
+    prompt = _build_filter_prompt(document, custom_prompt)
+    response = complete(prompt, system=DEFAULT_SYSTEM_PROMPT, use_cache=use_cache)
+
+    if response is None:
+        # LLM not configured - default to ingest
+        return FilterResult(
+            action=FilterAction.INGEST,
+            reason="LLM not configured, defaulting to ingest",
+        )
+
+    try:
+        return _parse_filter_response(response.content)
+    except ValueError as e:
+        # Parse error - default to ingest with warning
+        return FilterResult(
+            action=FilterAction.INGEST,
+            reason=f"Failed to parse LLM response: {e}",
+        )
+
+
+def filter_documents(
+    documents: list[Document],
+    custom_prompt: str | None = None,
+    use_cache: bool = True,
+) -> list[tuple[Document, FilterResult]]:
+    """Filter multiple documents.
+
+    Args:
+        documents: List of documents to filter
+        custom_prompt: Optional custom classification instructions
+        use_cache: Whether to use cached responses
+
+    Returns:
+        List of (document, filter_result) tuples
+    """
+    results = []
+    for doc in documents:
+        result = filter_document(doc, custom_prompt=custom_prompt, use_cache=use_cache)
+        results.append((doc, result))
+    return results