mirror of
https://github.com/lfnovo/open-notebook.git
synced 2026-10-19 09:13:16 +00:00
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.
This commit is contained in:
Luis Novo
GitHub
parent
c4749ad041
commit
11ca9c138c
13 changed files with 514 additions and 764 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)}")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)}")
|
||||
|
|
|
|||
229
tests/test_context_endpoint_characterization.py
Normal file
229
tests/test_context_endpoint_characterization.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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"}),
|
||||
|
|
|
|||
|
|
@ -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__":
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue