main.py

# main.py
# Runs the full TraceOps pipeline: ingestion → indexing → graph → agent → evaluation.

import argparse
import logging
import os
import time

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
)
logger = logging.getLogger("traceops.main")


def parse_args():
    p = argparse.ArgumentParser(description="TraceOps — Hallucination-Aware Ops Agent")
    # Dataset selection
    p.add_argument("--dataset", default="HDFS_v1",
                   help="LogHub dataset name(s), comma-separated. e.g. HDFS_v1 or HDFS_v1,Hadoop")
    p.add_argument("--data-dir", default="data",
                   help="Directory to store/load LogHub datasets")
    p.add_argument("--download", action="store_true",
                   help="Download specified dataset(s) before running")
    p.add_argument("--list-datasets", action="store_true",
                   help="Print available datasets and exit")
    p.add_argument("--session-chunking", action="store_true", default=True,
                   help="For HDFS: chunk by block session instead of sliding window")
    # Legacy single-file paths (still work)
    p.add_argument("--log", default=None, help="(Legacy) direct path to a .log file")
    p.add_argument("--labels", default=None, help="(Legacy) direct path to anomaly_label.csv")
    # Run options
    p.add_argument("--demo", action="store_true", help="Run with synthetic data")
    p.add_argument("--query", type=str, default=None, help="Single query to run")
    p.add_argument("--eval-only", action="store_true", help="Run evaluation only")
    p.add_argument("--n-eval-queries", type=int, default=50)
    p.add_argument("--max-lines", type=int, default=100_000,
                   help="Max log lines to ingest (use 0 for all)")
    p.add_argument("--llm-provider", default="ollama",
                   choices=["ollama", "google", "gemini", "huggingface", "hf",
                            "openai", "anthropic", "claude", "lmstudio"],
                   help="LLM backend. ollama=local (default), google=Gemini, "
                        "lmstudio=local OpenAI-compatible server")
    p.add_argument("--llm-model", default="",
                   help="Model name. Defaults per provider if empty. "
                        "ollama: mistral | google: gemini-1.5-flash | "
                        "anthropic: claude-3-5-haiku-20241022")
    p.add_argument("--output-dir", default="outputs")

    # live monitoring args
    live = p.add_argument_group("live monitoring")
    live.add_argument(
        "--live", metavar="LOG_PATH",
        default=None,
        help="Path to a live/growing log file, or '-' to read from stdin.",
    )
    live.add_argument(
        "--live-mode",
        choices=["interactive", "auto", "alert-only"],
        default="interactive",
        help="interactive: type queries at runtime. "
             "auto: run --live-query list on new chunks. "
             "alert-only: trigger on ERROR/FATAL lines. (default: interactive)",
    )
    live.add_argument(
        "--live-format",
        choices=["auto", "hdfs", "syslog", "generic"],
        default="auto",
        help="Log format for live stream (default: auto)",
    )
    live.add_argument(
        "--live-interval", type=float, default=5.0, metavar="SECONDS",
        help="Polling interval in seconds (default: 5)",
    )
    live.add_argument(
        "--live-min-chunks", type=int, default=3, metavar="N",
        help="Minimum new chunks required before auto-queries fire (default: 3)",
    )
    live.add_argument(
        "--live-query", action="append", dest="live_queries", metavar="QUERY",
        help="Query for auto mode. Repeat to add multiple. Defaults to built-in HDFS queries.",
    )
    live.add_argument(
        "--live-output", metavar="FILE", default=None,
        help="Append all agent answers to this file in addition to stdout.",
    )
    live.add_argument(
        "--live-replay", action="store_true",
        help="Replay --live path at --live-replay-speed (demo without a live cluster).",
    )
    live.add_argument(
        "--live-replay-speed", type=float, default=100.0, metavar="LINES_PER_SEC",
        help="Lines per second for replay mode (default: 100)",
    )

    return p.parse_args()


