from collections.abc import AsyncIterator from typing import Any PATH_TREE_ID = "__path_tree__" PAGE_CHUNK_BATCH_SIZE = 100 async def fetch_path_tree(accessor) -> str: collection = await accessor.get_collection() result = await collection.get(ids=[PATH_TREE_ID]) documents = result.get("documents") or [] if not documents: raise FileNotFoundError(PATH_TREE_ID) value = documents[0] if value is None: raise FileNotFoundError(PATH_TREE_ID) if isinstance(value, str): return value return str(value) async def fetch_page_chunks(accessor, slug: str) -> str: chunks = await page_chunks(accessor, slug) return "\n".join(chunk["document"] for chunk in chunks) async def iter_page_chunks(accessor, slug: str) -> AsyncIterator[str]: chunks = await page_chunks(accessor, slug) for chunk in chunks: yield chunk["document"] async def page_chunks(accessor, slug: str) -> list[dict[str, Any]]: collection = await accessor.get_collection() chunks: list[dict[str, Any]] = [] offset = 0 while True: result = await collection.get( where={accessor.config.slug_field: slug}, include=["documents", "metadatas"], limit=PAGE_CHUNK_BATCH_SIZE, offset=offset, ) documents = result.get("documents") or [] metadatas = result.get("metadatas") or [{} for _ in documents] for document, metadata in zip(documents, metadatas, strict=True): chunks.append({ "document": "" if document is None else str(document), "metadata": metadata if isinstance(metadata, dict) else {}, }) if len(documents) < PAGE_CHUNK_BATCH_SIZE: break offset += PAGE_CHUNK_BATCH_SIZE return sorted(chunks, key=lambda item: chunk_index( item["metadata"], accessor.config.chunk_index_field)) async def query_contains( accessor, pattern: str, candidate_slugs: list[str], *, regex: bool = False, ) -> list[str]: if not candidate_slugs: return [] collection = await accessor.get_collection() result = await collection.get( where={accessor.config.slug_field: { "$in": candidate_slugs }}, where_document={"$regex" if regex else "$contains": pattern}, include=["metadatas"], ) matched: set[str] = set() for metadata in result.get("metadatas") or []: if isinstance(metadata, dict): slug = metadata.get(accessor.config.slug_field) if slug is not None: matched.add(str(slug)) return sorted(matched) def chunk_index(metadata: dict[str, Any], field: str) -> int: value = metadata.get(field, 0) if isinstance(value, bool): return 0 if isinstance(value, int): return value if isinstance(value, float): return int(value) if isinstance(value, str) and value.isdigit(): return int(value) return 0