chore: import upstream snapshot with attribution

open_notebook/utils/README.md

@@ -0,0 +1,188 @@
+# ContextBuilder
+
+A flexible and generic ContextBuilder class for the Open Notebook project that can handle any parameters and build context from sources, notebooks, insights, and notes.
+
+## Features
+
+- **Flexible Parameters**: Accepts any parameters via `**kwargs` for future extensibility
+- **Priority-based Management**: Automatic prioritization and sorting of context items
+- **Token Counting**: Built-in token counting and truncation to fit limits
+- **Deduplication**: Automatic removal of duplicate items based on ID
+- **Type-based Grouping**: Separates sources, notes, and insights in output
+- **Async Support**: Fully async for database operations
+
+## Basic Usage
+
+```python
+from open_notebook.utils.context_builder import ContextBuilder, ContextConfig
+
+# Simple notebook context
+builder = ContextBuilder(notebook_id="notebook:123")
+context = await builder.build()
+
+# Single source with insights
+builder = ContextBuilder(
+    source_id="source:456",
+    include_insights=True,
+    max_tokens=2000
+)
+context = await builder.build()
+```
+
+## Convenience Functions
+
+```python
+from open_notebook.utils.context_builder import (
+    build_notebook_context,
+    build_source_context,
+    build_mixed_context
+)
+
+# Build notebook context
+context = await build_notebook_context(
+    notebook_id="notebook:123",
+    max_tokens=5000
+)
+
+# Build single source context
+context = await build_source_context(
+    source_id="source:456",
+    include_insights=True
+)
+
+# Build mixed context
+context = await build_mixed_context(
+    source_ids=["source:1", "source:2"],
+    note_ids=["note:1", "note:2"],
+    max_tokens=3000
+)
+```
+
+## Advanced Configuration
+
+```python
+from open_notebook.utils.context_builder import ContextConfig
+
+# Custom configuration
+config = ContextConfig(
+    sources={
+        "source:doc1": "insights",
+        "source:doc2": "full content", 
+        "source:doc3": "not in"  # Exclude
+    },
+    notes={
+        "note:summary": "full content",
+        "note:draft": "not in"  # Exclude
+    },
+    include_insights=True,
+    max_tokens=3000,
+    priority_weights={
+        "source": 120,  # Higher priority
+        "note": 80,     # Medium priority  
+        "insight": 100  # High priority
+    }
+)
+
+builder = ContextBuilder(
+    notebook_id="notebook:project",
+    context_config=config
+)
+context = await builder.build()
+```
+
+## Programmatic Item Management
+
+```python
+from open_notebook.utils.context_builder import ContextItem
+
+builder = ContextBuilder()
+
+# Add custom items
+item = ContextItem(
+    id="source:important",
+    type="source",
+    content={"title": "Key Document", "summary": "..."},
+    priority=150  # Very high priority
+)
+builder.add_item(item)
+
+# Apply management operations
+builder.remove_duplicates()
+builder.prioritize()
+builder.truncate_to_fit(1000)
+
+context = builder._format_response()
+```
+
+## Flexible Parameters
+
+The ContextBuilder accepts any parameters via `**kwargs`, making it extensible for future features:
+
+```python
+builder = ContextBuilder(
+    notebook_id="notebook:123",
+    include_insights=True,
+    max_tokens=2000,
+    
+    # Custom parameters for future extensions
+    user_id="user:456",
+    custom_filter="advanced",
+    experimental_feature=True
+)
+
+# Access custom parameters
+user_id = builder.params.get('user_id')
+```
+
+## Output Format
+
+The ContextBuilder returns a structured response:
+
+```python
+{
+    "sources": [...],           # List of source contexts
+    "notes": [...],             # List of note contexts  
+    "insights": [...],          # List of insight contexts
+    "total_tokens": 1234,       # Total token count
+    "total_items": 10,          # Total number of items
+    "notebook_id": "notebook:123",  # If provided
+    "metadata": {
+        "source_count": 5,
+        "note_count": 3,
+        "insight_count": 2,
+        "config": {
+            "include_insights": true,
+            "include_notes": true,
+            "max_tokens": 2000
+        }
+    }
+}
+```
+
+## Architecture
+
+The ContextBuilder follows these design principles:
+
+1. **Separation of Concerns**: Context building, item management, and formatting are separate
+2. **Extensibility**: Uses `**kwargs` and flexible configuration for future features
+3. **Performance**: Token-aware truncation and efficient deduplication
+4. **Type Safety**: Proper type hints and data classes for structure
+5. **Error Handling**: Graceful handling of missing items and database errors
+
+## Integration
+
+The ContextBuilder integrates seamlessly with the existing Open Notebook architecture:
+
+- Uses existing domain models (`Source`, `Notebook`, `Note`)
+- Leverages the repository pattern for database access
+- Follows the same async patterns as other services
+- Integrates with the token counting utilities
+
+## Error Handling
+
+The ContextBuilder handles errors gracefully:
+
+- Missing notebooks/sources/notes are logged but don't stop execution
+- Database errors are wrapped in `DatabaseOperationError`
+- Invalid parameters raise `InvalidInputError`
+- All errors include detailed context information

open_notebook/utils/__init__.py

@@ -0,0 +1,72 @@
+"""
+Utils package for Open Notebook.
+
+To avoid circular imports, import functions directly:
+- from open_notebook.utils.context_builder import build_notebook_context, build_source_context
+- from open_notebook.utils import token_count, compare_versions
+- from open_notebook.utils.chunking import chunk_text, detect_content_type, ContentType
+- from open_notebook.utils.embedding import generate_embedding, generate_embeddings
+- from open_notebook.utils.encryption import encrypt_value, decrypt_value
+"""
+
+from .chunking import (
+    CHUNK_SIZE,
+    ContentType,
+    chunk_text,
+    detect_content_type,
+    detect_content_type_from_extension,
+    detect_content_type_from_heuristics,
+)
+from .embedding import (
+    generate_embedding,
+    generate_embeddings,
+    mean_pool_embeddings,
+)
+from .encryption import (
+    decrypt_value,
+    encrypt_value,
+)
+from .model_utils import full_model_dump
+from .text_utils import (
+    clean_thinking_content,
+    parse_thinking_content,
+    remove_non_ascii,
+    remove_non_printable,
+)
+from .token_utils import token_cost, token_count
+from .version_utils import (
+    compare_versions,
+    get_installed_version,
+    get_version_from_github,
+)
+
+__all__ = [
+    # Chunking
+    "CHUNK_SIZE",
+    "ContentType",
+    "chunk_text",
+    "detect_content_type",
+    "detect_content_type_from_extension",
+    "detect_content_type_from_heuristics",
+    # Embedding
+    "generate_embedding",
+    "generate_embeddings",
+    "mean_pool_embeddings",
+    # Text utils
+    "remove_non_ascii",
+    "remove_non_printable",
+    "parse_thinking_content",
+    "clean_thinking_content",
+    # Token utils
+    "token_count",
+    "token_cost",
+    # Version utils
+    "compare_versions",
+    "get_installed_version",
+    "get_version_from_github",
+    # Encryption utils
+    "decrypt_value",
+    "encrypt_value",
+    # Model utils
+    "full_model_dump",
+]

open_notebook/utils/chunking.py

@@ -0,0 +1,494 @@
+"""
+Chunking utilities for Open Notebook.
+
+Provides content-type detection and smart text chunking for embedding operations.
+Supports HTML, Markdown, and plain text with appropriate splitters for each type.
+
+Key functions:
+- detect_content_type(): Detects content type from file extension or content heuristics
+- chunk_text(): Splits text into chunks using appropriate splitter for content type
+
+Environment Variables:
+    OPEN_NOTEBOOK_CHUNK_SIZE: Maximum chunk size in tokens (default: 400)
+    OPEN_NOTEBOOK_CHUNK_OVERLAP: Overlap between chunks in tokens (default: 15% of CHUNK_SIZE)
+    OPEN_NOTEBOOK_MIN_CHUNK_SIZE: Minimum chunk size in tokens (default: 5)
+"""
+
+import os
+import re
+from enum import Enum
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from langchain_text_splitters import (
+    HTMLHeaderTextSplitter,
+    MarkdownHeaderTextSplitter,
+    RecursiveCharacterTextSplitter,
+)
+from loguru import logger
+
+from .token_utils import token_count
+
+
+def _get_chunk_size() -> int:
+    """Get chunk size from environment variable or use default."""
+    chunk_size_str = os.getenv("OPEN_NOTEBOOK_CHUNK_SIZE")
+    if chunk_size_str:
+        try:
+            chunk_size = int(chunk_size_str)
+            if chunk_size < 100:
+                logger.warning(
+                    f"OPEN_NOTEBOOK_CHUNK_SIZE ({chunk_size}) is too small. "
+                    f"Using minimum value of 100."
+                )
+                return 100
+            if chunk_size > 8192:
+                logger.warning(
+                    f"OPEN_NOTEBOOK_CHUNK_SIZE ({chunk_size}) is very large. "
+                    f"This may cause issues with some embedding models."
+                )
+            logger.info(f"Using custom chunk size: {chunk_size} tokens")
+            return chunk_size
+        except ValueError:
+            logger.warning(
+                f"Invalid OPEN_NOTEBOOK_CHUNK_SIZE value: '{chunk_size_str}'. "
+                f"Using default: 400"
+            )
+    return 400
+
+
+def _get_chunk_overlap(chunk_size: int) -> int:
+    """Get chunk overlap from environment variable or calculate default (15% of chunk size)."""
+    overlap_str = os.getenv("OPEN_NOTEBOOK_CHUNK_OVERLAP")
+    if overlap_str:
+        try:
+            overlap = int(overlap_str)
+            if overlap < 0:
+                logger.warning(
+                    f"OPEN_NOTEBOOK_CHUNK_OVERLAP ({overlap}) cannot be negative. "
+                    f"Using 0."
+                )
+                return 0
+            if overlap >= chunk_size:
+                logger.warning(
+                    f"OPEN_NOTEBOOK_CHUNK_OVERLAP ({overlap}) cannot be >= chunk size ({chunk_size}). "
+                    f"Using 15% of chunk size: {int(chunk_size * 0.15)}"
+                )
+                return int(chunk_size * 0.15)
+            logger.info(f"Using custom chunk overlap: {overlap} tokens")
+            return overlap
+        except ValueError:
+            logger.warning(
+                f"Invalid OPEN_NOTEBOOK_CHUNK_OVERLAP value: '{overlap_str}'. "
+                f"Using default: 15% of chunk size"
+            )
+    return int(chunk_size * 0.15)
+
+
+def _get_min_chunk_size() -> int:
+    """Get minimum chunk size from environment variable or use default.
+
+    Chunks below this token count are dropped. Some splitters (notably the
+    HTML header splitter on complex pages) can emit single-character or
+    punctuation-only chunks that produce useless or null embeddings —
+    llama.cpp's OpenAI-compatible endpoint, for example, returns null vector
+    elements for such inputs and crashes downstream parsing.
+    """
+    raw = os.getenv("OPEN_NOTEBOOK_MIN_CHUNK_SIZE")
+    if raw is None:
+        return 5
+    try:
+        value = int(raw)
+        if value < 0:
+            logger.warning(
+                f"OPEN_NOTEBOOK_MIN_CHUNK_SIZE ({value}) cannot be negative. Using 0."
+            )
+            return 0
+        return value
+    except ValueError:
+        logger.warning(
+            f"Invalid OPEN_NOTEBOOK_MIN_CHUNK_SIZE value: '{raw}'. Using default: 5"
+        )
+        return 5
+
+
+# Constants (computed at import time from environment variables)
+CHUNK_SIZE = _get_chunk_size()
+CHUNK_OVERLAP = _get_chunk_overlap(CHUNK_SIZE)
+MIN_CHUNK_SIZE = _get_min_chunk_size()
+HIGH_CONFIDENCE_THRESHOLD = 0.8  # Threshold for heuristics to override extension
+
+logger.debug(
+    f"Chunking configuration: CHUNK_SIZE={CHUNK_SIZE}, "
+    f"CHUNK_OVERLAP={CHUNK_OVERLAP}, MIN_CHUNK_SIZE={MIN_CHUNK_SIZE}"
+)
+
+
+class ContentType(Enum):
+    """Content type for chunking strategy selection."""
+
+    HTML = "html"
+    MARKDOWN = "markdown"
+    PLAIN = "plain"
+
+
+# File extension mappings
+_EXTENSION_TO_CONTENT_TYPE = {
+    # HTML
+    ".html": ContentType.HTML,
+    ".htm": ContentType.HTML,
+    ".xhtml": ContentType.HTML,
+    # Markdown
+    ".md": ContentType.MARKDOWN,
+    ".markdown": ContentType.MARKDOWN,
+    ".mdown": ContentType.MARKDOWN,
+    ".mkd": ContentType.MARKDOWN,
+    # Plain text (explicit)
+    ".txt": ContentType.PLAIN,
+    ".text": ContentType.PLAIN,
+    # Code files (treat as plain)
+    ".py": ContentType.PLAIN,
+    ".js": ContentType.PLAIN,
+    ".ts": ContentType.PLAIN,
+    ".java": ContentType.PLAIN,
+    ".c": ContentType.PLAIN,
+    ".cpp": ContentType.PLAIN,
+    ".go": ContentType.PLAIN,
+    ".rs": ContentType.PLAIN,
+    ".rb": ContentType.PLAIN,
+    ".php": ContentType.PLAIN,
+    ".sh": ContentType.PLAIN,
+    ".bash": ContentType.PLAIN,
+    ".zsh": ContentType.PLAIN,
+    ".sql": ContentType.PLAIN,
+    ".json": ContentType.PLAIN,
+    ".yaml": ContentType.PLAIN,
+    ".yml": ContentType.PLAIN,
+    ".xml": ContentType.PLAIN,
+    ".csv": ContentType.PLAIN,
+    ".tsv": ContentType.PLAIN,
+}
+
+
+def detect_content_type_from_extension(
+    file_path: Optional[str],
+) -> Optional[ContentType]:
+    """
+    Detect content type from file extension.
+
+    Args:
+        file_path: Path to the file (can be full path or just filename)
+
+    Returns:
+        ContentType if extension is recognized, None otherwise
+    """
+    if not file_path:
+        return None
+
+    try:
+        extension = Path(file_path).suffix.lower()
+        return _EXTENSION_TO_CONTENT_TYPE.get(extension)
+    except Exception:
+        return None
+
+
+def detect_content_type_from_heuristics(text: str) -> Tuple[ContentType, float]:
+    """
+    Detect content type using content heuristics.
+
+    Args:
+        text: The text content to analyze
+
+    Returns:
+        Tuple of (ContentType, confidence_score) where confidence is 0.0-1.0
+    """
+    if not text or len(text) < 10:
+        return ContentType.PLAIN, 0.5
+
+    # Sample first 5000 chars for efficiency
+    sample = text[:5000]
+
+    # Check HTML first (most specific patterns)
+    html_score = _calculate_html_score(sample)
+    if html_score >= HIGH_CONFIDENCE_THRESHOLD:
+        return ContentType.HTML, html_score
+
+    # Check Markdown
+    markdown_score = _calculate_markdown_score(sample)
+    if markdown_score >= HIGH_CONFIDENCE_THRESHOLD:
+        return ContentType.MARKDOWN, markdown_score
+
+    # Return the higher scoring type, or PLAIN if both are low
+    if html_score > markdown_score and html_score > 0.3:
+        return ContentType.HTML, html_score
+    elif markdown_score > 0.3:
+        return ContentType.MARKDOWN, markdown_score
+    else:
+        return ContentType.PLAIN, 0.6
+
+
+def _calculate_html_score(text: str) -> float:
+    """Calculate confidence score for HTML content."""
+    score = 0.0
+    indicators = 0
+
+    # Strong indicators
+    if re.search(r"<!DOCTYPE\s+html", text, re.IGNORECASE):
+        score += 0.4
+        indicators += 1
+
+    if re.search(r"<html[\s>]", text, re.IGNORECASE):
+        score += 0.3
+        indicators += 1
+
+    # Structural tags
+    structural_tags = ["<head", "<body", "<div", "<span", "<p>", "<table", "<form"]
+    for tag in structural_tags:
+        if tag.lower() in text.lower():
+            score += 0.1
+            indicators += 1
+            if indicators >= 5:
+                break
+
+    # Header tags
+    if re.search(r"<h[1-6][\s>]", text, re.IGNORECASE):
+        score += 0.15
+        indicators += 1
+
+    # Closing tags pattern
+    if re.search(r"</\w+>", text):
+        score += 0.1
+        indicators += 1
+
+    return min(score, 1.0)
+
+
+def _calculate_markdown_score(text: str) -> float:
+    """Calculate confidence score for Markdown content."""
+    score = 0.0
+    indicators = 0
+
+    # Headers (# ## ###) - strong indicator
+    header_matches = len(re.findall(r"^#{1,6}\s+.+", text, re.MULTILINE))
+    if header_matches >= 3:
+        score += 0.35
+        indicators += 1
+    elif header_matches >= 1:
+        score += 0.2
+        indicators += 1
+
+    # Links [text](url) - strong indicator
+    link_matches = len(re.findall(r"\[.+?\]\(.+?\)", text))
+    if link_matches >= 2:
+        score += 0.25
+        indicators += 1
+    elif link_matches >= 1:
+        score += 0.15
+        indicators += 1
+
+    # Code blocks ``` - strong indicator
+    if re.search(r"^```", text, re.MULTILINE):
+        score += 0.2
+        indicators += 1
+
+    # Inline code `code`
+    if re.search(r"`[^`]+`", text):
+        score += 0.1
+        indicators += 1
+
+    # Lists (-, *, +, or numbered)
+    list_matches = len(re.findall(r"^[\*\-\+]\s+", text, re.MULTILINE))
+    list_matches += len(re.findall(r"^\d+\.\s+", text, re.MULTILINE))
+    if list_matches >= 3:
+        score += 0.15
+        indicators += 1
+    elif list_matches >= 1:
+        score += 0.08
+        indicators += 1
+
+    # Bold/italic
+    if re.search(r"\*\*.+?\*\*|__.+?__", text):
+        score += 0.1
+        indicators += 1
+
+    # Blockquotes
+    if re.search(r"^>\s+", text, re.MULTILINE):
+        score += 0.1
+        indicators += 1
+
+    return min(score, 1.0)
+
+
+def detect_content_type(text: str, file_path: Optional[str] = None) -> ContentType:
+    """
+    Detect content type using file extension (primary) and heuristics (fallback).
+
+    Strategy:
+    1. If file extension is available and recognized, use it as primary
+    2. If no extension or generic extension (.txt), use heuristics
+    3. Heuristics can override extension only with very high confidence
+
+    Args:
+        text: The text content
+        file_path: Optional file path for extension-based detection
+
+    Returns:
+        Detected ContentType
+    """
+    # Try extension-based detection first
+    extension_type = detect_content_type_from_extension(file_path)
+
+    # Get heuristic-based detection
+    heuristic_type, confidence = detect_content_type_from_heuristics(text)
+
+    # If no extension or generic extension, use heuristics
+    if extension_type is None:
+        logger.debug(
+            f"No file extension, using heuristics: {heuristic_type.value} "
+            f"(confidence: {confidence:.2f})"
+        )
+        return heuristic_type
+
+    # If extension suggests plain text but heuristics are very confident, override
+    if extension_type == ContentType.PLAIN and confidence >= HIGH_CONFIDENCE_THRESHOLD:
+        logger.debug(
+            f"Extension suggests plain, but heuristics override with "
+            f"{heuristic_type.value} (confidence: {confidence:.2f})"
+        )
+        return heuristic_type
+
+    # Otherwise trust the extension
+    logger.debug(f"Using extension-based content type: {extension_type.value}")
+    return extension_type
+
+
+def _get_html_splitter() -> HTMLHeaderTextSplitter:
+    """Get HTML header splitter configured for h1, h2, h3."""
+    headers_to_split_on = [
+        ("h1", "Header 1"),
+        ("h2", "Header 2"),
+        ("h3", "Header 3"),
+    ]
+    return HTMLHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
+
+
+def _get_markdown_splitter() -> MarkdownHeaderTextSplitter:
+    """Get Markdown header splitter configured for #, ##, ###."""
+    headers_to_split_on = [
+        ("#", "Header 1"),
+        ("##", "Header 2"),
+        ("###", "Header 3"),
+    ]
+    return MarkdownHeaderTextSplitter(
+        headers_to_split_on=headers_to_split_on,
+        strip_headers=False,
+    )
+
+
+def _get_plain_splitter() -> RecursiveCharacterTextSplitter:
+    """Get plain text splitter using CHUNK_SIZE and CHUNK_OVERLAP constants."""
+    return RecursiveCharacterTextSplitter(
+        chunk_size=CHUNK_SIZE,
+        chunk_overlap=CHUNK_OVERLAP,
+        length_function=token_count,
+        separators=["\n\n", "\n", ". ", ", ", " ", ""],
+    )
+
+
+def _apply_secondary_chunking(chunks: List[str]) -> List[str]:
+    """
+    Apply secondary chunking to ensure no chunk exceeds CHUNK_SIZE tokens.
+
+    Used when primary splitters (HTML/Markdown) produce oversized chunks.
+    """
+    result = []
+    secondary_splitter = _get_plain_splitter()
+
+    for chunk in chunks:
+        if token_count(chunk) > CHUNK_SIZE:
+            # Split oversized chunk
+            sub_chunks = secondary_splitter.split_text(chunk)
+            result.extend(sub_chunks)
+        else:
+            result.append(chunk)
+
+    return result
+
+
+def chunk_text(
+    text: str,
+    content_type: Optional[ContentType] = None,
+    file_path: Optional[str] = None,
+) -> List[str]:
+    """
+    Split text into chunks using appropriate splitter for content type.
+
+    Args:
+        text: The text to chunk
+        content_type: Optional explicit content type (auto-detected if not provided)
+        file_path: Optional file path for content type detection
+
+    Returns:
+        List of text chunks, each approximately <= CHUNK_SIZE tokens
+    """
+    if not text or not text.strip():
+        return []
+
+    # Short text doesn't need chunking
+    text_tokens = token_count(text)
+    if text_tokens <= CHUNK_SIZE:
+        return [text]
+
+    # Detect content type if not provided
+    if content_type is None:
+        content_type = detect_content_type(text, file_path)
+
+    logger.debug(f"Chunking text with content type: {content_type.value}")
+
+    # Select appropriate splitter
+    chunks: List[str]
+    if content_type == ContentType.HTML:
+        html_splitter = _get_html_splitter()
+        # HTML splitter returns Document objects
+        docs = html_splitter.split_text(text)
+        chunks = [
+            doc.page_content if hasattr(doc, "page_content") else str(doc)
+            for doc in docs
+        ]
+    elif content_type == ContentType.MARKDOWN:
+        md_splitter = _get_markdown_splitter()
+        # Markdown splitter returns Document objects
+        docs = md_splitter.split_text(text)
+        chunks = [
+            doc.page_content if hasattr(doc, "page_content") else str(doc)
+            for doc in docs
+        ]
+    else:
+        # Plain text - use recursive splitter directly
+        chunks = _get_plain_splitter().split_text(text)
+
+    # Apply secondary chunking if needed (for HTML/Markdown that may produce large chunks)
+    if content_type in (ContentType.HTML, ContentType.MARKDOWN):
+        chunks = _apply_secondary_chunking(chunks)
+
+    # Filter out empty chunks
+    chunks = [c.strip() for c in chunks if c and c.strip()]
+
+    # Drop chunks below the minimum token threshold. These are typically
+    # punctuation or single-character fragments left over from header-based
+    # splitters; embedding them is wasteful and some providers return null
+    # vector elements for such inputs (which then crash response parsing).
+    # Only filter when more than one chunk exists and at least one chunk
+    # would survive — never return an empty list because of this filter.
+    if MIN_CHUNK_SIZE > 0 and len(chunks) > 1:
+        kept = [c for c in chunks if token_count(c) >= MIN_CHUNK_SIZE]
+        if kept:
+            dropped = len(chunks) - len(kept)
+            if dropped > 0:
+                logger.debug(
+                    f"Dropped {dropped} chunk(s) below MIN_CHUNK_SIZE={MIN_CHUNK_SIZE} tokens"
+                )
+            chunks = kept
+
+    logger.debug(f"Created {len(chunks)} chunks from {text_tokens} tokens")
+    return chunks

open_notebook/utils/context_builder.py

@@ -0,0 +1,208 @@
+"""Context building for chat and podcast generation.
+
+This is the single implementation behind:
+
+- ``POST /api/chat/context`` (`api/routers/chat.py`) — assembles notebook
+  context from a source/note inclusion config, via
+  :func:`build_notebook_context`.
+- the source-chat graph (`open_notebook/graphs/source_chat.py`) — assembles
+  a single source plus its insights under a token budget, via
+  :func:`build_source_context`.
+
+The inclusion config uses string matching on human-readable status values
+("not in context", "insights", "full content"). That protocol is shared with
+the frontend — do not change it here.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple
+
+from loguru import logger
+
+from open_notebook.domain.notebook import (
+    Note,
+    Notebook,
+    Source,
+    SourceInsight,
+)
+from open_notebook.exceptions import DatabaseOperationError, NotFoundError
+
+from .token_utils import token_count
+
+
+def _ensure_prefix(table: str, record_id: str) -> str:
+    """Ensure a record ID carries its table prefix (`table:id`)."""
+    prefix = f"{table}:"
+    return record_id if record_id.startswith(prefix) else f"{prefix}{record_id}"
+
+
+async def build_notebook_context(
+    notebook: Notebook,
+    context_config: Optional[Dict[str, Any]],
+) -> Tuple[Dict[str, list], str]:
+    """Assemble source/note context for a notebook.
+
+    With a config, each entry's status string decides inclusion: "not in"
+    skips it, "insights" includes the short source context, "full content"
+    includes the long context (notes only support "full content"). Without a
+    config, every source and note is included with its short context.
+
+    Failures on individual items are logged and skipped — one broken record
+    never fails the whole request.
+
+    Returns:
+        ({"sources": [...], "notes": [...]}, concatenated str() of every
+        included context dict — used for token/char counting).
+    """
+    context_data: Dict[str, list] = {"sources": [], "notes": []}
+    total_content = ""
+
+    if context_config:
+        for source_id, status in context_config.get("sources", {}).items():
+            if "not in" in status:
+                continue
+
+            try:
+                full_source_id = _ensure_prefix("source", source_id)
+
+                try:
+                    source = await Source.get(full_source_id)
+                except Exception:
+                    continue
+
+                if "insights" in status:
+                    source_context = await source.get_context(context_size="short")
+                    context_data["sources"].append(source_context)
+                    total_content += str(source_context)
+                elif "full content" in status:
+                    source_context = await source.get_context(context_size="long")
+                    context_data["sources"].append(source_context)
+                    total_content += str(source_context)
+            except Exception as e:
+                logger.warning(f"Error processing source {source_id}: {str(e)}")
+                continue
+
+        for note_id, status in context_config.get("notes", {}).items():
+            if "not in" in status:
+                continue
+
+            try:
+                full_note_id = _ensure_prefix("note", note_id)
+                note = await Note.get(full_note_id)
+                if not note:
+                    continue
+
+                if "full content" in status:
+                    note_context = note.get_context(context_size="long")
+                    context_data["notes"].append(note_context)
+                    total_content += str(note_context)
+            except Exception as e:
+                logger.warning(f"Error processing note {note_id}: {str(e)}")
+                continue
+    else:
+        # Default behavior - include all sources and notes with short context
+        sources = await notebook.get_sources()
+        try:
+            insights_by_source = await SourceInsight.get_for_sources(
+                [source.id for source in sources if source.id]
+            )
+        except Exception as e:
+            # Match the per-source fallback below: a hiccup fetching
+            # insights shouldn't fail the whole context request.
+            logger.warning(f"Error batch-fetching source insights: {str(e)}")
+            insights_by_source = {}
+        for source in sources:
+            try:
+                source_context = await source.get_context(
+                    context_size="short",
+                    insights=insights_by_source.get(source.id or "", []),
+                )
+                context_data["sources"].append(source_context)
+                total_content += str(source_context)
+            except Exception as e:
+                logger.warning(f"Error processing source {source.id}: {str(e)}")
+                continue
+
+        notes = await notebook.get_notes()
+        for note in notes:
+            try:
+                note_context = note.get_context(context_size="short")
+                context_data["notes"].append(note_context)
+                total_content += str(note_context)
+            except Exception as e:
+                logger.warning(f"Error processing note {note.id}: {str(e)}")
+                continue
+
+    return context_data, total_content
+
+
+async def build_source_context(
+    source_id: str, max_tokens: Optional[int] = None
+) -> Dict[str, Any]:
+    """Assemble a single source's short context plus its insights.
+
+    Used by the source-chat graph. If `max_tokens` is given, insights are
+    dropped (last-fetched first) until the total fits — the source itself is
+    always kept.
+
+    Returns a dict with "sources", "notes" (always empty), "insights",
+    "total_tokens", "total_items" and per-type counts in "metadata".
+    """
+    try:
+        sources: list = []
+        insights: list = []
+        item_tokens: list[int] = []
+
+        try:
+            full_source_id = _ensure_prefix("source", source_id)
+            source = await Source.get(full_source_id)
+        except NotFoundError:
+            source = None
+
+        if source:
+            source_context = await source.get_context(context_size="short")
+            sources.append(source_context)
+            item_tokens.append(token_count(str(source_context)))
+
+            for insight in await source.get_insights():
+                insight_content = {
+                    "id": insight.id,
+                    "source_id": source.id,
+                    "insight_type": insight.insight_type,
+                    "content": insight.content,
+                }
+                insights.append(insight_content)
+                item_tokens.append(token_count(str(insight_content)))
+        else:
+            logger.warning(f"Source {source_id} not found")
+
+        # Truncate to the token budget: drop insights from the end (the
+        # source, added first, is dropped only if it alone exceeds the budget).
+        total_tokens = sum(item_tokens)
+        if max_tokens:
+            while total_tokens > max_tokens and item_tokens:
+                total_tokens -= item_tokens.pop()
+                if insights:
+                    insights.pop()
+                else:
+                    sources.pop()
+
+        total_items = len(sources) + len(insights)
+        logger.info(f"Built context with {total_items} items, {total_tokens} tokens")
+
+        return {
+            "sources": sources,
+            "notes": [],
+            "insights": insights,
+            "total_tokens": total_tokens,
+            "total_items": total_items,
+            "metadata": {
+                "source_count": len(sources),
+                "note_count": 0,
+                "insight_count": len(insights),
+            },
+        }
+    except Exception as e:
+        logger.error(f"Error building context: {str(e)}")
+        raise DatabaseOperationError(f"Failed to build context: {str(e)}")

open_notebook/utils/embedding.py

@@ -0,0 +1,269 @@
+"""
+Unified embedding utilities for Open Notebook.
+
+Provides centralized embedding generation with support for:
+- Single text embedding (with automatic chunking and mean pooling for large texts)
+- Batch text embedding (multiple texts with automatic batching)
+- Mean pooling for combining multiple embeddings into one
+
+All embedding operations in the application should use these functions
+to ensure consistent behavior and proper handling of large content.
+"""
+
+import asyncio
+import os
+from typing import List, Optional
+
+import numpy as np
+from loguru import logger
+
+from .chunking import CHUNK_SIZE, ContentType, chunk_text
+from .token_utils import token_count
+
+
+def _get_embedding_batch_size() -> int:
+    """
+    Read the embedding batch size from the environment.
+
+    This is intentionally configurable because provider limits vary widely, and
+    CPU-only local embedding endpoints often need smaller batches than cloud APIs.
+    """
+    raw = os.getenv("OPEN_NOTEBOOK_EMBEDDING_BATCH_SIZE", "50").strip()
+    try:
+        value = int(raw)
+        if value < 1:
+            raise ValueError
+        return value
+    except ValueError:
+        logger.warning(
+            "Invalid OPEN_NOTEBOOK_EMBEDDING_BATCH_SIZE='{}'; falling back to 50",
+            raw,
+        )
+        return 50
+
+
+EMBEDDING_BATCH_SIZE = _get_embedding_batch_size()
+EMBEDDING_MAX_RETRIES = 3
+EMBEDDING_RETRY_DELAY = 2  # seconds
+
+
+async def mean_pool_embeddings(embeddings: List[List[float]]) -> List[float]:
+    """
+    Combine multiple embeddings into a single embedding using mean pooling.
+
+    Algorithm:
+    1. Normalize each embedding to unit length
+    2. Compute element-wise mean
+    3. Normalize the result to unit length
+
+    This approach ensures the final embedding has the same properties as
+    individual embeddings (unit length) regardless of input count.
+
+    Args:
+        embeddings: List of embedding vectors (each is a list of floats)
+
+    Returns:
+        Single embedding vector (mean pooled and normalized)
+
+    Raises:
+        ValueError: If embeddings list is empty or embeddings have different dimensions
+    """
+    if not embeddings:
+        raise ValueError("Cannot mean pool empty list of embeddings")
+
+    if len(embeddings) == 1:
+        # Single embedding - just normalize and return
+        arr = np.array(embeddings[0], dtype=np.float64)
+        norm = np.linalg.norm(arr)
+        if norm > 0:
+            arr = arr / norm
+        return arr.tolist()
+
+    # Convert to numpy array
+    arr = np.array(embeddings, dtype=np.float64)
+
+    # Verify all embeddings have same dimension
+    if arr.ndim != 2:
+        raise ValueError(f"Expected 2D array, got shape {arr.shape}")
+
+    # Normalize each embedding to unit length
+    norms = np.linalg.norm(arr, axis=1, keepdims=True)
+    # Avoid division by zero
+    norms = np.where(norms > 0, norms, 1.0)
+    normalized = arr / norms
+
+    # Compute mean
+    mean = np.mean(normalized, axis=0)
+
+    # Normalize the result
+    mean_norm = np.linalg.norm(mean)
+    if mean_norm > 0:
+        mean = mean / mean_norm
+
+    return mean.tolist()
+
+
+async def generate_embeddings(
+    texts: List[str], command_id: Optional[str] = None
+) -> List[List[float]]:
+    """
+    Generate embeddings for multiple texts with automatic batching and retry.
+
+    Texts are split into batches of EMBEDDING_BATCH_SIZE to avoid exceeding
+    provider payload limits. Each batch is retried up to EMBEDDING_MAX_RETRIES
+    times on transient failures.
+
+    Args:
+        texts: List of text strings to embed
+        command_id: Optional command ID for error logging context
+
+    Returns:
+        List of embedding vectors, one per input text
+
+    Raises:
+        ValueError: If no embedding model is configured
+        RuntimeError: If embedding generation fails
+    """
+    if not texts:
+        return []
+
+    # Lazy import to avoid circular dependency
+    from open_notebook.ai.models import model_manager
+
+    embedding_model = await model_manager.get_embedding_model()
+    if not embedding_model:
+        raise ValueError(
+            "No embedding model configured. Please configure one in the Models section."
+        )
+
+    model_name = getattr(embedding_model, "model_name", "unknown")
+
+    # Log text sizes for debugging
+    metrics: tuple[int, int, int, int] | None = None
+
+    def _get_size_metrics() -> tuple[int, int, int, int]:
+        nonlocal metrics
+        if metrics is None:
+            token_sizes = [token_count(t) for t in texts]
+            metrics = (
+                min(token_sizes),
+                max(token_sizes),
+                sum(token_sizes),
+                sum(len(t) for t in texts),
+            )
+        return metrics
+
+    logger.opt(lazy=True).debug(
+        "Generating embeddings for {} texts "
+        "(tokens: min={}, max={}, total={}; chars: total={})",
+        lambda: len(texts),
+        lambda: _get_size_metrics()[0],
+        lambda: _get_size_metrics()[1],
+        lambda: _get_size_metrics()[2],
+        lambda: _get_size_metrics()[3],
+    )
+
+    all_embeddings: List[List[float]] = []
+    total_batches = (len(texts) + EMBEDDING_BATCH_SIZE - 1) // EMBEDDING_BATCH_SIZE
+
+    for batch_idx in range(total_batches):
+        start = batch_idx * EMBEDDING_BATCH_SIZE
+        end = start + EMBEDDING_BATCH_SIZE
+        batch = texts[start:end]
+
+        for attempt in range(1, EMBEDDING_MAX_RETRIES + 1):
+            try:
+                batch_embeddings = await embedding_model.aembed(batch)
+                all_embeddings.extend(batch_embeddings)
+                break
+            except Exception as e:
+                cmd_context = f" (command: {command_id})" if command_id else ""
+                if attempt < EMBEDDING_MAX_RETRIES:
+                    logger.debug(
+                        f"Embedding batch {batch_idx + 1}/{total_batches} "
+                        f"attempt {attempt}/{EMBEDDING_MAX_RETRIES} failed "
+                        f"using model '{model_name}'{cmd_context}: {e}. Retrying..."
+                    )
+                    await asyncio.sleep(EMBEDDING_RETRY_DELAY)
+                else:
+                    logger.debug(
+                        f"Embedding batch {batch_idx + 1}/{total_batches} "
+                        f"failed after {EMBEDDING_MAX_RETRIES} attempts "
+                        f"using model '{model_name}'{cmd_context}: {e}"
+                    )
+                    raise RuntimeError(
+                        f"Failed to generate embeddings using model '{model_name}' "
+                        f"(batch {batch_idx + 1}/{total_batches}, "
+                        f"{len(batch)} texts): {e}"
+                    ) from e
+
+    logger.debug(f"Generated {len(all_embeddings)} embeddings in {total_batches} batch(es)")
+    return all_embeddings
+
+
+async def generate_embedding(
+    text: str,
+    content_type: Optional[ContentType] = None,
+    file_path: Optional[str] = None,
+    command_id: Optional[str] = None,
+) -> List[float]:
+    """
+    Generate a single embedding for text, handling large content via chunking and mean pooling.
+
+    For short text (<= CHUNK_SIZE tokens):
+        - Embeds directly and returns the embedding
+
+    For long text (> CHUNK_SIZE tokens):
+        - Chunks the text using appropriate splitter for content type
+        - Embeds all chunks in batches
+        - Combines embeddings via mean pooling
+
+    Args:
+        text: The text to embed
+        content_type: Optional explicit content type for chunking
+        file_path: Optional file path for content type detection
+        command_id: Optional command ID for error logging context
+
+    Returns:
+        Single embedding vector (list of floats)
+
+    Raises:
+        ValueError: If text is empty or no embedding model configured
+        RuntimeError: If embedding generation fails
+    """
+    if not text or not text.strip():
+        raise ValueError("Cannot generate embedding for empty text")
+
+    text = text.strip()
+    text_tokens = token_count(text)
+
+    # Check if chunking is needed
+    if text_tokens <= CHUNK_SIZE:
+        # Short text - embed directly
+        logger.debug(f"Embedding short text ({text_tokens} tokens) directly")
+        embeddings = await generate_embeddings([text], command_id=command_id)
+        return embeddings[0]
+
+    # Long text - chunk and mean pool
+    logger.debug(f"Text exceeds chunk size ({text_tokens} tokens), chunking...")
+
+    chunks = chunk_text(text, content_type=content_type, file_path=file_path)
+
+    if not chunks:
+        raise ValueError("Text chunking produced no chunks")
+
+    if len(chunks) == 1:
+        # Single chunk after splitting
+        embeddings = await generate_embeddings(chunks, command_id=command_id)
+        return embeddings[0]
+
+    logger.debug(f"Embedding {len(chunks)} chunks and mean pooling")
+
+    # Embed all chunks in batches
+    embeddings = await generate_embeddings(chunks, command_id=command_id)
+
+    # Mean pool to get single embedding
+    pooled = await mean_pool_embeddings(embeddings)
+
+    logger.debug(f"Mean pooled {len(embeddings)} embeddings into single vector")
+    return pooled

open_notebook/utils/encryption.py

@@ -0,0 +1,198 @@
+"""
+Field-level encryption for sensitive data using API keys.
+
+This module provides encryption/decryption for API keys stored in the database.
+Fernet uses AES-128-CBC with HMAC-SHA256 for authenticated encryption.
+
+OPEN_NOTEBOOK_ENCRYPTION_KEY accepts **any string**. A Fernet key is derived
+from it via SHA-256, so users can set a simple passphrase like
+``OPEN_NOTEBOOK_ENCRYPTION_KEY=my-secret`` and it will work.
+
+Usage:
+    # Encrypt before storing
+    encrypted = encrypt_value(api_key)
+
+    # Decrypt when reading
+    decrypted = decrypt_value(encrypted)
+"""
+
+import base64
+import hashlib
+import os
+from pathlib import Path
+from typing import Optional
+
+from cryptography.fernet import Fernet, InvalidToken
+from loguru import logger
+
+
+def get_secret_from_env(var_name: str) -> Optional[str]:
+    """
+    Get a secret from environment, supporting Docker secrets pattern.
+
+    Checks for VAR_FILE first (Docker secrets), then falls back to VAR.
+
+    Args:
+        var_name: Base name of the environment variable (e.g., "OPEN_NOTEBOOK_ENCRYPTION_KEY")
+
+    Returns:
+        The secret value, or None if not configured.
+    """
+    # Check for _FILE variant first (Docker secrets)
+    file_path = os.environ.get(f"{var_name}_FILE")
+    if file_path:
+        try:
+            path = Path(file_path)
+            if path.exists() and path.is_file():
+                secret = path.read_text().strip()
+                if secret:
+                    logger.debug(f"Loaded {var_name} from file: {file_path}")
+                    return secret
+                else:
+                    logger.warning(f"{var_name}_FILE points to empty file: {file_path}")
+            else:
+                logger.warning(f"{var_name}_FILE path does not exist: {file_path}")
+        except Exception as e:
+            logger.error(f"Failed to read {var_name} from file {file_path}: {e}")
+
+    # Fall back to direct environment variable
+    return os.environ.get(var_name)
+
+
+def _get_or_create_encryption_key() -> str:
+    """
+    Get encryption key from environment, requires explicit configuration.
+
+    Priority:
+    1. OPEN_NOTEBOOK_ENCRYPTION_KEY_FILE (Docker secrets)
+    2. OPEN_NOTEBOOK_ENCRYPTION_KEY (environment variable)
+
+    For production deployments, you MUST set OPEN_NOTEBOOK_ENCRYPTION_KEY explicitly!
+
+    Returns:
+        Encryption key string.
+
+    Raises:
+        ValueError: If no encryption key is configured.
+    """
+    # First check environment/Docker secrets
+    key = get_secret_from_env("OPEN_NOTEBOOK_ENCRYPTION_KEY")
+    if key:
+        return key
+
+    raise ValueError(
+        "OPEN_NOTEBOOK_ENCRYPTION_KEY is not set. "
+        "Set this environment variable to any secret string to enable "
+        "encrypted storage of API keys in the database."
+    )
+
+
+# Lazy-loaded encryption key: initialized on first use, not at import time.
+# This prevents the entire app from crashing if the key is not yet configured
+# when other modules import from this file.
+_ENCRYPTION_KEY: Optional[str] = None
+
+
+def _get_encryption_key() -> str:
+    """Get the encryption key, initializing lazily on first call."""
+    global _ENCRYPTION_KEY
+    if _ENCRYPTION_KEY is None:
+        _ENCRYPTION_KEY = _get_or_create_encryption_key()
+    return _ENCRYPTION_KEY
+
+
+def _ensure_fernet_key(key: str) -> str:
+    """
+    Derive a valid Fernet key from an arbitrary string via SHA-256.
+
+    Any string is accepted as input. The key is derived by hashing it with
+    SHA-256 and encoding the result as URL-safe base64.
+    """
+    derived = hashlib.sha256(key.encode()).digest()
+    return base64.urlsafe_b64encode(derived).decode()
+
+
+def get_fernet() -> Fernet:
+    """
+    Get Fernet instance with the configured encryption key.
+
+    Returns:
+        Fernet instance.
+
+    Raises:
+        ValueError: If encryption key is not configured.
+    """
+    return Fernet(_ensure_fernet_key(_get_encryption_key()).encode())
+
+
+def encrypt_value(value: str) -> str:
+    """
+    Encrypt a string value using Fernet symmetric encryption.
+
+    Args:
+        value: The plain text string to encrypt.
+
+    Returns:
+        Base64-encoded encrypted string.
+
+    Raises:
+        ValueError: If encryption is not configured.
+    """
+    fernet = get_fernet()
+    return fernet.encrypt(value.encode()).decode()
+
+
+def looks_like_fernet_token(s: str) -> bool:
+    """
+    Check if string looks like a Fernet encrypted token.
+
+    Fernet tokens are versioned (1 byte) + timestamp (8 bytes) + IV (16 bytes)
+    + ciphertext (variable, multiple of 16 with PKCS7 padding) + HMAC (32 bytes).
+    Minimum decoded size is 73 bytes (1+8+16+16+32) for the smallest payload.
+    """
+    if len(s) < 100:  # Base64 of 73 bytes = ~100 chars minimum
+        return False
+    try:
+        decoded = base64.urlsafe_b64decode(s)
+        # Fernet: version(1) + timestamp(8) + IV(16) + ciphertext(>=16) + HMAC(32)
+        # Minimum 73 bytes, ciphertext must be multiple of 16 (AES block size)
+        if len(decoded) < 73:
+            return False
+        ciphertext_len = len(decoded) - 1 - 8 - 16 - 32
+        return ciphertext_len > 0 and ciphertext_len % 16 == 0
+    except Exception:
+        return False
+
+
+def decrypt_value(value: str) -> str:
+    """
+    Decrypt a Fernet-encrypted string value.
+
+    Handles graceful fallback for legacy unencrypted data.
+
+    Args:
+        value: The encrypted string (or plain text for legacy data).
+
+    Returns:
+        Decrypted plain text string, or original value if not encrypted.
+
+    Raises:
+        ValueError: If encryption is not configured or if decryption fails
+            for what appears to be encrypted data (wrong key).
+    """
+    fernet = get_fernet()
+
+    try:
+        return fernet.decrypt(value.encode()).decode()
+    except InvalidToken:
+        if looks_like_fernet_token(value):
+            # Looks like encrypted data but failed to decrypt - likely wrong key
+            raise ValueError(
+                "Decryption failed: data appears to be encrypted but key is incorrect. "
+                "Check OPEN_NOTEBOOK_ENCRYPTION_KEY configuration."
+            )
+        # Not a valid token - treat as legacy plaintext
+        return value
+    except Exception as e:
+        logger.error(f"Decryption failed: {e}")
+        raise ValueError(f"Decryption failed: {str(e)}")

open_notebook/utils/error_classifier.py

@@ -0,0 +1,103 @@
+"""
+Error classification utility for LLM provider errors.
+
+Maps raw exceptions from AI providers/Esperanto/LangChain to user-friendly
+error messages and appropriate exception types.
+"""
+
+from loguru import logger
+
+from open_notebook.exceptions import (
+    AuthenticationError,
+    ConfigurationError,
+    ExternalServiceError,
+    NetworkError,
+    OpenNotebookError,
+    RateLimitError,
+)
+
+# Classification rules: (keywords, exception_class, user_message or None to pass through)
+_CLASSIFICATION_RULES: list[tuple[list[str], type[OpenNotebookError], str | None]] = [
+    # Authentication errors
+    (
+        ["authentication", "unauthorized", "invalid api key", "invalid_api_key", "401"],
+        AuthenticationError,
+        "Authentication failed. Please check your API key in Settings -> Credentials.",
+    ),
+    # Rate limit errors
+    (
+        ["rate limit", "rate_limit", "429", "too many requests", "quota exceeded"],
+        RateLimitError,
+        "Rate limit exceeded. Please wait a moment and try again.",
+    ),
+    # Model not found (pass through original message)
+    (
+        ["model not found", "does not exist", "model_not_found"],
+        ConfigurationError,
+        None,
+    ),
+    # Configuration errors from provision.py (pass through)
+    (
+        ["no model configured", "please go to settings"],
+        ConfigurationError,
+        None,
+    ),
+    # Network errors
+    (
+        ["connecterror", "timeoutexception", "connection refused", "connection error", "timed out", "timeout"],
+        NetworkError,
+        "Could not connect to the AI provider. Please check your network connection and provider URL.",
+    ),
+    # Context length errors
+    (
+        ["context length", "token limit", "maximum context", "context_length_exceeded", "max_tokens"],
+        ExternalServiceError,
+        "Content too large for the selected model. Try using a smaller selection or a model with a larger context window.",
+    ),
+    # Payload too large errors
+    (
+        ["413", "payload too large", "request entity too large"],
+        ExternalServiceError,
+        "The request payload is too large for the AI provider. Try reducing the content size or using a different model.",
+    ),
+    # Provider availability errors
+    (
+        ["500", "502", "503", "service unavailable", "overloaded", "internal server error"],
+        ExternalServiceError,
+        "The AI provider is temporarily unavailable. Please try again in a few minutes.",
+    ),
+]
+
+
+def classify_error(exception: BaseException) -> tuple[type[OpenNotebookError], str]:
+    """
+    Classify a raw exception into a user-friendly error type and message.
+
+    Args:
+        exception: Any exception from LLM providers/Esperanto/LangChain
+
+    Returns:
+        Tuple of (exception_class, user_friendly_message)
+    """
+    error_str = str(exception).lower()
+    error_type_name = type(exception).__name__.lower()
+    combined = f"{error_type_name}: {error_str}"
+
+    for keywords, exc_class, message in _CLASSIFICATION_RULES:
+        for keyword in keywords:
+            if keyword in combined:
+                user_message = message if message is not None else _truncate(str(exception))
+                return exc_class, user_message
+
+    # Unclassified error - log for future improvement
+    logger.warning(
+        f"Unclassified LLM error ({type(exception).__name__}): {exception}"
+    )
+    return ExternalServiceError, f"AI service error: {_truncate(str(exception))}"
+
+
+def _truncate(text: str, max_length: int = 200) -> str:
+    """Truncate text to max_length to avoid leaking verbose internal details."""
+    if len(text) <= max_length:
+        return text
+    return text[:max_length] + "..."

open_notebook/utils/graph_utils.py

@@ -0,0 +1,23 @@
+import asyncio
+
+from langchain_core.runnables import RunnableConfig
+from loguru import logger
+
+
+async def get_session_message_count(graph, session_id: str) -> int:
+    """Get message count from LangGraph state, returns 0 on error."""
+    try:
+        # Use sync get_state() in a thread (SqliteSaver doesn't support async)
+        thread_state = await asyncio.to_thread(
+            graph.get_state,
+            config=RunnableConfig(configurable={"thread_id": session_id}),
+        )
+        if (
+            thread_state
+            and thread_state.values
+            and "messages" in thread_state.values
+        ):
+            return len(thread_state.values["messages"])
+    except Exception as e:
+        logger.warning(f"Could not fetch message count for session {session_id}: {e}")
+    return 0

open_notebook/utils/model_utils.py

@@ -0,0 +1,15 @@
+"""Utilities for working with Pydantic models."""
+
+from pydantic import BaseModel
+
+
+def full_model_dump(model):
+    """Recursively dump Pydantic models nested inside dicts/lists to plain data."""
+    if isinstance(model, BaseModel):
+        return model.model_dump()
+    elif isinstance(model, dict):
+        return {k: full_model_dump(v) for k, v in model.items()}
+    elif isinstance(model, list):
+        return [full_model_dump(item) for item in model]
+    else:
+        return model

open_notebook/utils/text_utils.py

@@ -0,0 +1,145 @@
+"""
+Text utilities for Open Notebook.
+Extracted from main utils to avoid circular imports.
+"""
+
+import re
+import unicodedata
+from typing import Tuple
+
+# Patterns for matching thinking content in AI responses
+# Standard pattern: <think>...</think>
+THINK_PATTERN = re.compile(r"<think>(.*?)</think>", re.DOTALL)
+# Pattern for malformed output: content</think> (missing opening tag)
+THINK_PATTERN_NO_OPEN = re.compile(r"^(.*?)</think>", re.DOTALL)
+
+
+def remove_non_ascii(text: str) -> str:
+    """Remove non-ASCII characters from text."""
+    return re.sub(r"[^\x00-\x7F]+", "", text)
+
+
+def remove_non_printable(text: str) -> str:
+    """Remove non-printable characters from text."""
+    # Replace any special Unicode whitespace characters with a regular space
+    text = re.sub(r"[\u2000-\u200B\u202F\u205F\u3000]", " ", text)
+
+    # Replace unusual line terminators with a single newline
+    text = re.sub(r"[\u2028\u2029\r]", "\n", text)
+
+    # Remove control characters, except newlines and tabs
+    text = "".join(
+        char for char in text if unicodedata.category(char)[0] != "C" or char in "\n\t"
+    )
+
+    # Replace non-breaking spaces with regular spaces
+    text = text.replace("\xa0", " ").strip()
+
+    # Keep letters (including accented ones), numbers, spaces, newlines, tabs, and basic punctuation
+    return re.sub(r"[^\w\s.,!?\-\n\t]", "", text, flags=re.UNICODE)
+
+
+def parse_thinking_content(content: str) -> Tuple[str, str]:
+    """
+    Parse message content to extract thinking content from <think> tags.
+
+    Handles both well-formed tags and malformed output where the opening
+    <think> tag is missing but </think> is present.
+
+    Args:
+        content (str): The original message content
+
+    Returns:
+        Tuple[str, str]: (thinking_content, cleaned_content)
+            - thinking_content: Content from within <think> tags
+            - cleaned_content: Original content with <think> blocks removed
+
+    Example:
+        >>> content = "<think>Let me analyze this</think>Here's my answer"
+        >>> thinking, cleaned = parse_thinking_content(content)
+        >>> print(thinking)
+        "Let me analyze this"
+        >>> print(cleaned)
+        "Here's my answer"
+    """
+    # Input validation
+    if not isinstance(content, str):
+        return "", str(content) if content is not None else ""
+
+    # Limit processing for very large content (100KB limit)
+    if len(content) > 100000:
+        return "", content
+
+    # Find all well-formed thinking blocks
+    thinking_matches = THINK_PATTERN.findall(content)
+
+    if thinking_matches:
+        # Join all thinking content with double newlines
+        thinking_content = "\n\n".join(match.strip() for match in thinking_matches)
+
+        # Remove all <think>...</think> blocks from the original content
+        cleaned_content = THINK_PATTERN.sub("", content)
+
+        # Clean up extra whitespace
+        cleaned_content = re.sub(r"\n\s*\n\s*\n", "\n\n", cleaned_content).strip()
+
+        return thinking_content, cleaned_content
+
+    # Handle malformed output: content</think> (missing opening tag)
+    # Some models like Nemotron output thinking without the opening <think> tag
+    malformed_match = THINK_PATTERN_NO_OPEN.match(content)
+    if malformed_match:
+        thinking_content = malformed_match.group(1).strip()
+        # Remove the thinking content and </think> tag
+        cleaned_content = content[malformed_match.end() :].strip()
+        return thinking_content, cleaned_content
+
+    return "", content
+
+
+def clean_thinking_content(content: str) -> str:
+    """
+    Remove thinking content from AI responses, returning only the cleaned content.
+
+    This is a convenience function for cases where you only need the cleaned
+    content and don't need access to the thinking process.
+
+    Args:
+        content (str): The original message content with potential <think> tags
+
+    Returns:
+        str: Content with <think> blocks removed and whitespace cleaned
+
+    Example:
+        >>> content = "<think>Let me think...</think>Here's the answer"
+        >>> clean_thinking_content(content)
+        "Here's the answer"
+    """
+    _, cleaned_content = parse_thinking_content(content)
+    return cleaned_content
+
+
+def extract_text_content(content) -> str:
+    """Extract text from LLM response content.
+
+    Handles both plain string responses and structured content formats
+    (e.g. Gemini's envelope format):
+    [{'type': 'text', 'text': '...', 'extras': {...}}]
+
+    Args:
+        content: The content from an AI message, either a string or a list of parts.
+
+    Returns:
+        The extracted text content as a string.
+    """
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        text_parts = []
+        for part in content:
+            if isinstance(part, dict) and "text" in part:
+                text_parts.append(part["text"])
+            elif isinstance(part, str):
+                text_parts.append(part)
+        return "".join(text_parts)
+    return str(content)

open_notebook/utils/token_utils.py

@@ -0,0 +1,57 @@
+"""
+Token utilities for Open Notebook.
+Handles token counting and cost calculations for language models.
+"""
+
+import os
+
+from open_notebook.config import TIKTOKEN_CACHE_DIR
+
+# Set tiktoken cache directory before importing tiktoken to ensure
+# tokenizer encodings are cached persistently in the data folder
+os.environ["TIKTOKEN_CACHE_DIR"] = TIKTOKEN_CACHE_DIR
+
+
+def token_count(input_string: str) -> int:
+    """
+    Count the number of tokens in the input string using the 'o200k_base' encoding.
+
+    Args:
+        input_string (str): The input string to count tokens for.
+
+    Returns:
+        int: The number of tokens in the input string.
+    """
+    try:
+        import tiktoken
+
+        encoding = tiktoken.get_encoding("o200k_base")
+        # disallowed_special=() treats sequences like "<|endoftext|>" as ordinary
+        # text instead of raising ValueError. User/source content can legitimately
+        # contain these substrings, and we only need a token count here.
+        tokens = encoding.encode(input_string, disallowed_special=())
+        return len(tokens)
+    except (ImportError, OSError) as e:
+        # Fallback: handles ImportError (tiktoken not installed) AND network/OS
+        # errors such as urllib.error.URLError or ConnectionError raised in
+        # offline environments when the encoding file cannot be downloaded.
+        from loguru import logger
+
+        logger.warning(
+            "tiktoken unavailable, falling back to word-count estimation: {}", e
+        )
+        return int(len(input_string.split()) * 1.3)
+
+
+def token_cost(token_count: int, cost_per_million: float = 0.150) -> float:
+    """
+    Calculate the cost of tokens based on the token count and cost per million tokens.
+
+    Args:
+        token_count (int): The number of tokens.
+        cost_per_million (float): The cost per million tokens. Default is 0.150.
+
+    Returns:
+        float: The calculated cost for the given token count.
+    """
+    return cost_per_million * (token_count / 1_000_000)

open_notebook/utils/url_validation.py

@@ -0,0 +1,136 @@
+"""
+URL validation shared by the credentials API and the AI provisioning layer.
+
+Lives in open_notebook.utils (not api/) so it can be imported both by the API
+layer (credential create/update, source-URL ingestion) and by open_notebook's
+own AI layer (connection_tester.py, ModelManager) to re-validate a
+provider-configured URL immediately before it is actually used - closing most
+of the DNS-rebinding TOCTOU window left by validating only once, at save time.
+"""
+
+import asyncio
+import ipaddress
+import socket
+from urllib.parse import urlparse
+
+# AWS's IMDSv6 metadata endpoint. Unlike the IPv4 metadata address
+# (169.254.169.254), this is a Unique Local Address, not link-local, so it is
+# not caught by ip.is_link_local and must be checked explicitly.
+_AWS_IMDS_V6_ADDRESS = ipaddress.ip_address("fd00:ec2::254")
+
+
+async def validate_url(url: str, provider: str) -> None:
+    """
+    Validate URL format for API endpoints.
+
+    This is a self-hosted application, so we allow:
+    - Private IPs (10.x, 172.16-31.x, 192.168.x) for self-hosted services
+    - Localhost for local services (Ollama, LM Studio, etc.)
+
+    We only block:
+    - Invalid schemes (must be http or https)
+    - Malformed URLs
+    - Link-local addresses (169.254.x.x) - used for cloud metadata endpoints
+    - AWS's IPv6 metadata address (fd00:ec2::254)
+    - Hostnames that resolve to any of the above
+
+    Args:
+        url: The URL to validate
+        provider: The provider name (for logging/context)
+
+    Raises:
+        ValueError: If the URL is invalid
+    """
+    if not url or not url.strip():
+        return  # Empty URLs handled elsewhere
+
+    try:
+        parsed = urlparse(url.strip())
+
+        # Validate scheme - only http/https allowed
+        if parsed.scheme not in ("http", "https"):
+            raise ValueError(
+                f"Invalid URL scheme: '{parsed.scheme}'. Only http and https are allowed."
+            )
+
+        # Extract hostname
+        hostname = parsed.hostname
+        if not hostname:
+            raise ValueError("Invalid URL: hostname could not be determined.")
+
+        # Try to parse as IP address to check for dangerous addresses
+        try:
+            ip = ipaddress.ip_address(hostname)
+            _reject_dangerous_ip(ip, hostname)
+
+        except ValueError as ve:
+            # Re-raise our own ValueErrors
+            if "Link-local" in str(ve) or "Invalid URL" in str(ve) or "metadata" in str(ve):
+                raise
+            # Not an IP address, it's a hostname - need to resolve and check
+            try:
+                # Resolve hostname to IP address. This is a blocking call -
+                # run it off the event loop so a slow/hanging DNS lookup
+                # doesn't stall every other concurrent request (this is
+                # called on the hot path of model provisioning, potentially
+                # once per chat message/transformation).
+                resolved_ips = await asyncio.to_thread(socket.getaddrinfo, hostname, None)
+                for family, _, _, _, sockaddr in resolved_ips:
+                    ip_addr = sockaddr[0]
+                    try:
+                        parsed_ip = ipaddress.ip_address(ip_addr)
+                        _reject_dangerous_ip(parsed_ip, hostname, resolved=True)
+                    except ValueError as inner_ve:
+                        if "link-local" in str(inner_ve).lower() or "metadata" in str(inner_ve).lower():
+                            raise
+                        # Skip non-IP addresses (e.g., IPv6 zones)
+                        continue
+            except socket.gaierror:
+                # Could not resolve hostname - allow it since the URL may be
+                # valid in the deployment environment (e.g., Azure endpoints,
+                # internal DNS names). We only block link-local addresses.
+                pass
+
+    except ValueError:
+        raise
+    except Exception:
+        raise ValueError("Invalid URL format. Check server logs for details.")
+
+
+def _reject_dangerous_ip(
+    ip: "ipaddress.IPv4Address | ipaddress.IPv6Address",
+    hostname: str,
+    resolved: bool = False,
+) -> None:
+    """Raise ValueError if `ip` is a link-local or cloud-metadata address."""
+    is_ipv4_mapped_link_local = (
+        hasattr(ip, "ipv4_mapped") and ip.ipv4_mapped and ip.ipv4_mapped.is_link_local
+    )
+
+    # Block link-local addresses (169.254.x.x / fe80::/10) - used for cloud
+    # metadata - including IPv4-mapped IPv6 addresses pointing to link-local
+    # (e.g. ::ffff:169.254.169.254 bypasses IPv6 is_link_local check).
+    if ip.is_link_local or is_ipv4_mapped_link_local:
+        if resolved:
+            raise ValueError(
+                f"Hostname '{hostname}' resolves to a link-local address (169.254.x.x) "
+                "which is not allowed for security reasons. These addresses are used "
+                "for cloud metadata endpoints."
+            )
+        raise ValueError(
+            "Link-local addresses (169.254.x.x) are not allowed for security reasons. "
+            "These addresses are used for cloud metadata endpoints."
+        )
+
+    # Block AWS's IMDSv6 metadata address - a Unique Local Address, not
+    # link-local, so it needs its own explicit check.
+    if ip == _AWS_IMDS_V6_ADDRESS:
+        if resolved:
+            raise ValueError(
+                f"Hostname '{hostname}' resolves to the AWS IMDSv6 metadata address "
+                "(fd00:ec2::254), which is not allowed for security reasons."
+            )
+        raise ValueError(
+            "The AWS IMDSv6 metadata address (fd00:ec2::254) is not allowed for "
+            "security reasons."
+        )

open_notebook/utils/version_utils.py

@@ -0,0 +1,153 @@
+"""
+Version utilities for Open Notebook.
+Handles version comparison, GitHub version fetching, and package version management.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+from urllib.parse import urlparse
+
+import requests  # type: ignore
+import tomli
+from packaging.version import parse as parse_version
+
+
+async def get_version_from_github_async(repo_url: str, branch: str = "main") -> str:
+    """
+    Fetch and parse the version from pyproject.toml in a public GitHub repository (async).
+    """
+    from urllib.parse import urlparse
+
+    import httpx
+    import tomli
+
+    # Parse the GitHub URL
+    parsed_url = urlparse(repo_url)
+    if "github.com" not in parsed_url.netloc:
+        raise ValueError("Not a GitHub URL")
+
+    # Extract owner and repo name from path
+    path_parts = parsed_url.path.strip("/").split("/")
+    if len(path_parts) < 2:
+        raise ValueError("Invalid GitHub repository URL")
+
+    owner, repo = path_parts[0], path_parts[1]
+
+    # Construct raw content URL for pyproject.toml
+    raw_url = f"https://raw.githubusercontent.com/{owner}/{repo}/{branch}/pyproject.toml"
+
+    # Fetch the file with timeout using httpx
+    async with httpx.AsyncClient(timeout=10.0) as client:
+        response = await client.get(raw_url)
+        response.raise_for_status()
+
+    # Parse TOML content
+    pyproject_data = tomli.loads(response.text)
+
+    # Try to find version
+    try:
+        # Check tool.poetry.version
+        version_str = pyproject_data["tool"]["poetry"]["version"]
+    except KeyError:
+        try:
+            # Check project.version
+            version_str = pyproject_data["project"]["version"]
+        except KeyError:
+            raise KeyError("Version not found in pyproject.toml")
+
+    return version_str
+
+def get_version_from_github(repo_url: str, branch: str = "main") -> str:
+    """
+    Fetch and parse the version from pyproject.toml in a public GitHub repository.
+
+    Args:
+        repo_url (str): URL of the GitHub repository
+        branch (str): Branch name to fetch from (defaults to "main")
+
+    Returns:
+        str: Version string from pyproject.toml
+
+    Raises:
+        ValueError: If the URL is not a valid GitHub repository URL
+        requests.RequestException: If there's an error fetching the file
+        KeyError: If version information is not found in pyproject.toml
+    """
+    # Parse the GitHub URL
+    parsed_url = urlparse(repo_url)
+    if "github.com" not in parsed_url.netloc:
+        raise ValueError("Not a GitHub URL")
+
+    # Extract owner and repo name from path
+    path_parts = parsed_url.path.strip("/").split("/")
+    if len(path_parts) < 2:
+        raise ValueError("Invalid GitHub repository URL")
+
+    owner, repo = path_parts[0], path_parts[1]
+
+    # Construct raw content URL for pyproject.toml
+    raw_url = (
+        f"https://raw.githubusercontent.com/{owner}/{repo}/{branch}/pyproject.toml"
+    )
+
+    # Fetch the file with timeout
+    response = requests.get(raw_url, timeout=10)
+    response.raise_for_status()
+
+    # Parse TOML content
+    pyproject_data = tomli.loads(response.text)
+
+    # Try to find version in different possible locations
+    try:
+        # Check project.version first (poetry style)
+        version = pyproject_data["tool"]["poetry"]["version"]
+    except KeyError:
+        try:
+            # Check project.version (standard style)
+            version = pyproject_data["project"]["version"]
+        except KeyError:
+            raise KeyError("Version not found in pyproject.toml")
+
+    return version
+
+
+def get_installed_version(package_name: str) -> str:
+    """
+    Get the version of an installed package.
+
+    Args:
+        package_name (str): Name of the installed package
+
+    Returns:
+        str: Version string of the installed package
+
+    Raises:
+        PackageNotFoundError: If the package is not installed
+    """
+    try:
+        return version(package_name)
+    except PackageNotFoundError:
+        raise PackageNotFoundError(f"Package '{package_name}' not found")
+
+
+def compare_versions(version1: str, version2: str) -> int:
+    """
+    Compare two semantic versions.
+
+    Args:
+        version1 (str): First version string
+        version2 (str): Second version string
+
+    Returns:
+        int: -1 if version1 < version2
+              0 if version1 == version2
+              1 if version1 > version2
+    """
+    v1 = parse_version(version1)
+    v2 = parse_version(version2)
+
+    if v1 < v2:
+        return -1
+    elif v1 > v2:
+        return 1
+    else:
+        return 0