| import asyncio |
| from collections.abc import Callable |
| from typing import Any |
|
|
| from src.server.config.logfire_config import get_logger, safe_logfire_info |
|
|
| logger = get_logger(__name__) |
|
|
|
|
| class ChunkingUtils: |
| """Handles logic for splitting markdown documents into chunks without DB interaction.""" |
|
|
| def __init__(self, supabase_client): |
| from src.server.services.storage.storage_services import DocumentStorageService |
|
|
| self.doc_storage_service = DocumentStorageService(supabase_client) |
|
|
| async def prepare_document_chunks( |
| self, |
| crawl_results: list[dict], |
| request: dict[str, Any], |
| crawl_type: str, |
| original_source_id: str, |
| cancellation_check: Callable | None = None, |
| ) -> tuple[list[str], list[int], list[str], list[dict], dict[str, int], dict[str, str], int]: |
| """ |
| Splits documents into chunks and prepares metadatas. |
| |
| Returns: |
| (all_urls, all_chunk_numbers, all_contents, all_metadatas, source_word_counts, url_to_full_document, processed_docs) |
| """ |
| all_urls = [] |
| all_chunk_numbers = [] |
| all_contents = [] |
| all_metadatas = [] |
| source_word_counts: dict[str, int] = {} |
| url_to_full_document = {} |
| processed_docs = 0 |
|
|
| for doc_index, doc in enumerate(crawl_results): |
| if cancellation_check: |
| cancellation_check() |
|
|
| doc_url = (doc.get("url") or "").strip() |
| markdown_content = (doc.get("markdown") or "").strip() |
|
|
| if not markdown_content or not doc_url: |
| logger.debug(f"Skipping document {doc_index}: empty {'URL' if not doc_url else 'content'}") |
| continue |
|
|
| processed_docs += 1 |
| url_to_full_document[doc_url] = markdown_content |
|
|
| |
| chunks = await self.doc_storage_service.smart_chunk_text_async(markdown_content, chunk_size=5000) |
|
|
| source_id = original_source_id |
| safe_logfire_info(f"Using original source_id '{source_id}' for URL '{doc_url}'") |
|
|
| for i, chunk in enumerate(chunks): |
| if cancellation_check and i % 10 == 0: |
| cancellation_check() |
|
|
| all_urls.append(doc_url) |
| all_chunk_numbers.append(i) |
| all_contents.append(chunk) |
|
|
| word_count = len(chunk.split()) |
| metadata = { |
| "url": doc_url, |
| "title": doc.get("title", ""), |
| "description": doc.get("description", ""), |
| "source_id": source_id, |
| "knowledge_type": request.get("knowledge_type", "documentation"), |
| "crawl_type": crawl_type, |
| "word_count": word_count, |
| "char_count": len(chunk), |
| "chunk_index": i, |
| "tags": request.get("tags", []), |
| } |
| all_metadatas.append(metadata) |
|
|
| source_word_counts[source_id] = source_word_counts.get(source_id, 0) + word_count |
|
|
| if i > 0 and i % 10 == 0: |
| await asyncio.sleep(0) |
|
|
| if doc_index > 0 and doc_index % 5 == 0: |
| await asyncio.sleep(0) |
|
|
| return ( |
| all_urls, |
| all_chunk_numbers, |
| all_contents, |
| all_metadatas, |
| source_word_counts, |
| url_to_full_document, |
| processed_docs, |
| ) |
|
|
