agent_interface.py

"""AI Agent interface for seamless book access."""
import numpy as np
from typing import List, Dict, Any, Optional, Union
from dataclasses import dataclass

from embeddings import EmbeddingEngine, VectorArithmetic
from vector_store import VectorStore, SearchResult
from config import DEFAULT_CONFIG


@dataclass
class AgentResponse:
    """Structured response for AI agent queries."""
    query: str
    results: List[SearchResult]
    context: str  # Combined context for the agent
    metadata_summary: Dict[str, Any]


class CybersecurityBookAgent:
    """
    High-level interface for AI agents to seamlessly access book knowledge.
    
    Capabilities:
    - Natural language queries
    - Vector arithmetic (concept algebra)
    - Category-filtered search
    - Severity-filtered search
    - Analogical reasoning
    """
    
    def __init__(
        self,
        vector_store: Optional[VectorStore] = None,
        embedding_engine: Optional[EmbeddingEngine] = None,
        config=DEFAULT_CONFIG
    ):
        self.config = config
        
        # Initialize or use provided components
        self._vector_store = vector_store
        self._embedding_engine = embedding_engine
        self._vector_arithmetic = None
    
    @property
    def vector_store(self) -> VectorStore:
        if self._vector_store is None:
            self._vector_store = VectorStore(
                collection_name=self.config.vector_db.collection_name,
                persist_directory=self.config.vector_db.persist_directory
            )
        return self._vector_store
    
    @property
    def embedding_engine(self) -> EmbeddingEngine:
        if self._embedding_engine is None:
            self._embedding_engine = EmbeddingEngine(
                model_name=self.config.embedding.model_name,
                device=self.config.embedding.device
            )
        return self._embedding_engine
    
    @property
    def vector_arithmetic(self) -> VectorArithmetic:
        if self._vector_arithmetic is None:
            self._vector_arithmetic = VectorArithmetic(self.embedding_engine)
        return self._vector_arithmetic
    
    def query(
        self,
        question: str,
        n_results: int = 5,
        use_hybrid: bool = True,
        include_parent_context: bool = True
    ) -> AgentResponse:
        """
        Query the book with a natural language question.
        
        Args:
            question: Natural language query
            n_results: Number of results to return
            use_hybrid: Use hybrid (dense + sparse) search
            include_parent_context: Include parent chunks for full context
            
        Returns:
            AgentResponse with results and combined context
        """
        # Generate query embedding
        query_embedding = self.embedding_engine.embed_query(question)
        
        # Search
        if use_hybrid:
            results = self.vector_store.hybrid_search(
                query_embedding=query_embedding,
                query_text=question,
                n_results=n_results
            )
        else:
            results = self.vector_store.search(
                query_embedding=query_embedding,
                n_results=n_results,
                include_parent=include_parent_context
            )
        
        # Build combined context
        context = self._build_context(results, include_parent_context)
        
        # Summarize metadata
        metadata_summary = self._summarize_metadata(results)
        
        return AgentResponse(
            query=question,
            results=results,
            context=context,
            metadata_summary=metadata_summary
        )
    
    def concept_search(
        self,
        expression: str,
        n_results: int = 5
    ) -> AgentResponse:
        """
        Search using vector arithmetic (concept algebra).
        
        Example expressions:
        - "Zero Trust + Cloud Architecture"
        - "SQL Injection - Web Application + Database"
        - "Ransomware + Prevention"
        
        This enables the famous King - Man + Woman = Queen style reasoning.
        """
        # Compute the concept vector
        concept_vector = self.vector_arithmetic.compute(expression)
        
        # Search for nearest concepts
        results = self.vector_store.search(
            query_embedding=concept_vector,
            n_results=n_results,
            include_parent=True
        )
        
        context = self._build_context(results, include_parent=True)
        metadata_summary = self._summarize_metadata(results)
        
        return AgentResponse(
            query=f"Concept: {expression}",
            results=results,
            context=context,
            metadata_summary=metadata_summary
        )
    
    def analogy_search(
        self,
        a: str,
        b: str,
        c: str,
        n_results: int = 5
    ) -> AgentResponse:
        """
        Search using analogical reasoning.
        
        "A is to B as C is to ?"
        
        Example:
        - analogy("SQL Injection", "Web Application", "Memory Corruption")
        - Returns content related to "Memory Corruption" exploits in applications
        """
        # Compute analogy vector
        analogy_vector = self.vector_arithmetic.analogy(a, b, c)
        
        # Search
        results = self.vector_store.search(
            query_embedding=analogy_vector,
            n_results=n_results,
            include_parent=True
        )
        
        context = self._build_context(results, include_parent=True)
        metadata_summary = self._summarize_metadata(results)
        
        return AgentResponse(
            query=f"Analogy: {a} → {b} :: {c} → ?",
            results=results,
            context=context,
            metadata_summary=metadata_summary
        )
    
    def search_by_category(
        self,
        question: str,
        category: str,
        n_results: int = 5
    ) -> AgentResponse:
        """
        Search within a specific cybersecurity category.
        
        Categories:
        - network_security
        - web_security
        - cryptography
        - malware
        - authentication
        - incident_response
        - compliance
        - cloud_security
        - memory_safety
        """
        query_embedding = self.embedding_engine.embed_query(question)
        
        results = self.vector_store.filter_by_category(
            query_embedding=query_embedding,
            category=category,
            n_results=n_results
        )
        
        context = self._build_context(results, include_parent=True)
        metadata_summary = self._summarize_metadata(results)
        
        return AgentResponse(
            query=f"{question} [Category: {category}]",
            results=results,
            context=context,
            metadata_summary=metadata_summary
        )
    
    def search_by_severity(
        self,
        question: str,
        severity: str,
        n_results: int = 5
    ) -> AgentResponse:
        """
        Search for content with specific severity level.
        
        Severities: critical, high, medium, low
        """
        query_embedding = self.embedding_engine.embed_query(question)
        
        results = self.vector_store.filter_by_severity(
            query_embedding=query_embedding,
            severity=severity,
            n_results=n_results
        )
        
        context = self._build_context(results, include_parent=True)
        metadata_summary = self._summarize_metadata(results)
        
        return AgentResponse(
            query=f"{question} [Severity: {severity}]",
            results=results,
            context=context,
            metadata_summary=metadata_summary
        )
    
    def _build_context(
        self,
        results: List[SearchResult],
        include_parent: bool = True
    ) -> str:
        """Build combined context string for AI agent consumption."""
        context_parts = []
        
        for i, result in enumerate(results, 1):
            part = f"[Source {i}]"
            
            # Add metadata breadcrumbs
            if result.metadata.get('chapter'):
                part += f"\nChapter: {result.metadata['chapter']}"
            if result.metadata.get('section'):
                part += f"\nSection: {result.metadata['section']}"
            if result.metadata.get('severity'):
                part += f"\nSeverity: {result.metadata['severity']}"
            
            part += f"\nRelevance Score: {result.score:.3f}"
            
            # Add content (prefer parent context if available)
            if include_parent and result.parent_content:
                part += f"\n\nContext:\n{result.parent_content}"
                part += f"\n\nRelevant excerpt:\n{result.content}"
            else:
                part += f"\n\nContent:\n{result.content}"
            
            context_parts.append(part)
        
        return "\n\n" + "="*60 + "\n\n".join(context_parts)
    
    def _summarize_metadata(self, results: List[SearchResult]) -> Dict[str, Any]:
        """Summarize metadata across all results."""
        chapters = set()
        sections = set()
        categories = set()
        severities = set()
        cves = set()
        
        for result in results:
            meta = result.metadata
            if meta.get('chapter'):
                chapters.add(meta['chapter'])
            if meta.get('section'):
                sections.add(meta['section'])
            if meta.get('categories'):
                cats = meta['categories']
                if isinstance(cats, str):
                    categories.update(cats.split(','))
                else:
                    categories.update(cats)
            if meta.get('severity'):
                severities.add(meta['severity'])
            if meta.get('cves'):
                cve_list = meta['cves']
                if isinstance(cve_list, str):
                    cves.update(cve_list.split(','))
                else:
                    cves.update(cve_list)
        
        return {
            'chapters': list(chapters),
            'sections': list(sections),
            'categories': list(categories),
            'severities': list(severities),
            'cves': list(cves),
            'result_count': len(results)
        }
    
    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the encoded book."""
        return {
            'total_chunks': self.vector_store.get_chunk_count(),
            'embedding_model': self.config.embedding.model_name,
            'embedding_dimension': self.embedding_engine.dimension,
            'collection_name': self.config.vector_db.collection_name
        }


# Convenience functions for direct use
def create_agent(
    db_path: str = "./vector_db",
    collection_name: str = "cybersecurity_book"
) -> CybersecurityBookAgent:
    """Create an agent connected to an existing encoded book."""
    from config import PipelineConfig, VectorDBConfig
    
    config = PipelineConfig(
        vector_db=VectorDBConfig(
            collection_name=collection_name,
            persist_directory=db_path
        )
    )
    
    return CybersecurityBookAgent(config=config)


def quick_query(
    question: str,
    db_path: str = "./vector_db",
    n_results: int = 5
) -> str:
    """Quick query function that returns just the context string."""
    agent = create_agent(db_path)
    response = agent.query(question, n_results=n_results)
    return response.context