diff --git a/mosaic.py b/mosaic.py index 943daca..84088e8 100755 --- a/mosaic.py +++ b/mosaic.py @@ -224,10 +224,28 @@ def collect(ctx, source, url, name, profile_path, limit, since, date_range, upda console.print("[yellow]No source specified.[/yellow] Use --source, --url, or --profile.") raise SystemExit(1) - label = source or url or profile_path + from collectors.profile_loader import load_profile, create_adhoc_profile + from collectors.custom import CustomSource + from utils.cache import DownloadCache + from utils.http_limiter import HTTPRateLimiter + + # Load profile + try: + if profile_path: + profile = load_profile(profile_path) + elif url: + pname = name or url.split('/')[2].replace('.', '_') + profile = create_adhoc_profile(url, pname) + else: + profile = load_profile(source) + except (FileNotFoundError, ValueError) as e: + console.print(f"[red]Profile error: {e}[/red]") + raise SystemExit(1) + + label = profile['name'] console.print(Panel( f"[bold cyan]Collect[/bold cyan]\n\n" - f" Source : {label}\n" + f" Source : {label} — {profile.get('description', '')}\n" f" Limit : {limit or 'all'}\n" f" Since : {since or 'n/a'}\n" f" Range : {' to '.join(date_range) if date_range else 'n/a'}\n" @@ -235,9 +253,40 @@ def collect(ctx, source, url, name, profile_path, limit, since, date_range, upda title="[bold]mosaic collect[/bold]", border_style="cyan", )) - console.print( - "\n[dim]Collection not yet implemented — see collectors/[/dim]" - ) + + db = open_db(ctx.obj["db_path"]) + try: + cache = DownloadCache(os.path.join(ctx.obj["output_dir"], "raw")) + cfg = ctx.obj["cfg"] + http_cfg = { + 'timeout': cfg.get('http', 'timeout', fallback='30'), + 'max_retries': cfg.get('http', 'max_retries', fallback='3'), + 'user_agent': cfg.get('http', 'user_agent', fallback='Mozilla/5.0 (compatible; research-crawler/1.0)'), + } + collector = CustomSource(profile, db, cache, config=http_cfg) + + since_date = since + until_date = None + if date_range: + since_date, until_date = date_range + + stats = collector.collect( + since=since_date, until=until_date, + limit=limit, update=update, + ) + + table = Table(title="Collection Results", box=box.ROUNDED) + table.add_column("Metric", style="bold") + table.add_column("Count", justify="right") + table.add_row("Discovered", str(stats.get('discovered', 0))) + table.add_row("Downloaded", str(stats.get('downloaded', 0))) + table.add_row("Already cached", str(stats.get('cached', 0))) + table.add_row("Failed", str(stats.get('failed', 0))) + if update: + table.add_row("Updated", str(stats.get('updated', 0))) + console.print(table) + finally: + db.close() # --------------------------------------------------------------------------- @@ -253,6 +302,9 @@ def parse(ctx, source, retry_failed): source = source or ctx.obj.get("source") retry_failed = retry_failed or ctx.obj.get("retry_failed", False) + from parsers.base import ParserManager + from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn + console.print(Panel( f"[bold green]Parse[/bold green]\n\n" f" Source : {source or 'all'}\n" @@ -260,9 +312,80 @@ def parse(ctx, source, retry_failed): title="[bold]mosaic parse[/bold]", border_style="green", )) - console.print( - "\n[dim]Parsing not yet implemented — see parsers/[/dim]" - ) + + db = open_db(ctx.obj["db_path"]) + try: + cfg = ctx.obj["cfg"] + parser = ParserManager( + memory_limit_mb=int(cfg.get('parser', 'subprocess_memory_limit_mb', fallback='512')), + timeout_seconds=int(cfg.get('parser', 'subprocess_timeout_seconds', fallback='60')), + min_chars_per_page=int(cfg.get('parser', 'min_chars_per_page', fallback='50')), + max_pdf_pages=int(cfg.get('parser', 'max_pdf_pages', fallback='200')), + ) + + # Get documents to parse + if source: + docs = db.get_documents_by_source(source, status='CACHED') + else: + docs = db.get_documents_by_source('', status='CACHED') + # Get all cached docs across sources + cursor = db.conn.execute( + "SELECT * FROM documents WHERE status = 'CACHED' AND parse_status IN ('PENDING', ?)", + ('PARSE_FAILED' if retry_failed else '__NONE__',) + ) + from storage.models import Document + docs = [Document.from_row(row) for row in cursor.fetchall()] + + if not docs: + console.print("[yellow]No documents to parse.[/yellow]") + return + + parsed = 0 + failed = 0 + needs_ocr = 0 + + with Progress( + SpinnerColumn(), TextColumn("{task.description}"), + BarColumn(), TaskProgressColumn(), + ) as progress: + task = progress.add_task(f"Parsing {len(docs)} docs", total=len(docs)) + for doc in docs: + if doc.parse_status in ('PARSED', 'NEEDS_OCR') and not retry_failed: + progress.advance(task) + continue + + db.update_document_status(doc.id, parse_status='IN_PROGRESS') + result = parser.parse(doc.raw_path, doc.doc_type, doc.content_hash) + + if result.success: + doc.text = result.text + doc.title = result.title or doc.title + doc.char_count = result.char_count + if result.metadata: + doc.metadata.update(result.metadata) + ps = 'NEEDS_OCR' if result.needs_ocr else 'PARSED' + doc.parse_status = ps + db.insert_document(doc) + db.update_document_status(doc.id, parse_status=ps) + if result.needs_ocr: + needs_ocr += 1 + else: + parsed += 1 + else: + db.update_document_status(doc.id, parse_status='PARSE_FAILED') + failed += 1 + + progress.advance(task) + + table = Table(title="Parse Results", box=box.ROUNDED) + table.add_column("Metric", style="bold") + table.add_column("Count", justify="right") + table.add_row("Parsed", str(parsed)) + table.add_row("Needs OCR", str(needs_ocr)) + table.add_row("Failed", str(failed)) + console.print(table) + finally: + db.close() # ---------------------------------------------------------------------------