query.py

#!/usr/bin/env python3
"""
RAG query CLI tool for TextStruct.

This tool allows interactive querying of indexed textbook content with
metadata-aware filtering by pedagogical role, chapter, and section.

Usage:
    python src/query.py <vector_store_path> [--query "your question"]
    python src/query.py <vector_store_path> --interactive
"""

import argparse
import json
from pathlib import Path
from typing import Optional

from src.rag_retriever import load_retriever, format_results
from src.embeddings import EmbeddingEngine


PEDAGOGICAL_ROLES = [
    "definition",
    "explanation",
    "procedure",
    "example",
    "question",
    "result",
    "reference",
    "evidence",
    "visual_proxy",
    "context",
]


def print_separator():
    """Print a visual separator."""
    print("\n" + "=" * 80 + "\n")


def query_interactive(retriever, embedding_engine):
    """
    Interactive query mode with role-based filtering.

    Args:
        retriever: RAGRetriever instance
        embedding_engine: EmbeddingEngine instance
    """
    print("=" * 80)
    print("TextStruct RAG Interactive Query")
    print("=" * 80)

    # Show stats
    stats = retriever.get_stats()
    print(f"\nIndexed chunks: {stats['num_chunks']}")

    if "chunks_by_role" in stats:
        print("\nChunks by pedagogical role:")
        for role, count in sorted(stats["chunks_by_role"].items(), key=lambda x: -x[1])[:5]:
            pct = (count / stats["num_chunks"]) * 100
            print(f"  {role}: {count} ({pct:.1f}%)")

    print("\nCommands:")
    print("  - Type your query and press Enter")
    print("  - Use 'role:<role_name>' to filter by pedagogical role")
    print("  - Use 'section:<section_name>' to filter by section")
    print("  - Use 'top:<n>' to get top n results (default: 5)")
    print("  - Use 'sql:<WHERE clause>' for advanced SQL filtering (LanceDB only)")
    print("  - Type 'help' to see available roles")
    print("  - Type 'stats' to see index statistics")
    print("  - Type 'exit' or 'quit' to exit")

    print_separator()

    while True:
        try:
            # Get user input
            query = input("Query: ").strip()

            if not query:
                continue

            # Handle special commands
            if query.lower() in ["exit", "quit", "q"]:
                print("Goodbye!")
                break

            if query.lower() == "help":
                print("\nAvailable pedagogical roles:")
                for role in PEDAGOGICAL_ROLES:
                    print(f"  - {role}")
                print_separator()
                continue

            if query.lower() == "stats":
                stats = retriever.get_stats()
                print(f"\nTotal chunks: {stats['num_chunks']}")
                print(f"Embedding dimension: {stats['embedding_dim']}")

                if "chunks_by_role" in stats:
                    print("\nChunks by pedagogical role:")
                    for role, count in sorted(stats["chunks_by_role"].items(), key=lambda x: -x[1]):
                        pct = (count / stats["num_chunks"]) * 100
                        print(f"  {role}: {count} ({pct:.1f}%)")

                if "chunks_by_section" in stats:
                    print(f"\nUnique sections: {len(stats['chunks_by_section'])}")

                print_separator()
                continue

            # Parse query modifiers
            role_filter = None
            section_filter = None
            sql_filter = None
            top_k = 5
            query_text = query

            # Extract modifiers
            parts = query.split()
            remaining_parts = []

            for part in parts:
                if part.startswith("role:"):
                    role_filter = part[5:]
                elif part.startswith("section:"):
                    section_filter = part[8:]
                elif part.startswith("sql:"):
                    # SQL filter can span multiple words, so grab rest of query
                    sql_idx = query.index("sql:")
                    sql_filter = query[sql_idx + 4:].strip()
                    # Remove SQL part from query
                    query_text = query[:sql_idx].strip()
                    break
                elif part.startswith("top:"):
                    try:
                        top_k = int(part[4:])
                    except ValueError:
                        print(f"[WARN] Invalid top value: {part[4:]}, using default (5)")
                else:
                    remaining_parts.append(part)

            if not sql_filter:
                query_text = " ".join(remaining_parts)

            if not query_text:
                print("[ERROR] Please provide a query text")
                continue

            # Execute query
            print(f"\nSearching for: '{query_text}'")
            if role_filter:
                print(f"Filtered by role: {role_filter}")
            if section_filter:
                print(f"Filtered by section: {section_filter}")
            if sql_filter:
                print(f"SQL filter: {sql_filter}")
            print(f"Top {top_k} results:\n")

            # Use SQL query if available, otherwise use regular query
            if sql_filter and hasattr(retriever, 'query_sql'):
                results = retriever.query_sql(
                    query_text=query_text,
                    top_k=top_k,
                    sql_filter=sql_filter,
                )
            else:
                if sql_filter:
                    print("[WARN] SQL filters not supported by current vector store, using dict filters")
                results = retriever.query(
                    query_text=query_text,
                    top_k=top_k,
                    pedagogical_role=role_filter,
                    section=section_filter,
                )

            # Display results
            if not results:
                print("No results found.")
            else:
                print(format_results(results, max_text_length=300))

            print_separator()

        except KeyboardInterrupt:
            print("\n\nGoodbye!")
            break
        except Exception as e:
            print(f"\n[ERROR] {e}")
            print_separator()


def query_single(
    retriever,
    query_text: str,
    top_k: int = 5,
    role: Optional[str] = None,
    section: Optional[str] = None,
    sql_filter: Optional[str] = None,
    output_json: Optional[Path] = None,
):
    """
    Execute a single query and print results.

    Args:
        retriever: RAGRetriever instance
        query_text: Query text
        top_k: Number of results
        role: Optional role filter
        section: Optional section filter
        sql_filter: Optional SQL WHERE clause (LanceDB only)
        output_json: Optional path to save results as JSON
    """
    print(f"Query: {query_text}")
    if role:
        print(f"Role filter: {role}")
    if section:
        print(f"Section filter: {section}")
    if sql_filter:
        print(f"SQL filter: {sql_filter}")
    print()

    # Use SQL query if available, otherwise use regular query
    if sql_filter and hasattr(retriever, 'query_sql'):
        results = retriever.query_sql(
            query_text=query_text,
            top_k=top_k,
            sql_filter=sql_filter,
        )
    else:
        if sql_filter:
            print("[WARN] SQL filters not supported by current vector store, using dict filters")
        results = retriever.query(
            query_text=query_text,
            top_k=top_k,
            pedagogical_role=role,
            section=section,
        )

    if not results:
        print("No results found.")
        return

    # Print results
    print(format_results(results, max_text_length=300))

    # Save to JSON if requested
    if output_json:
        output_data = []
        for chunk, score in results:
            output_data.append({
                "chunk": chunk,
                "similarity_score": score,
            })

        output_json.parent.mkdir(parents=True, exist_ok=True)
        with open(output_json, "w", encoding="utf-8") as f:
            json.dump(output_data, f, indent=2, ensure_ascii=False)

        print(f"\nResults saved to {output_json}")


def main():
    parser = argparse.ArgumentParser(
        description="Query indexed textbook content using RAG",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Interactive mode
  python src/query.py data/output/book1/vector_store.lancedb --interactive

  # Single query
  python src/query.py data/output/book1/vector_store.lancedb --query "What is a triangle?"

  # Query with role filter
  python src/query.py data/output/book1/vector_store.lancedb \\
    --query "triangle" --role definition

  # Query with SQL filter (LanceDB only)
  python src/query.py data/output/book1/vector_store.lancedb \\
    --query "triangles" --sql-filter "page_start >= 5 AND page_start <= 10"

  # Query with multiple filters
  python src/query.py data/output/book1/vector_store.lancedb \\
    --query "solve equations" --role procedure --top-k 3
        """
    )

    parser.add_argument(
        "vector_store",
        type=Path,
        help="Path to vector store file (.pkl or .lancedb)",
    )

    parser.add_argument(
        "--query",
        "-q",
        type=str,
        help="Query text (use --interactive for interactive mode)",
    )

    parser.add_argument(
        "--interactive",
        "-i",
        action="store_true",
        help="Start interactive query mode",
    )

    parser.add_argument(
        "--top-k",
        "-k",
        type=int,
        default=5,
        help="Number of results to return (default: 5)",
    )

    parser.add_argument(
        "--role",
        "-r",
        type=str,
        choices=PEDAGOGICAL_ROLES,
        help="Filter by pedagogical role",
    )

    parser.add_argument(
        "--section",
        "-s",
        type=str,
        help="Filter by section name",
    )

    parser.add_argument(
        "--sql-filter",
        type=str,
        help="SQL WHERE clause for advanced filtering (LanceDB only, e.g., 'page_start >= 5 AND page_start <= 10')",
    )

    parser.add_argument(
        "--output",
        "-o",
        type=Path,
        help="Save results to JSON file",
    )

    args = parser.parse_args()

    # Validate vector store path
    if not args.vector_store.exists():
        print(f"[ERROR] Vector store not found: {args.vector_store}")
        print("\nTo build a vector store, run:")
        print("  python src/main.py <pdf> --refine --semantic-chunk --classify-roles --build-index")
        return

    # Load retriever
    print(f"[INFO] Loading vector store from {args.vector_store}...")

    # Set up embedding engine with cache
    cache_dir = args.vector_store.parent.parent / ".doctr_cache" / "embeddings"
    embedding_engine = EmbeddingEngine(cache_dir=cache_dir)

    retriever = load_retriever(args.vector_store, embedding_engine)
    print(f"[INFO] Loaded {retriever.get_stats()['num_chunks']} chunks")
    print()

    # Execute query
    if args.interactive:
        query_interactive(retriever, embedding_engine)
    elif args.query:
        query_single(
            retriever,
            query_text=args.query,
            top_k=args.top_k,
            role=args.role,
            section=args.section,
            sql_filter=args.sql_filter,
            output_json=args.output,
        )
    else:
        parser.print_help()
        print("\n[ERROR] Please provide --query or --interactive")


if __name__ == "__main__":
    main()