src/kbdebugger/vector_main.py

from __future__ import annotations

"""
Minimal entry point to test the Vector Similarity Filter component in isolation.

What this script validates
--------------------------
This script sanity-checks the end-to-end "Vector Similarity Filter" MVP:

1) Retrieve a keyword-guided subgraph from Neo4j (Graph Retriever).
2) Build a vector index over the KG subgraph relation sentences (r.sentence).
3) Load corpus (raw documents) and run only the Decomposer LLM to produce candidate qualities.
4) Filter qualities by vector similarity against the KG subgraph index.
5) Pretty-print kept/dropped results and optionally write them to JSON logs.

Important design decision
-------------------------
We intentionally DO NOT run triplet extraction here.

Triplet extraction is an expensive LLM step and is deferred until after
vector filtering. This script tests the cheaper filtering stage only.

Configuration
-------------
Driven by environment variables (see VectorMainConfig.from_env).

Typical usage:
    KB_KEYWORD=requirement KB_SOURCE_KIND=TEXT python -m kbdebugger.subgraph_similarity_main
"""

import os
from dataclasses import dataclass
from typing import cast

from rich.console import Console

from kbdebugger.extraction.types import SourceKind

from kbdebugger.extraction import chunk_corpus, decompose_documents
from kbdebugger.extraction.types import Qualities

from kbdebugger.graph.retriever import KnowledgeGraphRetriever

from kbdebugger.subgraph_similarity.encoder import SentenceTransformerEncoder
from kbdebugger.subgraph_similarity.similarity_filter import SubgraphSimilarityFilter

console = Console()


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class VectorMainConfig:
    """
    Runtime configuration for vector_main.py.

    Environment variables
    ---------------------
    KB_KEYWORD:
        Keyword to retrieve a KG subgraph (Graph Retriever input).
        Default: "requirement"

    KB_SOURCE_KIND:
        One of:
            - "TEXT"
            - "PDF_SENTENCES"
            - "PDF_CHUNKS"
        Default: "TEXT"

    KB_TEXT_PATH:
        Text file path used when KB_SOURCE_KIND == "TEXT"
        Default: "data/DSA/DSA_knowledge.txt"

    KB_PDF_PATH:
        PDF path used when KB_SOURCE_KIND starts with "PDF"
        Default: "data/SDS/InstructCIR.pdf"

    KB_LIMIT_PER_PATTERN:
        How many KG relations to retrieve per pattern in KG retriever.
        Default: 50

    KB_TOP_K:
        How many nearest KG relations to retrieve per candidate quality.
        Default: 5

    KB_THRESHOLD:
        Minimum cosine similarity threshold.
        Default: 0.55

    KB_LOG_PATH:
        Optional JSON output path.
        Default: "logs/vector_similarity_results.json"
        Set empty to disable logging.
    """
    keyword: str
    source_kind: SourceKind
    text_path: str
    pdf_path: str

    limit_per_pattern: int
    top_k: int
    threshold: float

    log_path: str | None

    @classmethod
    def from_env(cls) -> VectorMainConfig:
        keyword = os.getenv("KB_KEYWORD", "requirement").strip()

        source_raw = os.getenv("KB_SOURCE_KIND", "TEXT").upper().strip()
        if source_raw not in {"TEXT", "PDF_SENTENCES", "PDF_CHUNKS"}:
            raise ValueError(f"Invalid KB_SOURCE_KIND={source_raw!r}")
        source_raw = cast(SourceKind, source_raw)

        text_path = os.getenv("KB_TEXT_PATH", "data/DSA/DSA_knowledge.txt").strip()
        pdf_path = os.getenv("KB_PDF_PATH", "data/SDS/Handout for inspection of examinations.pdf").strip()

        limit_per_pattern = int(os.getenv("KB_LIMIT_PER_PATTERN", "50").strip())
        top_k = int(os.getenv("KB_TOP_K", "5").strip())
        threshold = float(os.getenv("KB_THRESHOLD", "0.55").strip())

        log_raw = os.getenv("KB_LOG_PATH", "logs/vector_similarity_results.json").strip()
        log_path = log_raw or None

        return cls(
            keyword=keyword,
            source_kind=source_raw,
            text_path=text_path,
            pdf_path=pdf_path,
            limit_per_pattern=max(1, limit_per_pattern),
            top_k=max(1, top_k),
            threshold=threshold,
            log_path=log_path,
        )

# ---------------------------------------------------------------------------
def run_extractor(cfg: VectorMainConfig) -> Qualities:
    """
    Produce candidate qualities using ONLY the Decomposer module.

    This is intentionally cheaper than the full extractor pipeline and matches
    our improved design (triplet extraction happens later, only for kept qualities).
    """
    # 1. Load corpus & chunk into LangChain Documents
    docs, mode = chunk_corpus(
        source_kind=cfg.source_kind,
        path=cfg.text_path if cfg.source_kind == SourceKind.TEXT else cfg.pdf_path,
    )

    # 2. Decompose each document into qualities
    qualities, decomposer_log = decompose_documents(
        docs=docs,
        mode=mode,
    )

    if not qualities:
        raise ValueError("Decomposition produced no qualities.")

    return qualities


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> None:
    cfg = VectorMainConfig.from_env()
    console.rule("[bold green]KBDEBUGGER Vector Similarity Filter")

    console.print(f"[bold]Keyword:[/bold] [cyan]{cfg.keyword}[/cyan]")
    console.print(f"[bold]Source kind:[/bold] [cyan]{cfg.source_kind}[/cyan]")
    console.print(f"[bold]Limit per retriever pattern:[/bold] [yellow]{cfg.limit_per_pattern}[/yellow]")
    console.print(f"[bold]top_k:[/bold] [yellow]{cfg.top_k}[/yellow]")
    console.print(f"[bold]threshold:[/bold] [yellow]{cfg.threshold}[/yellow]")
    if cfg.log_path:
        console.print(f"[bold]Log path:[/bold] [green]{cfg.log_path}[/green]")
    console.print()

    # 1) Retrieve KG subgraph relations around the keyword
    retriever = KnowledgeGraphRetriever(limit_per_pattern=cfg.limit_per_pattern)
    hits = retriever.retrieve(cfg.keyword)
    relations = [h["relation"] for h in hits]

    if not relations:
        console.print("[bold yellow]No KG relations retrieved. Try a different keyword.[/bold yellow]")
        return

    console.print(f"[bold]Retrieved KG relations:[/bold] [cyan]{len(relations)}[/cyan]")

    # 2) Produce candidate qualities (decomposer output)
    qualities = run_extractor(cfg)
    console.print(f"[bold]Candidate qualities (from decomposer):[/bold] [cyan]{len(qualities)}[/cyan]\n")

    # 3) Create encoder + filter component
    encoder = SentenceTransformerEncoder(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        device=None,          # let sentence-transformers choose
        normalize=True,       # recommended for cosine similarity
    )

    filt = SubgraphSimilarityFilter(
        encoder=encoder,
        top_k=cfg.top_k,
        threshold=cfg.threshold,
    )

    # 4) Build index over KG subgraph relation sentences
    index = filt.build_index(relations) # Here we internally encode the KG relations & build the index

    # 5) Filter qualities
    kept, dropped = filt.filter_qualities(index=index, qualities=qualities)

    # 6) Print results
    filt.pretty_print(kept=kept, dropped=dropped)

if __name__ == "__main__":
    main()