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) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
61cadf7935
commit
90458b6f08
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user