From 11ca9c138c6c1c5ac2ad975141a0fbe35a41ba0e Mon Sep 17 00:00:00 2001 From: Luis Novo Date: Sat, 11 Jul 2026 19:55:46 -0300 Subject: [PATCH] refactor(api): single context-building implementation (#1079) Consolidate the three copies of context assembly into open_notebook/utils/context_builder.py: - POST /api/chat/context now delegates to build_notebook_context() (same request/response shapes, same string-matching config semantics) - The source-chat graph now calls build_source_context() instead of the 495-line generalized ContextBuilder class, which had exactly one caller and whose notebook/notes/priority-config flexibility was dead - POST /api/notebooks/{notebook_id}/context removed: it duplicated /api/chat/context with a slightly different envelope and had zero callers (frontend, docs, tests) Behavior is pinned by new characterization tests written before the refactor (tests/test_context_endpoint_characterization.py) plus unit tests for build_source_context. --- CHANGELOG.md | 2 + api/main.py | 2 - api/models.py | 24 - api/routers/chat.py | 107 +-- api/routers/context.py | 132 ---- docs/7-DEVELOPMENT/api-reference.md | 2 +- docs/7-DEVELOPMENT/content-processing.md | 7 +- open_notebook/graphs/source_chat.py | 19 +- open_notebook/utils/__init__.py | 2 +- open_notebook/utils/context_builder.py | 641 +++++------------- .../test_context_endpoint_characterization.py | 229 +++++++ tests/test_typed_exceptions_reach_handlers.py | 1 - tests/test_utils.py | 110 ++- 13 files changed, 514 insertions(+), 764 deletions(-) delete mode 100644 api/routers/context.py create mode 100644 tests/test_context_endpoint_characterization.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c4087ae6..91d73455 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed +- Context building consolidated into a single implementation (`open_notebook/utils/context_builder.py`): the copy-pasted source/note assembly loops behind `POST /api/chat/context` and the removed notebook-context endpoint, plus the 495-line generalized `ContextBuilder` class (whose only caller was the source-chat graph), are now two focused functions — `build_notebook_context()` (backs `POST /api/chat/context`, unchanged request/response shapes and config semantics) and `build_source_context()` (backs the source-chat graph, same context shape and 50k-token budget). Pinned by new characterization tests — no behavior change for the surviving paths +- **Removed** `POST /api/notebooks/{notebook_id}/context`: it duplicated `POST /api/chat/context` (same assembly logic, slightly different response envelope) and had zero callers — frontend, docs and tests only use `/api/chat/context`. If you called it programmatically, switch to `POST /api/chat/context` (body: `{notebook_id, context_config}`; response fields: `context.sources`/`context.notes`, `token_count`, `char_count`) - Backend provider metadata now lives in a single registry (`open_notebook/ai/provider_registry.py`): env var config, modalities, connection-test models, OpenAI-compatible discovery URLs and docs links are defined once per provider, and `PROVIDER_ENV_CONFIG`, `PROVIDER_MODALITIES`, `TEST_MODELS` and `OPENAI_COMPAT_PROVIDERS` are derived from it. Adding a provider drops from ~6 hand-synced dicts to the registry plus two manual copies (the `SupportedProvider` Literal and the frontend provider table), both enforced by tests - Frontend convention cleanup (no user-facing change): hook files unified to kebab-case (`useNotebookChat.ts`/`useSourceChat.ts` → `use-notebook-chat.ts`/`use-source-chat.ts`), `src/components/source/` merged into `src/components/sources/`, the localStorage auth-token parsing ritual extracted into a single `getAuthToken()` helper (`src/lib/auth-token.ts`), and non-streaming raw `fetch` calls routed through `apiClient` (podcast audio download, auth-status check). SSE/streaming paths and the login/checkAuth credential probes deliberately keep raw `fetch` - Pruned unused langchain packages: removed `langchain-community` and `langchain-deepseek` from the dependencies (nothing imports them — DeepSeek and xAI route through esperanto's OpenAI-compatible path, which uses `langchain-openai`). The remaining `langchain-*` provider packages are documented as runtime requirements of esperanto's dynamic `to_langchain()` and the whole langchain/langgraph family now carries explicit upper bounds; `langchain-core` and `langchain-text-splitters` (both directly imported but previously only transitive) are now declared explicitly diff --git a/api/main.py b/api/main.py index 24b206a6..2e7aa9e8 100644 --- a/api/main.py +++ b/api/main.py @@ -19,7 +19,6 @@ from api.routers import ( auth, chat, config, - context, credentials, embedding, embedding_rebuild, @@ -382,7 +381,6 @@ app.include_router( embedding_rebuild.router, prefix="/api/embeddings", tags=["embeddings"] ) app.include_router(settings.router, prefix="/api", tags=["settings"]) -app.include_router(context.router, prefix="/api", tags=["context"]) app.include_router(sources.router, prefix="/api", tags=["sources"]) app.include_router(insights.router, prefix="/api", tags=["insights"]) app.include_router(commands_router.router, prefix="/api", tags=["commands"]) diff --git a/api/models.py b/api/models.py index 5c7b7302..3e046cf5 100644 --- a/api/models.py +++ b/api/models.py @@ -386,30 +386,6 @@ class SourceListResponse(BaseModel): processing_info: Optional[Dict[str, Any]] = None -# Context API models -class ContextConfig(BaseModel): - sources: Dict[str, str] = Field( - default_factory=dict, description="Source inclusion config {source_id: level}" - ) - notes: Dict[str, str] = Field( - default_factory=dict, description="Note inclusion config {note_id: level}" - ) - - -class ContextRequest(BaseModel): - notebook_id: str = Field(..., description="Notebook ID to get context for") - context_config: Optional[ContextConfig] = Field( - None, description="Context configuration" - ) - - -class ContextResponse(BaseModel): - notebook_id: str - sources: List[Dict[str, Any]] = Field(..., description="Source context data") - notes: List[Dict[str, Any]] = Field(..., description="Note context data") - total_tokens: Optional[int] = Field(None, description="Estimated token count") - - # Insights API models class SourceInsightResponse(BaseModel): id: str diff --git a/api/routers/chat.py b/api/routers/chat.py index 19e1f05e..5cdd732a 100644 --- a/api/routers/chat.py +++ b/api/routers/chat.py @@ -12,21 +12,16 @@ from api.routers._chat_shared import ( SuccessResponse, extract_chat_messages, get_session_or_404, - normalize_record_id, ) from open_notebook.database.repository import ensure_record_id, repo_query -from open_notebook.domain.notebook import ( - ChatSession, - Note, - Notebook, - Source, - SourceInsight, -) +from open_notebook.domain.notebook import ChatSession, Notebook from open_notebook.exceptions import ( NotFoundError, OpenNotebookError, ) from open_notebook.graphs.chat import graph as chat_graph +from open_notebook.utils import token_count +from open_notebook.utils.context_builder import build_notebook_context from open_notebook.utils.graph_utils import get_session_message_count router = APIRouter() @@ -402,100 +397,12 @@ async def build_context(request: BuildContextRequest): if not notebook: raise HTTPException(status_code=404, detail="Notebook not found") - context_data: dict[str, list[dict[str, str]]] = {"sources": [], "notes": []} - total_content = "" + context_data, total_content = await build_notebook_context( + notebook, request.context_config + ) - # Process context configuration if provided - if request.context_config: - # Process sources - for source_id, status in request.context_config.get("sources", {}).items(): - if "not in" in status: - continue - - try: - # Add table prefix if not present - full_source_id = normalize_record_id("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 - - # Process notes - for note_id, status in request.context_config.get("notes", {}).items(): - if "not in" in status: - continue - - try: - # Add table prefix if not present - full_note_id = normalize_record_id("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 - - # Calculate character and token counts char_count = len(total_content) - # Use token count utility if available - try: - from open_notebook.utils import token_count - - estimated_tokens = token_count(total_content) if total_content else 0 - except ImportError: - # Fallback to simple estimation - estimated_tokens = char_count // 4 + estimated_tokens = token_count(total_content) if total_content else 0 return BuildContextResponse( context=context_data, token_count=estimated_tokens, char_count=char_count diff --git a/api/routers/context.py b/api/routers/context.py deleted file mode 100644 index d5625551..00000000 --- a/api/routers/context.py +++ /dev/null @@ -1,132 +0,0 @@ -from fastapi import APIRouter, HTTPException -from loguru import logger - -from api.models import ContextRequest, ContextResponse -from open_notebook.domain.notebook import Note, Notebook, Source, SourceInsight -from open_notebook.exceptions import ( - InvalidInputError, - OpenNotebookError, -) -from open_notebook.utils import token_count - -router = APIRouter() - - -@router.post("/notebooks/{notebook_id}/context", response_model=ContextResponse) -async def get_notebook_context(notebook_id: str, context_request: ContextRequest): - """Get context for a notebook based on configuration.""" - try: - # Verify notebook exists - notebook = await Notebook.get(notebook_id) - if not notebook: - raise HTTPException(status_code=404, detail="Notebook not found") - - context_data: dict[str, list[dict[str, str]]] = {"note": [], "source": []} - total_content = "" - - # Process context configuration if provided - if context_request.context_config: - # Process sources - for source_id, status in context_request.context_config.sources.items(): - if "not in" in status: - continue - - try: - # Add table prefix if not present - full_source_id = ( - source_id - if source_id.startswith("source:") - else f"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["source"].append(source_context) - total_content += str(source_context) - elif "full content" in status: - source_context = await source.get_context(context_size="long") - context_data["source"].append(source_context) - total_content += str(source_context) - except Exception as e: - logger.warning(f"Error processing source {source_id}: {str(e)}") - continue - - # Process notes - for note_id, status in context_request.context_config.notes.items(): - if "not in" in status: - continue - - try: - # Add table prefix if not present - full_note_id = ( - note_id if note_id.startswith("note:") else f"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["note"].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["source"].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["note"].append(note_context) - total_content += str(note_context) - except Exception as e: - logger.warning(f"Error processing note {note.id}: {str(e)}") - continue - - # Calculate estimated token count - estimated_tokens = token_count(total_content) if total_content else 0 - - return ContextResponse( - notebook_id=notebook_id, - sources=context_data["source"], - notes=context_data["note"], - total_tokens=estimated_tokens, - ) - - except HTTPException: - raise - except InvalidInputError as e: - raise HTTPException(status_code=400, detail=str(e)) - except OpenNotebookError: - raise - except Exception as e: - logger.error(f"Error getting context for notebook {notebook_id}: {str(e)}") - raise HTTPException(status_code=500, detail=f"Error getting context: {str(e)}") diff --git a/docs/7-DEVELOPMENT/api-reference.md b/docs/7-DEVELOPMENT/api-reference.md index 451f6742..8628cdfe 100644 --- a/docs/7-DEVELOPMENT/api-reference.md +++ b/docs/7-DEVELOPMENT/api-reference.md @@ -60,7 +60,7 @@ Instead of memorizing endpoints, use the interactive API docs: **Chat** - Conversational AI interface - `GET/POST /chat/sessions` - Manage chat sessions - `POST /chat/execute` - Send message and get response -- `POST /chat/context/build` - Prepare context for chat +- `POST /chat/context` - Prepare context for chat **Search** - Find content by text or semantic similarity - `POST /search` - Full-text or vector search diff --git a/docs/7-DEVELOPMENT/content-processing.md b/docs/7-DEVELOPMENT/content-processing.md index 47ace643..82f7c9f4 100644 --- a/docs/7-DEVELOPMENT/content-processing.md +++ b/docs/7-DEVELOPMENT/content-processing.md @@ -35,10 +35,11 @@ All embedding is fire-and-forget through the surreal-commands worker — nothing ## Context building (`utils/context_builder.py`) -`ContextBuilder` assembles LLM context from selected sources/notes/insights under a token budget: +The single implementation behind both context consumers: -- Each `ContextItem` counts its own tokens on construction; `build()` adds items in priority order (default weights: source 100 > insight 75 > note 50, see `ContextConfig`) and stops once `max_tokens` is exceeded (no prorating). -- Fetching is lazy (nothing is loaded until `build()`), and every call re-fetches — there is no cache layer. +- `build_notebook_context()` backs `POST /api/chat/context` (chat panel + podcast generation): it assembles source/note contexts from the inclusion config, whose status strings are matched textually ("not in" skips, "insights" → short context, "full content" → long context). Without a config, every source and note is included with its short context. Per-item failures are logged and skipped. +- `build_source_context()` backs the source-chat graph: one source's short context plus its insights, truncated to a token budget by dropping insights (last-fetched first). +- Every call re-fetches — there is no cache layer. - Token counting uses `o200k_base` via tiktoken and is an estimate (±5-10% vs. the actual model); `token_count()` falls back to a coarse estimate if tiktoken is unavailable. ## Encryption (`utils/encryption.py`) {#encryption} diff --git a/open_notebook/graphs/source_chat.py b/open_notebook/graphs/source_chat.py index fb9047c8..4e2dfa6f 100644 --- a/open_notebook/graphs/source_chat.py +++ b/open_notebook/graphs/source_chat.py @@ -15,7 +15,7 @@ from open_notebook.config import LANGGRAPH_CHECKPOINT_FILE from open_notebook.domain.notebook import Source, SourceInsight from open_notebook.exceptions import OpenNotebookError from open_notebook.utils import clean_thinking_content -from open_notebook.utils.context_builder import ContextBuilder +from open_notebook.utils.context_builder import build_source_context from open_notebook.utils.error_classifier import classify_error from open_notebook.utils.text_utils import extract_text_content @@ -37,7 +37,7 @@ def call_model_with_source_context( Main function that builds source context and calls the model. This function: - 1. Uses ContextBuilder to build source-specific context + 1. Uses build_source_context to build source-specific context 2. Applies the source_chat Jinja2 prompt template 3. Handles model provisioning with override support 4. Tracks context indicators for referenced insights/content @@ -58,19 +58,18 @@ def _call_model_with_source_context_inner( if not source_id: raise ValueError("source_id is required in state") - # Build source context using ContextBuilder (run async code in new loop) + # Build source context using build_source_context (run async code in new loop) def build_context(): """Build context in a new event loop""" new_loop = asyncio.new_event_loop() try: asyncio.set_event_loop(new_loop) - context_builder = ContextBuilder( - source_id=source_id, - include_insights=True, - include_notes=False, # Focus on source-specific content - max_tokens=50000, # Reasonable limit for source context + return new_loop.run_until_complete( + build_source_context( + source_id=source_id, + max_tokens=50000, # Reasonable limit for source context + ) ) - return new_loop.run_until_complete(context_builder.build()) finally: new_loop.close() asyncio.set_event_loop(None) @@ -192,7 +191,7 @@ def _format_source_context(context_data: Dict) -> str: Format the context data into a readable string for the prompt. Args: - context_data: Context data from ContextBuilder + context_data: Context data from build_source_context Returns: Formatted context string diff --git a/open_notebook/utils/__init__.py b/open_notebook/utils/__init__.py index 17e2b10b..72dd0945 100644 --- a/open_notebook/utils/__init__.py +++ b/open_notebook/utils/__init__.py @@ -2,7 +2,7 @@ Utils package for Open Notebook. To avoid circular imports, import functions directly: -- from open_notebook.utils.context_builder import ContextBuilder +- 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 diff --git a/open_notebook/utils/context_builder.py b/open_notebook/utils/context_builder.py index ed03712f..7d3717ca 100644 --- a/open_notebook/utils/context_builder.py +++ b/open_notebook/utils/context_builder.py @@ -1,495 +1,208 @@ -""" -Generic ContextBuilder for the Open Notebook project. +"""Context building for chat and podcast generation. -This module provides a flexible ContextBuilder class that can handle any parameters -and build context from sources, notebooks, insights, and notes. +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 dataclasses import dataclass -from typing import Any, Dict, List, Literal, Optional +from typing import Any, Dict, Optional, Tuple from loguru import logger -from open_notebook.domain.notebook import Note, Notebook, Source +from open_notebook.domain.notebook import ( + Note, + Notebook, + Source, + SourceInsight, +) from open_notebook.exceptions import DatabaseOperationError, NotFoundError from .token_utils import token_count -@dataclass -class ContextItem: - """Represents a single item in the context.""" - - id: str - type: Literal["source", "note", "insight"] - content: Dict[str, Any] - priority: int = 0 - token_count: Optional[int] = None - - def __post_init__(self): - """Calculate token count for the content if not provided.""" - if self.token_count is None: - content_str = str(self.content) - self.token_count = token_count(content_str) - - -@dataclass -class ContextConfig: - """Configuration for context building.""" - - sources: Optional[Dict[str, str]] = None # {source_id: inclusion_level} - notes: Optional[Dict[str, str]] = None # {note_id: inclusion_level} - include_insights: bool = True - include_notes: bool = True - max_tokens: Optional[int] = None - priority_weights: Optional[Dict[str, int]] = None # {type: weight} - - def __post_init__(self): - """Initialize default values.""" - if self.sources is None: - self.sources = {} - if self.notes is None: - self.notes = {} - if self.priority_weights is None: - self.priority_weights = {"source": 100, "note": 50, "insight": 75} - - -class ContextBuilder: - """ - Generic ContextBuilder that can handle any parameters and build context - from sources, notebooks, insights, and notes. - """ - - def __init__(self, **kwargs): - """ - Initialize ContextBuilder with flexible parameters. - - Supported parameters: - - source_id: str - Include specific source - - notebook_id: str - Include notebook content - - include_insights: bool - Include source insights - - include_notes: bool - Include notes - - context_config: ContextConfig - Custom context configuration - - max_tokens: int - Maximum token limit - - priority_order: List[str] - Custom priority order - """ - # Store all parameters for flexibility - self.params = kwargs - - # Extract commonly used parameters - self.source_id: Optional[str] = kwargs.get("source_id") - self.notebook_id: Optional[str] = kwargs.get("notebook_id") - self.include_insights: bool = kwargs.get("include_insights", True) - self.include_notes: bool = kwargs.get("include_notes", True) - self.max_tokens: Optional[int] = kwargs.get("max_tokens") - - # Context configuration - context_config_arg: Optional[ContextConfig] = kwargs.get("context_config") - self.context_config: ContextConfig - if context_config_arg is None: - self.context_config = ContextConfig( - include_insights=self.include_insights, - include_notes=self.include_notes, - max_tokens=self.max_tokens, - ) - else: - self.context_config = context_config_arg - - # Items storage - self.items: List[ContextItem] = [] - - logger.debug(f"ContextBuilder initialized with params: {list(kwargs.keys())}") - - async def build(self) -> Dict[str, Any]: - """ - Build context based on provided parameters. - - Returns: - Dict containing the built context with metadata - """ - try: - logger.info("Starting context building") - - # Clear existing items - self.items = [] - - # Build context based on parameters - if self.source_id: - await self._add_source_context(self.source_id) - - if self.notebook_id: - await self._add_notebook_context(self.notebook_id) - - # Process any additional custom parameters - await self._process_custom_params() - - # Apply post-processing - self.remove_duplicates() - self.prioritize() - - if self.max_tokens: - self.truncate_to_fit(self.max_tokens) - - # Format and return response - return self._format_response() - - except Exception as e: - logger.error(f"Error building context: {str(e)}") - raise DatabaseOperationError(f"Failed to build context: {str(e)}") - - async def _add_source_context( - self, source_id: str, inclusion_level: str = "insights" - ) -> None: - """ - Add source and its insights to context. - - Args: - source_id: ID of the source - inclusion_level: "insights", "full content", or "not in" - """ - if inclusion_level == "not in": - return - - try: - # Ensure source ID has table prefix - full_source_id = ( - source_id if source_id.startswith("source:") else f"source:{source_id}" - ) - - source = await Source.get(full_source_id) - if not source: - logger.warning(f"Source {source_id} not found") - return - - # Determine context size based on inclusion level - context_size: Literal["short", "long"] = ( - "long" if "full content" in inclusion_level else "short" - ) - source_context = await source.get_context(context_size=context_size) - - # Add source item - priority = (self.context_config.priority_weights or {}).get("source", 100) - item = ContextItem( - id=source.id or "", - type="source", - content=source_context, - priority=priority, - ) - self.add_item(item) - - # Add insights if requested and available - if self.include_insights and "insights" in inclusion_level: - insights = await source.get_insights() - for insight in insights: - insight_priority = (self.context_config.priority_weights or {}).get( - "insight", 75 - ) - insight_item = ContextItem( - id=insight.id or "", - type="insight", - content={ - "id": insight.id, - "source_id": source.id, - "insight_type": insight.insight_type, - "content": insight.content, - }, - priority=insight_priority, - ) - self.add_item(insight_item) - - logger.debug(f"Added source context for {source_id}") - - except NotFoundError: - logger.warning(f"Source {source_id} not found") - except Exception as e: - logger.error(f"Error adding source context for {source_id}: {str(e)}") - raise - - async def _add_notebook_context(self, notebook_id: str) -> None: - """ - Add notebook content based on context configuration. - - Args: - notebook_id: ID of the notebook - """ - try: - notebook = await Notebook.get(notebook_id) - if not notebook: - raise NotFoundError(f"Notebook {notebook_id} not found") - - # Process sources from context config or get all - config_sources = self.context_config.sources - if config_sources: - for source_id, status in config_sources.items(): - await self._add_source_context(source_id, status) - else: - # Default: get all sources with insights - sources = await notebook.get_sources() - for source in sources: - if source.id: - await self._add_source_context(source.id, "insights") - - # Process notes from context config or get all - if self.include_notes: - config_notes = self.context_config.notes - if config_notes: - for note_id, status in config_notes.items(): - if "not in" not in status: - await self._add_note_context(note_id, status) - else: - # Default: get all notes with short content - notes = await notebook.get_notes() - for note in notes: - if note.id: - await self._add_note_context(note.id, "full content") - - logger.debug(f"Added notebook context for {notebook_id}") - - except Exception as e: - logger.error(f"Error adding notebook context for {notebook_id}: {str(e)}") - raise - - async def _add_note_context( - self, note_id: str, inclusion_level: str = "full content" - ) -> None: - """ - Add note to context. - - Args: - note_id: ID of the note - inclusion_level: "full content" or "not in" - """ - if inclusion_level == "not in": - return - - try: - # Ensure note ID has table prefix - full_note_id = note_id if note_id.startswith("note:") else f"note:{note_id}" - - note = await Note.get(full_note_id) - if not note: - logger.warning(f"Note {note_id} not found") - return - - # Get note context - context_size: Literal["short", "long"] = ( - "long" if "full content" in inclusion_level else "short" - ) - note_context = note.get_context(context_size=context_size) - - # Add note item - priority = (self.context_config.priority_weights or {}).get("note", 50) - item = ContextItem( - id=note.id or "", type="note", content=note_context, priority=priority - ) - self.add_item(item) - - logger.debug(f"Added note context for {note_id}") - - except NotFoundError: - logger.warning(f"Note {note_id} not found") - except Exception as e: - logger.error(f"Error adding note context for {note_id}: {str(e)}") - - async def _process_custom_params(self) -> None: - """Process any additional custom parameters.""" - # Hook for future extensions - can be overridden in subclasses - # or used to process additional kwargs - for key, value in self.params.items(): - if key.startswith("custom_"): - logger.debug(f"Processing custom parameter: {key}={value}") - # Custom processing logic can be added here - - def add_item(self, item: ContextItem) -> None: - """ - Add a ContextItem to the builder. - - Args: - item: ContextItem to add - """ - self.items.append(item) - logger.debug(f"Added item {item.id} with priority {item.priority}") - - def prioritize(self) -> None: - """Sort items by priority (higher priority first).""" - self.items.sort(key=lambda x: x.priority, reverse=True) - logger.debug(f"Prioritized {len(self.items)} items") - - def truncate_to_fit(self, max_tokens: int) -> None: - """ - Remove items if total token count exceeds limit. - - Args: - max_tokens: Maximum allowed tokens - """ - if not max_tokens: - return - - total_tokens = sum(item.token_count or 0 for item in self.items) - - if total_tokens <= max_tokens: - logger.debug(f"Token count {total_tokens} within limit {max_tokens}") - return - - logger.info(f"Truncating from {total_tokens} to {max_tokens} tokens") - - # Remove items from the end (lowest priority) until under limit - current_tokens = total_tokens - removed_count = 0 - - while current_tokens > max_tokens and self.items: - removed_item = self.items.pop() - current_tokens -= removed_item.token_count or 0 - removed_count += 1 - - logger.info( - f"Removed {removed_count} items, final token count: {current_tokens}" - ) - - def remove_duplicates(self) -> None: - """Remove duplicate items based on ID.""" - seen_ids = set() - deduplicated_items = [] - - for item in self.items: - if item.id not in seen_ids: - deduplicated_items.append(item) - seen_ids.add(item.id) - - removed_count = len(self.items) - len(deduplicated_items) - self.items = deduplicated_items - - if removed_count > 0: - logger.debug(f"Removed {removed_count} duplicate items") - - def _format_response(self) -> Dict[str, Any]: - """ - Format the final response. - - Returns: - Formatted context response - """ - # Group items by type - sources = [] - notes = [] - insights = [] - - for item in self.items: - if item.type == "source": - sources.append(item.content) - elif item.type == "note": - notes.append(item.content) - elif item.type == "insight": - insights.append(item.content) - - # Calculate total tokens - total_tokens = sum(item.token_count or 0 for item in self.items) - - response = { - "sources": sources, - "notes": notes, - "insights": insights, - "total_tokens": total_tokens, - "total_items": len(self.items), - "metadata": { - "source_count": len(sources), - "note_count": len(notes), - "insight_count": len(insights), - "config": { - "include_insights": self.include_insights, - "include_notes": self.include_notes, - "max_tokens": self.max_tokens, - }, - }, - } - - # Add notebook_id if provided - if self.notebook_id: - response["notebook_id"] = self.notebook_id - - logger.info( - f"Built context with {len(self.items)} items, {total_tokens} tokens" - ) - - return response - - -# Convenience functions for common use cases +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_id: str, - context_config: Optional[ContextConfig] = None, - max_tokens: Optional[int] = None, -) -> Dict[str, Any]: - """ - Build context for a notebook. + notebook: Notebook, + context_config: Optional[Dict[str, Any]], +) -> Tuple[Dict[str, list], str]: + """Assemble source/note context for a notebook. - Args: - notebook_id: ID of the notebook - context_config: Optional context configuration - max_tokens: Optional token limit + 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: - Built context + ({"sources": [...], "notes": [...]}, concatenated str() of every + included context dict — used for token/char counting). """ - builder = ContextBuilder( - notebook_id=notebook_id, context_config=context_config, max_tokens=max_tokens - ) - return await builder.build() + 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, include_insights: bool = True, max_tokens: Optional[int] = None + 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". """ - Build context for a single source. + try: + sources: list = [] + insights: list = [] + item_tokens: list[int] = [] - Args: - source_id: ID of the source - include_insights: Whether to include insights - max_tokens: Optional token limit + try: + full_source_id = _ensure_prefix("source", source_id) + source = await Source.get(full_source_id) + except NotFoundError: + source = None - Returns: - Built context - """ - builder = ContextBuilder( - source_id=source_id, include_insights=include_insights, max_tokens=max_tokens - ) - return await builder.build() + 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") -async def build_mixed_context( - source_ids: Optional[List[str]] = None, - note_ids: Optional[List[str]] = None, - notebook_id: Optional[str] = None, - max_tokens: Optional[int] = None, -) -> Dict[str, Any]: - """ - Build context from mixed sources. + # 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() - Args: - source_ids: List of source IDs - note_ids: List of note IDs - notebook_id: Optional notebook ID - max_tokens: Optional token limit + total_items = len(sources) + len(insights) + logger.info(f"Built context with {total_items} items, {total_tokens} tokens") - Returns: - Built context - """ - context_config = ContextConfig(max_tokens=max_tokens) - - # Configure sources - if source_ids: - context_config.sources = {sid: "insights" for sid in source_ids} - - # Configure notes - if note_ids: - context_config.notes = {nid: "full content" for nid in note_ids} - - builder = ContextBuilder( - notebook_id=notebook_id, context_config=context_config, max_tokens=max_tokens - ) - return await builder.build() + 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)}") diff --git a/tests/test_context_endpoint_characterization.py b/tests/test_context_endpoint_characterization.py new file mode 100644 index 00000000..b5957220 --- /dev/null +++ b/tests/test_context_endpoint_characterization.py @@ -0,0 +1,229 @@ +"""Characterization tests for POST /api/chat/context. + +These pin down the exact response shape and the string-matching config +semantics ("not in" skips, "insights" -> short context, "full content" -> +long context) before the context-building loop is extracted out of the +router into open_notebook/utils/context_builder.py. They must pass +unchanged before and after the refactor. + +DB access is mocked following the style of +tests/test_chat_routers_characterization.py. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + + +@pytest.fixture +def client(): + from api.main import app + + return TestClient(app) + + +def _notebook(**overrides): + defaults = dict(id="notebook:1", name="My Notebook") + defaults.update(overrides) + nb = SimpleNamespace(**defaults) + nb.get_sources = AsyncMock(return_value=[]) + nb.get_notes = AsyncMock(return_value=[]) + return nb + + +def _source(source_id="source:s1", context=None): + src = SimpleNamespace(id=source_id) + src.get_context = AsyncMock( + return_value=context if context is not None else {"id": source_id, "title": "T"} + ) + return src + + +def _note(note_id="note:n1", context=None): + note = SimpleNamespace(id=note_id) + # Note.get_context is synchronous in the router code. + note.get_context = MagicMock( + return_value=context if context is not None else {"id": note_id, "title": "N"} + ) + return note + + +# Patch the domain classes at their definition site so the tests are +# independent of which module hosts the context-building loop. +PATCH_NOTEBOOK = "open_notebook.domain.notebook.Notebook.get" +PATCH_SOURCE = "open_notebook.domain.notebook.Source.get" +PATCH_NOTE = "open_notebook.domain.notebook.Note.get" +PATCH_INSIGHTS = "open_notebook.domain.notebook.SourceInsight.get_for_sources" + + +@pytest.mark.asyncio +async def test_config_string_matching_semantics(client): + """'not in' skips, 'insights' -> short, 'full content' -> long.""" + source_short = _source("source:s2", {"id": "source:s2", "insights": ["i"]}) + source_long = _source("source:s3", {"id": "source:s3", "full_text": "body"}) + note_full = _note("note:n2", {"id": "note:n2", "content": "note body"}) + + async def get_source(full_id): + return {"source:s2": source_short, "source:s3": source_long}[full_id] + + with ( + patch(PATCH_NOTEBOOK, new=AsyncMock(return_value=_notebook())), + patch(PATCH_SOURCE, new=AsyncMock(side_effect=get_source)), + patch(PATCH_NOTE, new=AsyncMock(return_value=note_full)), + ): + response = client.post( + "/api/chat/context", + json={ + "notebook_id": "notebook:1", + "context_config": { + "sources": { + "s1": "not in context", + "s2": "insights", + "s3": "full content", + }, + "notes": {"n1": "not in context", "n2": "full content"}, + }, + }, + ) + + assert response.status_code == 200 + body = response.json() + assert body["context"]["sources"] == [ + {"id": "source:s2", "insights": ["i"]}, + {"id": "source:s3", "full_text": "body"}, + ] + assert body["context"]["notes"] == [{"id": "note:n2", "content": "note body"}] + + # "insights" -> short context, "full content" -> long context + source_short.get_context.assert_awaited_once_with(context_size="short") + source_long.get_context.assert_awaited_once_with(context_size="long") + note_full.get_context.assert_called_once_with(context_size="long") + + # char_count is the length of the concatenated str() of every context dict + expected_content = ( + str({"id": "source:s2", "insights": ["i"]}) + + str({"id": "source:s3", "full_text": "body"}) + + str({"id": "note:n2", "content": "note body"}) + ) + assert body["char_count"] == len(expected_content) + assert isinstance(body["token_count"], int) + assert body["token_count"] > 0 + + +@pytest.mark.asyncio +async def test_config_bare_ids_get_table_prefix(client): + """Bare ids are prefixed with 'source:' / 'note:' before lookup.""" + source_get = AsyncMock(return_value=_source()) + note_get = AsyncMock(return_value=_note()) + + with ( + patch(PATCH_NOTEBOOK, new=AsyncMock(return_value=_notebook())), + patch(PATCH_SOURCE, new=source_get), + patch(PATCH_NOTE, new=note_get), + ): + response = client.post( + "/api/chat/context", + json={ + "notebook_id": "notebook:1", + "context_config": { + "sources": {"abc": "insights"}, + "notes": {"def": "full content"}, + }, + }, + ) + + assert response.status_code == 200 + source_get.assert_awaited_once_with("source:abc") + note_get.assert_awaited_once_with("note:def") + + +@pytest.mark.asyncio +async def test_config_missing_source_is_skipped(client): + """A source lookup failure skips that source, not the whole request.""" + ok_source = _source("source:ok") + + async def get_source(full_id): + if full_id == "source:missing": + raise Exception("not found") + return ok_source + + with ( + patch(PATCH_NOTEBOOK, new=AsyncMock(return_value=_notebook())), + patch(PATCH_SOURCE, new=AsyncMock(side_effect=get_source)), + ): + response = client.post( + "/api/chat/context", + json={ + "notebook_id": "notebook:1", + "context_config": { + "sources": {"missing": "insights", "ok": "insights"}, + "notes": {}, + }, + }, + ) + + assert response.status_code == 200 + assert response.json()["context"]["sources"] == [{"id": "source:ok", "title": "T"}] + + +@pytest.mark.asyncio +async def test_empty_config_defaults_to_all_short_contexts(client): + """Falsy context_config -> every source (batched insights) and note, short.""" + src = _source("source:s1", {"id": "source:s1", "title": "S"}) + note = _note("note:n1", {"id": "note:n1", "title": "N"}) + notebook = _notebook() + notebook.get_sources = AsyncMock(return_value=[src]) + notebook.get_notes = AsyncMock(return_value=[note]) + insights = [SimpleNamespace(id="source_insight:1")] + + with ( + patch(PATCH_NOTEBOOK, new=AsyncMock(return_value=notebook)), + patch( + PATCH_INSIGHTS, + new=AsyncMock(return_value={"source:s1": insights}), + ), + ): + response = client.post( + "/api/chat/context", + json={"notebook_id": "notebook:1", "context_config": {}}, + ) + + assert response.status_code == 200 + body = response.json() + assert body["context"]["sources"] == [{"id": "source:s1", "title": "S"}] + assert body["context"]["notes"] == [{"id": "note:n1", "title": "N"}] + src.get_context.assert_awaited_once_with(context_size="short", insights=insights) + note.get_context.assert_called_once_with(context_size="short") + + +@pytest.mark.asyncio +async def test_default_path_survives_insight_batch_failure(client): + """A failure batch-fetching insights falls back to empty insights.""" + src = _source("source:s1") + notebook = _notebook() + notebook.get_sources = AsyncMock(return_value=[src]) + + with ( + patch(PATCH_NOTEBOOK, new=AsyncMock(return_value=notebook)), + patch(PATCH_INSIGHTS, new=AsyncMock(side_effect=Exception("db hiccup"))), + ): + response = client.post( + "/api/chat/context", + json={"notebook_id": "notebook:1", "context_config": {}}, + ) + + assert response.status_code == 200 + src.get_context.assert_awaited_once_with(context_size="short", insights=[]) + + +@pytest.mark.asyncio +async def test_missing_notebook_returns_404(client): + with patch(PATCH_NOTEBOOK, new=AsyncMock(return_value=None)): + response = client.post( + "/api/chat/context", + json={"notebook_id": "notebook:missing", "context_config": {}}, + ) + + assert response.status_code == 404 diff --git a/tests/test_typed_exceptions_reach_handlers.py b/tests/test_typed_exceptions_reach_handlers.py index f09f4ca5..7eaf7e73 100644 --- a/tests/test_typed_exceptions_reach_handlers.py +++ b/tests/test_typed_exceptions_reach_handlers.py @@ -47,7 +47,6 @@ CASES = [ ("notes", "api.routers.notes.Note.get_all", "GET", "/api/notes", None), ("models", "api.routers.models.Model.get_all", "GET", "/api/models", None), ("commands", "api.routers.commands.CommandService.get_command_status", "GET", "/api/commands/jobs/command:abc", None), - ("context", "api.routers.context.Notebook.get", "POST", "/api/notebooks/notebook:1/context", {"notebook_id": "notebook:1"}), ("credentials", "api.routers.credentials.Credential.get_all", "GET", "/api/credentials", None), ("embedding", "api.routers.embedding.model_manager.get_embedding_model", "POST", "/api/embed", {"item_id": "source:1", "item_type": "source"}), ("embedding_rebuild", "api.routers.embedding_rebuild.repo_query", "POST", "/api/embeddings/rebuild", {"mode": "existing"}), diff --git a/tests/test_utils.py b/tests/test_utils.py index cf1ebc11..5917fe20 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -5,6 +5,9 @@ This test suite focuses on testing utility functions that perform actual logic without heavy mocking - string processing, validation, and algorithms. """ +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + import pytest from open_notebook.utils import ( @@ -16,7 +19,7 @@ from open_notebook.utils import ( remove_non_printable, token_count, ) -from open_notebook.utils.context_builder import ContextBuilder, ContextConfig +from open_notebook.utils.context_builder import build_source_context # ============================================================================ # TEST SUITE 1: Text Utilities @@ -246,39 +249,94 @@ class TestVersionUtilities: # ============================================================================ -# TEST SUITE 4: Context Builder Configuration +# TEST SUITE 4: Source Context Building # ============================================================================ -class TestContextBuilder: - """Test suite for ContextBuilder initialization and configuration.""" +def _mock_source(insights): + source = SimpleNamespace(id="source:123") + source.get_context = AsyncMock( + return_value={"id": "source:123", "title": "T", "full_text": "body"} + ) + source.get_insights = AsyncMock(return_value=insights) + return source - def test_context_config_defaults(self): - """Test ContextConfig default values.""" - config = ContextConfig() - assert config.sources == {} - assert config.notes == {} - assert config.include_insights is True - assert config.include_notes is True - assert config.priority_weights is not None - assert "source" in config.priority_weights - assert "note" in config.priority_weights - assert "insight" in config.priority_weights +def _insight(insight_id, content="insight content"): + return SimpleNamespace( + id=insight_id, insight_type="summary", content=content + ) - def test_context_builder_initialization(self): - """Test ContextBuilder initialization with various params.""" - builder = ContextBuilder( - source_id="source:123", - notebook_id="notebook:456", - max_tokens=1000, - include_insights=False, + +class TestBuildSourceContext: + """Test suite for build_source_context (used by the source-chat graph).""" + + @pytest.mark.asyncio + async def test_source_and_insights_shape(self): + """The response carries the source's short context and its insights.""" + source = _mock_source([_insight("source_insight:1")]) + + with patch( + "open_notebook.utils.context_builder.Source.get", + new=AsyncMock(return_value=source), + ) as mock_get: + result = await build_source_context("123") + + mock_get.assert_awaited_once_with("source:123") # bare id gets prefixed + source.get_context.assert_awaited_once_with(context_size="short") + assert result["sources"] == [ + {"id": "source:123", "title": "T", "full_text": "body"} + ] + assert result["insights"] == [ + { + "id": "source_insight:1", + "source_id": "source:123", + "insight_type": "summary", + "content": "insight content", + } + ] + assert result["notes"] == [] + assert result["total_items"] == 2 + assert result["total_tokens"] > 0 + assert result["metadata"] == { + "source_count": 1, + "note_count": 0, + "insight_count": 1, + } + + @pytest.mark.asyncio + async def test_truncates_insights_to_token_budget(self): + """Insights are dropped (last first) when over the token budget.""" + big = "word " * 500 + source = _mock_source( + [_insight("source_insight:1", big), _insight("source_insight:2", big)] ) - assert builder.source_id == "source:123" - assert builder.notebook_id == "notebook:456" - assert builder.max_tokens == 1000 - assert builder.include_insights is False + with patch( + "open_notebook.utils.context_builder.Source.get", + new=AsyncMock(return_value=source), + ): + result = await build_source_context("source:123", max_tokens=600) + + # Budget fits the source and the first insight, not the second + assert len(result["sources"]) == 1 + assert [i["id"] for i in result["insights"]] == ["source_insight:1"] + assert result["total_tokens"] <= 600 + + @pytest.mark.asyncio + async def test_missing_source_yields_empty_context(self): + """A missing source produces an empty context, not an error.""" + from open_notebook.exceptions import NotFoundError + + with patch( + "open_notebook.utils.context_builder.Source.get", + new=AsyncMock(side_effect=NotFoundError("nope")), + ): + result = await build_source_context("source:missing") + + assert result["sources"] == [] + assert result["insights"] == [] + assert result["total_tokens"] == 0 if __name__ == "__main__":