def run_pipeline(args):
    os.makedirs(args.output_dir, exist_ok=True)

    # --- eval-only ---
    if args.eval_only:
        logger.info("Running synthetic evaluation...")
        from src.evaluation.evaluator import run_synthetic_evaluation
        ret, cal = run_synthetic_evaluation(
            n_queries=args.n_eval_queries,
            save_dir=args.output_dir,
        )
        print("\n=== Retrieval Metrics ===")
        print(ret.summary())
        print("\n=== Calibration Metrics ===")
        print(cal.summary())
        return

    logger.info("=== Stage 1: Ingestion ===")
    from src.ingestion.loghub_loader import LogHubLoader

    loader = LogHubLoader(data_dir=args.data_dir)

    if args.list_datasets:
        loader.print_dataset_catalog()
        return

    if args.download:
        for ds in args.dataset.split(","):
            ds = ds.strip()
            logger.info(f"Downloading {ds}...")
            loader.download(ds)

    if args.demo:
        documents = loader._synthetic_fallback(chunk_size=20, overlap=5, max_lines=5000)
    elif "," in args.dataset:
        dataset_names = [d.strip() for d in args.dataset.split(",")]
        documents = loader.load_multiple(dataset_names, max_lines_each=args.max_lines)
    else:
        documents = loader.load(
            args.dataset,
            chunk_size=20,
            overlap=5,
            max_lines=args.max_lines,
            session_chunking=args.session_chunking,
        )

    # Legacy --log / --labels override
    if args.log and os.path.exists(args.log):
        logger.info(f"Using legacy --log path: {args.log}")
        from src.ingestion.ingestor import chunk_logs
        documents = chunk_logs(
            args.log,
            chunk_size=20, overlap=5,
            label_path=args.labels if args.labels and os.path.exists(args.labels) else None,
            max_lines=args.max_lines or None,
        )
        logger.info(f"Ingested {len(documents)} chunks from legacy path")
    elif not args.demo:
        pass  # already handled above

    logger.info(f"Total document chunks: {len(documents)}")

    logger.info("=== Stage 2: Building Hybrid Index ===")
    from src.retrieval.hybrid_retriever import build_retriever
    retriever = build_retriever(
        documents=documents,
        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
        reranker_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
        enable_compression=True,
        compression_max_tokens=512,
        bm25_weight=0.4,
        dense_weight=0.6,
    )

    logger.info("=== Stage 3: Building Knowledge Graph ===")
    from src.graph.knowledge_graph import LogKnowledgeGraph
    kg = LogKnowledgeGraph()
    kg.build(documents)
    logger.info(f"KG summary: {kg.summary()}")

    # --- LLM ---
    logger.info(f"=== Stage 4: Loading LLM ({args.llm_provider}) ===")
    from src.agent.llm_wrapper import TraceOpsLLM
    model = args.llm_model or None   # None lets TraceOpsLLM apply provider defaults
    llm = TraceOpsLLM(
        provider=args.llm_provider,
        model=model,
        temperature_reason=0.3,
        temperature_consistency=0.7,
    )
    logger.info(f"LLM ready: {llm.provider}/{llm.model}")

    # BerryEnhancedScorer is a drop-in for ConfidenceScorer; falls back transparently.
    try:
        from src.confidence.berry_adapter import BerryEnhancedScorer
        scorer = BerryEnhancedScorer(
            weight_grounding=0.40,
            weight_consistency=0.35,
            weight_citation=0.25,
            threshold_suggest=0.75,
            threshold_review=0.40,
            embedder=retriever.faiss_index.model,
        )
        logger.info("Using BerryEnhancedScorer (strawberry integration)")
    except ImportError:
        from src.confidence.scorer import ConfidenceScorer
        scorer = ConfidenceScorer(
            weight_grounding=0.40,
            weight_consistency=0.35,
            weight_citation=0.25,
            threshold_suggest=0.75,
            threshold_review=0.40,
            embedder=retriever.faiss_index.model,
        )
        logger.info("Using ConfidenceScorer (berry_adapter import failed)")

    logger.info("=== Stage 5: Building LangGraph Agent ===")
    from src.agent.graph_agent import build_traceops_graph, TraceOpsAgent
    compiled_graph = build_traceops_graph(
        retriever=retriever,
        kg=kg,
        llm=llm,
        scorer=scorer,
        threshold_suggest=0.75,
        threshold_review=0.40,
        top_k=5,
    )
    agent = TraceOpsAgent(compiled_graph)

    if args.live:
        logger.info("=== Stage 6: Starting Live Monitor ===")
        from src.ingestion.live_stream import LogStreamWatcher, LogReplayWatcher
        from src.live_monitor import LiveMonitor

        WatcherClass = LogReplayWatcher if args.live_replay else LogStreamWatcher
        watcher_kwargs = dict(
            path=args.live,
            format=args.live_format,
            chunk_size=20,
            overlap=5,
        )
        if args.live_replay:
            watcher_kwargs["lines_per_second"] = args.live_replay_speed

        watcher = WatcherClass(**watcher_kwargs)

        monitor = LiveMonitor(
            agent=agent,
            retriever=retriever,
            kg=kg,
            watcher=watcher,
            auto_queries=args.live_queries or None,
            output_file=args.live_output,
        )
        monitor.run(
            mode=args.live_mode,
            poll_interval=args.live_interval,
            min_new_chunks=args.live_min_chunks,
        )
        return  # skip static queries and evaluation in live mode

    # --- Run queries ---
    queries = []
    if args.query:
        queries = [args.query]
    else:
        queries = [
            "Why is the DataNode failing to write blocks?",
            "What caused the IOException on blk_-1233456789?",
            "Are there replication errors affecting multiple nodes?",
            "Summarize recent WARN-level events from the NameSystem.",
        ]

    logger.info("=== Stage 6: Running Agent Queries ===")
    rate_limit_sleep = 15 if args.llm_provider in ("google", "gemini") else 0
    for i, q in enumerate(queries):
        if i > 0 and rate_limit_sleep:
            logger.info(f"Rate-limit pause: sleeping {rate_limit_sleep}s before next query...")
            time.sleep(rate_limit_sleep)
        print(f"\n{'='*70}")
        print(f"Query: {q}")
        print(f"{'='*70}")
        result_text = agent.run_and_explain(q)
        print(result_text)

    # --- Evaluation ---
    logger.info("=== Stage 7: Running Evaluation ===")
    from src.evaluation.evaluator import run_synthetic_evaluation
    ret, cal = run_synthetic_evaluation(
        n_queries=args.n_eval_queries,
        save_dir=args.output_dir,
    )
    print("\n=== Evaluation Results ===")
    print(ret.summary())
    print(cal.summary())
    print(f"\nPlots saved to {args.output_dir}/")


if __name__ == "__main__":
    args = parse_args()
    run_pipeline(args)