main.py

"""
Cybersecurity Book Encoder - Main Entry Point

This script provides a complete pipeline for encoding an 800-page cybersecurity
book into a vector database that any AI agent can seamlessly access.

Features:
- Semantic chunking with parent-document retrieval
- High-dimensional dense embeddings (BGE-large)
- Hybrid search (dense + sparse with RRF fusion)
- Vector arithmetic for concept algebra
- Metadata-filtered search (category, severity, CVE)

Usage:
    # Encode a book
    python main.py encode path/to/book.pdf
    
    # Query the encoded book
    python main.py query "How do buffer overflow attacks work?"
    
    # Concept algebra query
    python main.py concept "Zero Trust + Cloud Architecture"
    
    # Analogy query
    python main.py analogy "SQL Injection" "Web Application" "Memory Corruption"
"""

import sys
import argparse
from pathlib import Path
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown

console = Console()


def encode_command(args):
    """Encode a book into the vector database."""
    from pipeline import encode_book
    from config import PipelineConfig, EmbeddingConfig, VectorDBConfig, ChunkingConfig
    
    # Build config from args
    config = PipelineConfig(
        chunking=ChunkingConfig(
            chunk_size=args.chunk_size,
            parent_chunk_size=args.parent_chunk_size,
            overlap=args.overlap
        ),
        embedding=EmbeddingConfig(
            model_name=args.model,
            device=args.device
        ),
        vector_db=VectorDBConfig(
            collection_name=args.collection,
            persist_directory=args.db_path
        ),
        use_hyperbolic_embeddings=args.hyperbolic
    )
    
    encode_book(args.book_path, config)


def query_command(args):
    """Query the encoded book."""
    from agent_interface import create_agent
    
    agent = create_agent(args.db_path, args.collection)
    
    console.print(f"\n[bold cyan]Query:[/bold cyan] {args.question}\n")
    
    response = agent.query(
        question=args.question,
        n_results=args.n_results,
        use_hybrid=not args.dense_only
    )
    
    # Display results
    console.print(Panel(
        f"[green]Found {len(response.results)} relevant passages[/green]",
        title="Search Results"
    ))
    
    for i, result in enumerate(response.results, 1):
        console.print(f"\n[bold]Result {i}[/bold] (Score: {result.score:.3f})")
        if result.metadata.get('chapter'):
            console.print(f"  Chapter: [yellow]{result.metadata['chapter']}[/yellow]")
        if result.metadata.get('severity'):
            console.print(f"  Severity: [red]{result.metadata['severity']}[/red]")
        console.print(f"\n{result.content[:500]}...")
        
        if result.parent_content and args.show_context:
            console.print(f"\n[dim]Full context available ({len(result.parent_content)} chars)[/dim]")
    
    # Show metadata summary
    console.print(Panel(
        f"Categories: {', '.join(response.metadata_summary['categories']) or 'N/A'}\n"
        f"CVEs mentioned: {', '.join(response.metadata_summary['cves']) or 'None'}",
        title="Metadata Summary"
    ))


def concept_command(args):
    """Perform vector arithmetic concept search."""
    from agent_interface import create_agent
    
    agent = create_agent(args.db_path, args.collection)
    
    console.print(f"\n[bold cyan]Concept Expression:[/bold cyan] {args.expression}\n")
    console.print("[dim]Computing vector arithmetic...[/dim]\n")
    
    response = agent.concept_search(
        expression=args.expression,
        n_results=args.n_results
    )
    
    console.print(Panel(
        f"[green]Found {len(response.results)} semantically related passages[/green]",
        title="Concept Algebra Results"
    ))
    
    for i, result in enumerate(response.results, 1):
        console.print(f"\n[bold]Result {i}[/bold] (Score: {result.score:.3f})")
        console.print(f"{result.content[:400]}...")


def analogy_command(args):
    """Perform analogical reasoning search."""
    from agent_interface import create_agent
    
    agent = create_agent(args.db_path, args.collection)
    
    console.print(f"\n[bold cyan]Analogy:[/bold cyan] {args.a} → {args.b} :: {args.c} → ?")
    console.print("[dim]Computing analogy vector...[/dim]\n")
    
    response = agent.analogy_search(
        a=args.a,
        b=args.b,
        c=args.c,
        n_results=args.n_results
    )
    
    console.print(Panel(
        f"[green]The analogical target relates to:[/green]",
        title="Analogy Results"
    ))
    
    for i, result in enumerate(response.results, 1):
        console.print(f"\n[bold]Result {i}[/bold] (Score: {result.score:.3f})")
        console.print(f"{result.content[:400]}...")


def stats_command(args):
    """Show statistics about the encoded book."""
    from agent_interface import create_agent
    
    agent = create_agent(args.db_path, args.collection)
    stats = agent.get_stats()
    
    console.print(Panel(
        f"Collection: [yellow]{stats['collection_name']}[/yellow]\n"
        f"Total Chunks: [green]{stats['total_chunks']}[/green]\n"
        f"Embedding Model: [cyan]{stats['embedding_model']}[/cyan]\n"
        f"Vector Dimension: [magenta]{stats['embedding_dimension']}[/magenta]",
        title="Vector Database Statistics"
    ))


