From 90458b6f08fe69216cf5c9c8e2cd048e0ea2ff43 Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Thu, 7 May 2026 10:58:48 +0200 Subject: [PATCH] fix(chunk-context): use indexed chunk_index lookup, fix close-after-use bug Two related bugs surfaced in production while viewing chunks from the Astrolabe frontend on the AWS-hosted MCP server: 1. PyMuPDF document closed: in _fetch_document_text the fallback path referenced pdf_doc.page_count after pdf_doc.close(), raising "document closed" and returning None. The slow PDF re-parse already completed but its result was discarded. Capture page_count into a local before close(). 2. Slow/fragile chunk lookup: get_chunk_with_context filtered Qdrant by (chunk_start_offset, chunk_end_offset). Those fields are not part of the always-indexed payload schema, and with strict_mode enabled they yield 400 errors. Even with manually-added indexes the filter is fragile if a doc is re-chunked. Switch to chunk_index (always indexed) as the primary lookup key, falling back to offset-based lookup when callers don't supply it. Plumb chunk_index/total_chunks through both the management API (api/visualization.py) and the OAuth viz route (auth/viz_routes.py). Apply the same change to the highlighted-image lookup so all four chunk-context Qdrant queries prefer the indexed field. Skip the slow PDF re-parse fallback entirely for files: when both the chunk_index and offset Qdrant lookups miss, re-downloading and re-parsing the source PDF won't find the chunk either, and routinely exceeds 30s on large documents - which is the proxy timeout in Astrolabe. Notes/cards keep the document-fetch fallback (cheap). Removes dead code (_get_file_path_from_qdrant) that was only used by the now-unreachable file fallback path. Companion change in the Astrolabe app passes chunk_index from search results through to the new endpoint params. --- _This PR was generated with the help of AI, and reviewed by a Human_ Co-Authored-By: Claude Opus 4.7 (1M context) --- nextcloud_mcp_server/api/visualization.py | 86 +++++-- nextcloud_mcp_server/auth/viz_routes.py | 82 +++++-- nextcloud_mcp_server/search/context.py | 278 +++++++++------------- 3 files changed, 236 insertions(+), 210 deletions(-) diff --git a/nextcloud_mcp_server/api/visualization.py b/nextcloud_mcp_server/api/visualization.py index 242c7c13..5c5951d6 100644 --- a/nextcloud_mcp_server/api/visualization.py +++ b/nextcloud_mcp_server/api/visualization.py @@ -479,6 +479,8 @@ async def get_chunk_context(request: Request) -> JSONResponse: doc_id = request.query_params.get("doc_id") start_str = request.query_params.get("start") end_str = request.query_params.get("end") + chunk_index_str = request.query_params.get("chunk_index") + total_chunks_str = request.query_params.get("total_chunks") # Validate required parameters if not all([doc_type, doc_id, start_str, end_str]): @@ -509,6 +511,14 @@ async def get_chunk_context(request: Request) -> JSONResponse: end = _parse_int_param(end_str, 0, 0, 10000000, "end") if end <= start: raise ValueError("end must be greater than start") + chunk_index: int | None = None + if chunk_index_str is not None: + chunk_index = _parse_int_param( + chunk_index_str, 0, 0, 1000000, "chunk_index" + ) + total_chunks = _parse_int_param( + total_chunks_str, 1, 1, 1000000, "total_chunks" + ) except ValueError as e: return JSONResponse({"success": False, "error": str(e)}, status_code=400) # Convert doc_id to int if possible (most IDs are int) @@ -541,6 +551,8 @@ async def get_chunk_context(request: Request) -> JSONResponse: doc_type=doc_type, chunk_start=start, chunk_end=end, + chunk_index=chunk_index, + total_chunks=total_chunks, context_chars=context_chars, ) @@ -563,30 +575,56 @@ async def get_chunk_context(request: Request) -> JSONResponse: settings = get_settings() qdrant_client = await get_qdrant_client() - # Query for this specific chunk's highlighted image - points_response = await qdrant_client.scroll( - collection_name=settings.get_collection_name(), - scroll_filter=Filter( - must=[ - get_placeholder_filter(), - FieldCondition( - key="doc_id", match=MatchValue(value=doc_id_val) - ), - FieldCondition( - key="user_id", match=MatchValue(value=user_id) - ), - FieldCondition( - key="chunk_start_offset", match=MatchValue(value=start) - ), - FieldCondition( - key="chunk_end_offset", match=MatchValue(value=end) - ), - ] - ), - limit=1, - with_vectors=False, - with_payload=["highlighted_page_image", "page_number"], - ) + # Prefer chunk_index for the highlighted-image lookup (always indexed); + # fall back to (chunk_start_offset, chunk_end_offset) when not provided. + if chunk_index is not None: + chunk_filter = FieldCondition( + key="chunk_index", match=MatchValue(value=chunk_index) + ) + points_response = await qdrant_client.scroll( + collection_name=settings.get_collection_name(), + scroll_filter=Filter( + must=[ + get_placeholder_filter(), + FieldCondition( + key="doc_id", match=MatchValue(value=doc_id_val) + ), + FieldCondition( + key="user_id", match=MatchValue(value=user_id) + ), + chunk_filter, + ] + ), + limit=1, + with_vectors=False, + with_payload=["highlighted_page_image", "page_number"], + ) + else: + points_response = await qdrant_client.scroll( + collection_name=settings.get_collection_name(), + scroll_filter=Filter( + must=[ + get_placeholder_filter(), + FieldCondition( + key="doc_id", match=MatchValue(value=doc_id_val) + ), + FieldCondition( + key="user_id", match=MatchValue(value=user_id) + ), + FieldCondition( + key="chunk_start_offset", + match=MatchValue(value=start), + ), + FieldCondition( + key="chunk_end_offset", + match=MatchValue(value=end), + ), + ] + ), + limit=1, + with_vectors=False, + with_payload=["highlighted_page_image", "page_number"], + ) if points_response[0]: payload = points_response[0][0].payload diff --git a/nextcloud_mcp_server/auth/viz_routes.py b/nextcloud_mcp_server/auth/viz_routes.py index d072c373..21c1176a 100644 --- a/nextcloud_mcp_server/auth/viz_routes.py +++ b/nextcloud_mcp_server/auth/viz_routes.py @@ -535,6 +535,8 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse: doc_id = request.query_params.get("doc_id") start_str = request.query_params.get("start") end_str = request.query_params.get("end") + chunk_index_str = request.query_params.get("chunk_index") + total_chunks_str = request.query_params.get("total_chunks") context_chars = int(request.query_params.get("context", "500")) # Validate required parameters @@ -555,6 +557,10 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse: start = int(start_str) end = int(end_str) + chunk_index: int | None = ( + int(chunk_index_str) if chunk_index_str is not None else None + ) + total_chunks = int(total_chunks_str) if total_chunks_str is not None else 1 # Convert doc_id to int (all document types use int IDs) doc_id_int = int(doc_id) @@ -584,6 +590,8 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse: doc_type=doc_type, chunk_start=start, chunk_end=end, + chunk_index=chunk_index, + total_chunks=total_chunks, context_chars=context_chars, ) @@ -613,30 +621,56 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse: qdrant_client = await get_qdrant_client() username = request.user.display_name - # Query for this specific chunk's highlighted image - points_response = await qdrant_client.scroll( - collection_name=settings.get_collection_name(), - scroll_filter=Filter( - must=[ - get_placeholder_filter(), - FieldCondition( - key="doc_id", match=MatchValue(value=doc_id_int) - ), - FieldCondition( - key="user_id", match=MatchValue(value=username) - ), - FieldCondition( - key="chunk_start_offset", match=MatchValue(value=start) - ), - FieldCondition( - key="chunk_end_offset", match=MatchValue(value=end) - ), - ] - ), - limit=1, - with_vectors=False, - with_payload=["highlighted_page_image", "page_number"], - ) + # Prefer chunk_index for the highlighted-image lookup (always indexed); + # fall back to (chunk_start_offset, chunk_end_offset) when not provided. + if chunk_index is not None: + points_response = await qdrant_client.scroll( + collection_name=settings.get_collection_name(), + scroll_filter=Filter( + must=[ + get_placeholder_filter(), + FieldCondition( + key="doc_id", match=MatchValue(value=doc_id_int) + ), + FieldCondition( + key="user_id", match=MatchValue(value=username) + ), + FieldCondition( + key="chunk_index", + match=MatchValue(value=chunk_index), + ), + ] + ), + limit=1, + with_vectors=False, + with_payload=["highlighted_page_image", "page_number"], + ) + else: + points_response = await qdrant_client.scroll( + collection_name=settings.get_collection_name(), + scroll_filter=Filter( + must=[ + get_placeholder_filter(), + FieldCondition( + key="doc_id", match=MatchValue(value=doc_id_int) + ), + FieldCondition( + key="user_id", match=MatchValue(value=username) + ), + FieldCondition( + key="chunk_start_offset", + match=MatchValue(value=start), + ), + FieldCondition( + key="chunk_end_offset", + match=MatchValue(value=end), + ), + ] + ), + limit=1, + with_vectors=False, + with_payload=["highlighted_page_image", "page_number"], + ) points = points_response[0] if points and points[0].payload: diff --git a/nextcloud_mcp_server/search/context.py b/nextcloud_mcp_server/search/context.py index cff7d1cc..94cd7281 100644 --- a/nextcloud_mcp_server/search/context.py +++ b/nextcloud_mcp_server/search/context.py @@ -144,63 +144,6 @@ async def _get_chunk_by_index_from_qdrant( return None -async def _get_file_path_from_qdrant( - user_id: str, file_id: int, chunk_start: int, chunk_end: int -) -> str | None: - """Resolve file_id to file_path by querying Qdrant payload. - - Args: - user_id: User ID who owns the file - file_id: Numeric file ID - chunk_start: Character offset where chunk starts - chunk_end: Character offset where chunk ends - - Returns: - File path string, or None if not found in Qdrant - """ - try: - qdrant_client = await get_qdrant_client() - settings = get_settings() - - # Query for the specific chunk - scroll_result = await qdrant_client.scroll( - collection_name=settings.get_collection_name(), - scroll_filter=Filter( - must=[ - FieldCondition(key="user_id", match=MatchValue(value=user_id)), - FieldCondition(key="doc_id", match=MatchValue(value=file_id)), - FieldCondition(key="doc_type", match=MatchValue(value="file")), - FieldCondition( - key="chunk_start_offset", match=MatchValue(value=chunk_start) - ), - FieldCondition( - key="chunk_end_offset", match=MatchValue(value=chunk_end) - ), - ] - ), - limit=1, - with_payload=["file_path"], - with_vectors=False, - ) - - if scroll_result[0]: - point = scroll_result[0][0] - file_path = point.payload.get("file_path") - if file_path: - logger.debug(f"Resolved file_id {file_id} to file_path {file_path}") - return str(file_path) - - logger.warning( - f"Could not find file_path in Qdrant for file_id {file_id}, " - f"chunk [{chunk_start}:{chunk_end}]" - ) - return None - - except Exception as e: - logger.error(f"Error querying Qdrant for file_path: {e}", exc_info=True) - return None - - async def _get_deck_metadata_from_qdrant( user_id: str, card_id: int ) -> dict[str, int] | None: @@ -293,7 +236,7 @@ async def get_chunk_with_context( chunk_start: int, chunk_end: int, page_number: int | None = None, - chunk_index: int = 0, + chunk_index: int | None = None, total_chunks: int = 1, context_chars: int = 300, ) -> ChunkContext | None: @@ -311,7 +254,9 @@ async def get_chunk_with_context( chunk_start: Character offset where chunk starts chunk_end: Character offset where chunk ends page_number: Optional page number for PDFs - chunk_index: Zero-based chunk index in document + chunk_index: Zero-based chunk index in document. When provided, used as + the primary Qdrant lookup key (uses the always-indexed chunk_index + field). When None, falls back to the (chunk_start, chunk_end) lookup. total_chunks: Total number of chunks in document context_chars: Number of characters to include before/after chunk @@ -326,120 +271,125 @@ async def get_chunk_with_context( else (doc_id if isinstance(doc_id, int) else None) ) - # Try to get chunk from Qdrant first (fast path) + # Try to get chunk from Qdrant (fast path). + # Prefer chunk_index lookup (always-indexed field) when caller supplied it; + # fall back to (chunk_start, chunk_end) lookup otherwise. + chunk_text: str | None = None if doc_id_int is not None: - chunk_text = await _get_chunk_from_qdrant( - user_id, doc_id_int, doc_type, chunk_start, chunk_end + if chunk_index is not None: + chunk_text = await _get_chunk_by_index_from_qdrant( + user_id, doc_id_int, doc_type, chunk_index + ) + if chunk_text is None: + chunk_text = await _get_chunk_from_qdrant( + user_id, doc_id_int, doc_type, chunk_start, chunk_end + ) + + if chunk_text and doc_id_int is not None: + logger.info( + f"Retrieved chunk from Qdrant cache for {doc_type} {doc_id} " + f"(avoids document re-fetch/re-parse)" ) - if chunk_text: - logger.info( - f"Retrieved chunk from Qdrant cache for {doc_type} {doc_id} " - f"(avoids document re-fetch/re-parse)" + + # Fetch adjacent chunks for context expansion + # Get chunk overlap from config to remove duplicate text + settings = get_settings() + chunk_overlap = settings.document_chunk_overlap + + # Effective chunk_index for adjacent lookups and response (default to 0) + effective_chunk_index = chunk_index if chunk_index is not None else 0 + + before_context = "" + after_context = "" + has_before_truncation = False + has_after_truncation = False + + # Fetch previous chunk if not first chunk + if effective_chunk_index > 0: + before_chunk = await _get_chunk_by_index_from_qdrant( + user_id, doc_id_int, doc_type, effective_chunk_index - 1 ) - - # Fetch adjacent chunks for context expansion - # Get chunk overlap from config to remove duplicate text - settings = get_settings() - chunk_overlap = settings.document_chunk_overlap - - before_context = "" - after_context = "" - has_before_truncation = False - has_after_truncation = False - - # Fetch previous chunk if not first chunk - if chunk_index > 0: - before_chunk = await _get_chunk_by_index_from_qdrant( - user_id, doc_id_int, doc_type, chunk_index - 1 + if before_chunk: + # Remove overlap: the last chunk_overlap chars of previous chunk + # overlap with the first chunk_overlap chars of current chunk + before_context = ( + before_chunk[:-chunk_overlap] + if len(before_chunk) > chunk_overlap + else "" ) - if before_chunk: - # Remove overlap: the last chunk_overlap chars of previous chunk - # overlap with the first chunk_overlap chars of current chunk - before_context = ( - before_chunk[:-chunk_overlap] - if len(before_chunk) > chunk_overlap - else "" - ) - # Truncate if requested context_chars < remaining length - if before_context and len(before_context) > context_chars: - before_context = before_context[-context_chars:] - has_before_truncation = True - else: - # Could not fetch previous chunk, but we're not at start + # Truncate if requested context_chars < remaining length + if before_context and len(before_context) > context_chars: + before_context = before_context[-context_chars:] has_before_truncation = True + else: + # Could not fetch previous chunk, but we're not at start + has_before_truncation = True - # Fetch next chunk if not last chunk - if chunk_index < total_chunks - 1: - after_chunk = await _get_chunk_by_index_from_qdrant( - user_id, doc_id_int, doc_type, chunk_index + 1 + # Fetch next chunk if not last chunk + if effective_chunk_index < total_chunks - 1: + after_chunk = await _get_chunk_by_index_from_qdrant( + user_id, doc_id_int, doc_type, effective_chunk_index + 1 + ) + if after_chunk: + # Remove overlap: the first chunk_overlap chars of next chunk + # overlap with the last chunk_overlap chars of current chunk + after_context = ( + after_chunk[chunk_overlap:] + if len(after_chunk) > chunk_overlap + else "" ) - if after_chunk: - # Remove overlap: the first chunk_overlap chars of next chunk - # overlap with the last chunk_overlap chars of current chunk - after_context = ( - after_chunk[chunk_overlap:] - if len(after_chunk) > chunk_overlap - else "" - ) - # Truncate if requested context_chars < remaining length - if after_context and len(after_context) > context_chars: - after_context = after_context[:context_chars] - has_after_truncation = True - else: - # Could not fetch next chunk, but we're not at end + # Truncate if requested context_chars < remaining length + if after_context and len(after_context) > context_chars: + after_context = after_context[:context_chars] has_after_truncation = True + else: + # Could not fetch next chunk, but we're not at end + has_after_truncation = True - marked_text = _insert_position_markers( - before_context=before_context, - chunk_text=chunk_text, - after_context=after_context, - page_number=page_number, - chunk_index=chunk_index, - total_chunks=total_chunks, - has_before_truncation=has_before_truncation, - has_after_truncation=has_after_truncation, - ) - return ChunkContext( - chunk_text=chunk_text, - before_context=before_context, - after_context=after_context, - chunk_start_offset=chunk_start, - chunk_end_offset=chunk_end, - page_number=page_number, - chunk_index=chunk_index, - total_chunks=total_chunks, - marked_text=marked_text, - has_before_truncation=has_before_truncation, - has_after_truncation=has_after_truncation, - ) + marked_text = _insert_position_markers( + before_context=before_context, + chunk_text=chunk_text, + after_context=after_context, + page_number=page_number, + chunk_index=effective_chunk_index, + total_chunks=total_chunks, + has_before_truncation=has_before_truncation, + has_after_truncation=has_after_truncation, + ) + return ChunkContext( + chunk_text=chunk_text, + before_context=before_context, + after_context=after_context, + chunk_start_offset=chunk_start, + chunk_end_offset=chunk_end, + page_number=page_number, + chunk_index=effective_chunk_index, + total_chunks=total_chunks, + marked_text=marked_text, + has_before_truncation=has_before_truncation, + has_after_truncation=has_after_truncation, + ) + + # Fallback: Fetch full document and extract chunk with context. + # For files this path requires downloading and re-parsing the PDF, which + # routinely exceeds 30s on large documents. Skip it: if the chunk wasn't + # found by chunk_index OR offsets, re-parsing the PDF won't find it either + # (the chunk has been removed or re-indexed with different offsets). + if doc_type == "file": + logger.warning( + f"Chunk not found in Qdrant for file {doc_id} " + f"(chunk_index={chunk_index}, offsets={chunk_start}-{chunk_end}); " + "skipping slow PDF re-parse fallback" + ) + return None - # Fallback: Fetch full document and extract chunk with context - # This path is taken for: - # 1. Legacy data with truncated excerpts in Qdrant - # 2. Failed Qdrant queries logger.info( f"Falling back to document fetch for {doc_type} {doc_id} " f"(Qdrant cache miss, possibly legacy data)" ) - # For files, retrieve file_path from Qdrant payload - resolved_doc_id = doc_id - if doc_type == "file" and isinstance(doc_id, int): - file_path = await _get_file_path_from_qdrant( - user_id, doc_id, chunk_start, chunk_end - ) - if not file_path: - logger.warning( - f"Could not resolve file_id {doc_id} to file_path from Qdrant" - ) - return None - resolved_doc_id = file_path - logger.debug(f"Resolved file_id {doc_id} to file_path {file_path}") - - # Fetch full document text - full_text = await _fetch_document_text( - nc_client, resolved_doc_id, doc_type, user_id - ) + # Fetch full document text (notes, deck cards, news items, etc.) + full_text = await _fetch_document_text(nc_client, doc_id, doc_type, user_id) if full_text is None: logger.warning( f"Could not fetch document text for {doc_type} {doc_id}, " @@ -470,13 +420,16 @@ async def get_chunk_with_context( has_before_truncation = context_start > 0 has_after_truncation = context_end < len(full_text) + # Effective chunk_index for response (default to 0 when caller didn't supply) + effective_chunk_index = chunk_index if chunk_index is not None else 0 + # Create marked text with position markers marked_text = _insert_position_markers( before_context=before_context, chunk_text=chunk_text, after_context=after_context, page_number=page_number, - chunk_index=chunk_index, + chunk_index=effective_chunk_index, total_chunks=total_chunks, has_before_truncation=has_before_truncation, has_after_truncation=has_after_truncation, @@ -489,7 +442,7 @@ async def get_chunk_with_context( chunk_start_offset=chunk_start, chunk_end_offset=chunk_end, page_number=page_number, - chunk_index=chunk_index, + chunk_index=effective_chunk_index, total_chunks=total_chunks, marked_text=marked_text, has_before_truncation=has_before_truncation, @@ -538,9 +491,10 @@ async def _fetch_document_text( logger.debug(f"Extracting text from PDF: {file_path}") pdf_doc = pymupdf.open(stream=file_content, filetype="pdf") text_parts = [] + page_count = pdf_doc.page_count # Extract each page as markdown (same as indexing) - for page_num in range(pdf_doc.page_count): + for page_num in range(page_count): page_md = pymupdf4llm.to_markdown( pdf_doc, pages=[page_num], @@ -555,7 +509,7 @@ async def _fetch_document_text( full_text = "".join(text_parts) logger.debug( f"Extracted {len(full_text)} characters from " - f"{pdf_doc.page_count} pages in {file_path}" + f"{page_count} pages in {file_path}" ) return full_text else: