from __future__ import annotations import argparse import json import subprocess import sys import time from pathlib import Path from app.processing import process_raw_data_with_summary from app.qa_suggestions import export_suggested_questions_json from app.vector_store import index_chunks PROJECT_ROOT = Path(__file__).resolve().parent.parent class PipelineConsole: def __init__(self) -> None: self._started_at = time.perf_counter() def intro(self, command: str) -> None: print("=" * 72) print("ChatVNS pipeline") print(f"Command: {command}") print("Flow: collect raw data -> process/chunk -> embed/index Qdrant -> Streamlit") print("=" * 72) def stage(self, number: int, total: int, title: str, detail: str | None = None) -> float: print() print(f"[{number}/{total}] {title}") if detail: print(f" {detail}") return time.perf_counter() def done(self, started_at: float, detail: str | None = None) -> None: duration = time.perf_counter() - started_at suffix = f" | {detail}" if detail else "" print(f" DONE in {duration:.1f}s{suffix}") def note(self, message: str) -> None: print(f" {message}") def finish(self) -> None: total = time.perf_counter() - self._started_at print() print(f"Pipeline ready in {total:.1f}s") def format_counts(values: dict[str, int]) -> str: if not values: return "none" return ", ".join(f"{key}={value}" for key, value in values.items()) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run ChatVNS pipeline commands.") subparsers = parser.add_subparsers(dest="command", required=True) collect = subparsers.add_parser("collect", help="Collect raw data into data/raw.") add_collect_args(collect) index = subparsers.add_parser("index", help="Process raw data and index chunks into Qdrant.") index.add_argument("--skip-qdrant", action="store_true") chat = subparsers.add_parser("chat", help="Open the Streamlit chatbot.") chat.add_argument("--server-port", default="8501") all_cmd = subparsers.add_parser("all", help="Collect, index, then open Streamlit.") add_collect_args(all_cmd) all_cmd.add_argument("--server-port", default="8501") all_cmd.add_argument("--skip-qdrant", action="store_true") return parser.parse_args() def add_collect_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--tickers", nargs="+", default=["HPG", "FPT", "VCB"]) parser.add_argument("--include-news", action="store_true") parser.add_argument("--include-charts", action="store_true") parser.add_argument("--max-reports", type=int, default=5) parser.add_argument("--delay-seconds", type=float, default=1.5) parser.add_argument("--timeout", type=int, default=30) parser.add_argument("--min-crawl-interval-minutes", type=float, default=30) def run_collect(args: argparse.Namespace) -> None: command = [ sys.executable, str(PROJECT_ROOT / "bot-collect-data" / "crawl_raw_data.py"), "--tickers", *args.tickers, "--max-reports", str(args.max_reports), "--delay-seconds", str(args.delay_seconds), "--timeout", str(args.timeout), "--min-crawl-interval-minutes", str(args.min_crawl_interval_minutes), ] if args.include_news: command.append("--include-news") if args.include_charts: command.append("--include-charts") subprocess.run(command, cwd=PROJECT_ROOT, check=True) def ensure_qdrant() -> None: subprocess.run( ["docker", "compose", "up", "-d", "qdrant"], cwd=PROJECT_ROOT, check=True, ) def run_index() -> None: suggestions = export_suggested_questions_json() print(f" Synced Q&A suggestions: groups={len(suggestions)}") chunks, summary = process_raw_data_with_summary() print( " Processed raw data: " f"documents={summary['document_count']}, " f"chunks={summary['chunk_count']}, " f"tickers={', '.join(summary['tickers']) or 'none'}" ) print(f" Modalities: {format_counts(summary['modalities'])}") print(f" Chunks by ticker: {format_counts(summary['chunks_by_ticker'])}") print(f" Chunks by scope: {format_counts(summary['chunks_by_scope'])}") print(f" Structures: {format_counts(summary['structures'])}") def on_index_progress(event: dict) -> None: action = "Embedding" if event["stage"] == "embedding" else "Upserted" cache_detail = "" if event["stage"] == "embedding" and "cache_hits" in event: cache_detail = f", cache_hits={event['cache_hits']}, api_chunks={event['cache_misses']}" print( " " f"{action} batch {event['batch_number']}/{event['total_batches']} " f"({event['batch_size']} chunks, indexed={event['indexed_so_far']}/{event['total_chunks']}" f"{cache_detail})" ) indexed_count = index_chunks(chunks, progress_callback=on_index_progress) print( json.dumps( { "chunk_count": len(chunks), "indexed_count": indexed_count, "documents_by_ticker": summary["documents_by_ticker"], "chunks_by_ticker": summary["chunks_by_ticker"], "documents_by_scope": summary["documents_by_scope"], "chunks_by_scope": summary["chunks_by_scope"], }, ensure_ascii=False, indent=2, ) ) def run_streamlit(server_port: str) -> None: subprocess.run( [ sys.executable, "-m", "streamlit", "run", str(PROJECT_ROOT / "app" / "streamlit_app.py"), "--server.port", server_port, ], cwd=PROJECT_ROOT, check=True, ) def main() -> int: args = parse_args() console = PipelineConsole() console.intro(args.command) if args.command == "collect": stage = console.stage( 1, 1, "Collect raw data", f"tickers={', '.join(args.tickers)} | news={args.include_news} | charts={args.include_charts}", ) run_collect(args) console.done(stage) console.finish() return 0 if args.command == "index": if not args.skip_qdrant: stage = console.stage(1, 2, "Start Qdrant", "docker compose up -d qdrant") ensure_qdrant() console.done(stage) stage = console.stage(2 if not args.skip_qdrant else 1, 2 if not args.skip_qdrant else 1, "Process and index", "parse raw files -> chunk -> embed -> upsert") run_index() console.done(stage) console.finish() return 0 if args.command == "chat": stage = console.stage(1, 1, "Open Streamlit", f"http://localhost:{args.server_port}") console.note("Launching Streamlit app. Keep this terminal open while using the UI.") console.done(stage, f"url=http://localhost:{args.server_port}") run_streamlit(args.server_port) return 0 if args.command == "all": if not args.skip_qdrant: stage = console.stage(1, 4, "Start Qdrant", "docker compose up -d qdrant") ensure_qdrant() console.done(stage) collect_stage = console.stage( 2 if not args.skip_qdrant else 1, 4 if not args.skip_qdrant else 3, "Collect raw data", f"tickers={', '.join(args.tickers)} | news={args.include_news} | charts={args.include_charts}", ) run_collect(args) console.done(collect_stage) index_stage = console.stage( 3 if not args.skip_qdrant else 2, 4 if not args.skip_qdrant else 3, "Process and index", "parse raw files -> chunk -> embed -> upsert", ) run_index() console.done(index_stage) streamlit_stage = console.stage( 4 if not args.skip_qdrant else 3, 4 if not args.skip_qdrant else 3, "Open Streamlit", f"http://localhost:{args.server_port}", ) console.note("Launching Streamlit app. Keep this terminal open while using the UI.") console.done(streamlit_stage, f"url=http://localhost:{args.server_port}") console.finish() run_streamlit(args.server_port) return 0 return 1 if __name__ == "__main__": raise SystemExit(main())