def interactive_command(args):
    """Start interactive query session."""
    from agent_interface import create_agent
    
    agent = create_agent(args.db_path, args.collection)
    
    console.print(Panel(
        "Enter queries to search the cybersecurity book.\n"
        "Commands:\n"
        "  /concept <expression>  - Vector arithmetic\n"
        "  /category <cat> <query> - Category filter\n"
        "  /severity <sev> <query> - Severity filter\n"
        "  /quit - Exit",
        title="Interactive Mode"
    ))
    
    while True:
        try:
            query = console.input("\n[bold cyan]Query>[/bold cyan] ").strip()
            
            if not query:
                continue
            
            if query.lower() in ['/quit', '/exit', '/q']:
                console.print("Goodbye!")
                break
            
            if query.startswith('/concept '):
                expression = query[9:]
                response = agent.concept_search(expression, n_results=3)
            elif query.startswith('/category '):
                parts = query[10:].split(' ', 1)
                if len(parts) == 2:
                    response = agent.search_by_category(parts[1], parts[0], n_results=3)
                else:
                    console.print("[red]Usage: /category <category> <query>[/red]")
                    continue
            elif query.startswith('/severity '):
                parts = query[10:].split(' ', 1)
                if len(parts) == 2:
                    response = agent.search_by_severity(parts[1], parts[0], n_results=3)
                else:
                    console.print("[red]Usage: /severity <level> <query>[/red]")
                    continue
            else:
                response = agent.query(query, n_results=3)
            
            # Display results
            for i, result in enumerate(response.results, 1):
                console.print(f"\n[bold]─── Result {i} ───[/bold] (Score: {result.score:.3f})")
                console.print(result.content[:300] + "...")
            
        except KeyboardInterrupt:
            console.print("\nGoodbye!")
            break
        except Exception as e:
            console.print(f"[red]Error: {e}[/red]")


def main():
    parser = argparse.ArgumentParser(
        description="Cybersecurity Book Encoder - Seamless AI access to book knowledge"
    )
    
    subparsers = parser.add_subparsers(dest='command', help='Commands')
    
    # Encode command
    encode_parser = subparsers.add_parser('encode', help='Encode a book into vector database')
    encode_parser.add_argument('book_path', help='Path to the book (PDF or text)')
    encode_parser.add_argument('--db-path', default='./vector_db', help='Vector database path')
    encode_parser.add_argument('--collection', default='cybersecurity_book', help='Collection name')
    encode_parser.add_argument('--model', default='BAAI/bge-large-en-v1.5', help='Embedding model')
    encode_parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)')
    encode_parser.add_argument('--chunk-size', type=int, default=512, help='Child chunk size')
    encode_parser.add_argument('--parent-chunk-size', type=int, default=2048, help='Parent chunk size')
    encode_parser.add_argument('--overlap', type=float, default=0.25, help='Chunk overlap ratio')
    encode_parser.add_argument('--hyperbolic', action='store_true', help='Use hyperbolic embeddings')
    
    # Query command
    query_parser = subparsers.add_parser('query', help='Query the encoded book')
    query_parser.add_argument('question', help='Natural language question')
    query_parser.add_argument('--db-path', default='./vector_db', help='Vector database path')
    query_parser.add_argument('--collection', default='cybersecurity_book', help='Collection name')
    query_parser.add_argument('--n-results', type=int, default=5, help='Number of results')
    query_parser.add_argument('--dense-only', action='store_true', help='Skip hybrid search')
    query_parser.add_argument('--show-context', action='store_true', help='Show parent context')
    
    # Concept command
    concept_parser = subparsers.add_parser('concept', help='Vector arithmetic search')
    concept_parser.add_argument('expression', help='Concept expression (e.g., "Zero Trust + Cloud")')
    concept_parser.add_argument('--db-path', default='./vector_db', help='Vector database path')
    concept_parser.add_argument('--collection', default='cybersecurity_book', help='Collection name')
    concept_parser.add_argument('--n-results', type=int, default=5, help='Number of results')
    
    # Analogy command
    analogy_parser = subparsers.add_parser('analogy', help='Analogical reasoning search')
    analogy_parser.add_argument('a', help='First term (A)')
    analogy_parser.add_argument('b', help='Second term (B)')
    analogy_parser.add_argument('c', help='Third term (C)')
    analogy_parser.add_argument('--db-path', default='./vector_db', help='Vector database path')
    analogy_parser.add_argument('--collection', default='cybersecurity_book', help='Collection name')
    analogy_parser.add_argument('--n-results', type=int, default=5, help='Number of results')
    
    # Stats command
    stats_parser = subparsers.add_parser('stats', help='Show database statistics')
    stats_parser.add_argument('--db-path', default='./vector_db', help='Vector database path')
    stats_parser.add_argument('--collection', default='cybersecurity_book', help='Collection name')
    
    # Interactive command
    interactive_parser = subparsers.add_parser('interactive', help='Interactive query mode')
    interactive_parser.add_argument('--db-path', default='./vector_db', help='Vector database path')
    interactive_parser.add_argument('--collection', default='cybersecurity_book', help='Collection name')
    
    args = parser.parse_args()
    
    if args.command == 'encode':
        encode_command(args)
    elif args.command == 'query':
        query_command(args)
    elif args.command == 'concept':
        concept_command(args)
    elif args.command == 'analogy':
        analogy_command(args)
    elif args.command == 'stats':
        stats_command(args)
    elif args.command == 'interactive':
        interactive_command(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()