"""Interactive Q&A mode for Mosaic — Rich-based chat over extracted intelligence.""" import logging from typing import Optional import anthropic from rich.console import Console from rich.panel import Panel from rich.markdown import Markdown from rich.prompt import Prompt from storage.db import MosaicDB from .context_builder import ContextBuilder logger = logging.getLogger(__name__) console = Console() # Import prompts — handle if not yet created try: from .prompts import ( QUERY_SYSTEM_PROMPT, QUERY_PROMPT, EMULATE_PROMPT, COMPARE_PROMPT, BRIEF_PROMPT, SCENARIO_PROMPT ) except ImportError: QUERY_SYSTEM_PROMPT = "You are an intelligence analyst. Answer questions based on provided documents." QUERY_PROMPT = "{context}\n\nQuestion: {question}" EMULATE_PROMPT = "Emulate threat actor {actor}:\n{context}" COMPARE_PROMPT = "Compare {tool_a} vs {tool_b}:\n{context}" BRIEF_PROMPT = "Brief on {topic}:\n{context}" SCENARIO_PROMPT = "Scenario for {target}:\n{context}" class InteractiveAnalyzer: """Interactive Q&A over extracted intelligence corpus.""" def __init__(self, db: MosaicDB, api_key: str, model: str = "claude-sonnet-4-20250514", source: Optional[str] = None): self.db = db self.client = anthropic.Anthropic(api_key=api_key) self.model = model self.source = source self.context_builder = ContextBuilder(db) self.mode = "default" self.conversation_history = [] def run(self): """Start interactive Q&A session.""" self._show_welcome() while True: try: user_input = Prompt.ask("\n[bold green]mosaic[/bold green]") except (KeyboardInterrupt, EOFError): console.print("\n[dim]Goodbye.[/dim]") break user_input = user_input.strip() if not user_input: continue # Handle commands if user_input.lower() in ('exit', 'quit', 'q'): console.print("[dim]Goodbye.[/dim]") break elif user_input.lower() == 'help': self._show_help() continue elif user_input.lower().startswith('mode '): self.mode = user_input[5:].strip() console.print(f"[cyan]Mode set to: {self.mode}[/cyan]") continue elif user_input.lower().startswith('source '): self.source = user_input[7:].strip() console.print(f"[cyan]Source filter: {self.source}[/cyan]") continue elif user_input.lower() == 'clear': self.conversation_history = [] console.print("[dim]Conversation cleared.[/dim]") continue elif user_input.lower() == 'status': self._show_status() continue # Process query based on mode self._handle_query(user_input) def _handle_query(self, query: str): """Process a query in the current mode.""" with console.status("[bold cyan]Searching corpus...[/bold cyan]"): # Build context based on mode if self.mode == 'emulate': context = self.context_builder.build_context(query, self.source) prompt = EMULATE_PROMPT.format(actor=query, context=context['context_text']) elif self.mode == 'compare' and ' vs ' in query: parts = query.split(' vs ', 1) tool_a = parts[0].strip() tool_b = parts[1].strip() ctx_a = self.context_builder.build_tool_context(tool_a) ctx_b = self.context_builder.build_tool_context(tool_b) context = {'context_text': ctx_a + '\n\n' + ctx_b, 'documents': [], 'total_tokens': 0} prompt = COMPARE_PROMPT.format(tool_a=tool_a, tool_b=tool_b, context=context['context_text']) elif self.mode == 'brief': context = self.context_builder.build_context(query, self.source) prompt = BRIEF_PROMPT.format(topic=query, context=context['context_text']) elif self.mode == 'scenario': context = self.context_builder.build_context(query, self.source) prompt = SCENARIO_PROMPT.format(target=query, context=context['context_text']) else: context = self.context_builder.build_context(query, self.source) prompt = QUERY_PROMPT.format(context=context['context_text'], question=query) if not context.get('context_text'): console.print("[yellow]No relevant documents found for your query.[/yellow]") return # Show source count doc_count = len(context.get('documents', [])) console.print(f"[dim]Found {doc_count} relevant documents, " f"~{context.get('total_tokens', 0):,} tokens of context[/dim]") # Call Claude with console.status("[bold cyan]Analyzing...[/bold cyan]"): try: messages = list(self.conversation_history) messages.append({"role": "user", "content": prompt}) response = self.client.messages.create( model=self.model, max_tokens=4096, system=QUERY_SYSTEM_PROMPT, messages=messages, ) answer = response.content[0].text # Save to conversation history self.conversation_history.append({"role": "user", "content": prompt}) self.conversation_history.append({"role": "assistant", "content": answer}) # Trim history if too long if len(self.conversation_history) > 20: self.conversation_history = self.conversation_history[-10:] except anthropic.APIError as e: console.print(f"[red]API error: {e}[/red]") return # Display response console.print() console.print(Panel( Markdown(answer), title="[bold]Analysis[/bold]", border_style="cyan", )) # Show citations if context.get('documents'): console.print("\n[dim]Sources:[/dim]") for doc in context['documents'][:5]: console.print(f" [dim]- {doc['source']}: {doc['title']}", markup=False) def _show_welcome(self): """Show welcome message.""" welcome = """[bold cyan]MOSAIC Interactive Intelligence Query[/bold cyan] Ask questions about the extracted intelligence corpus. Type [bold]help[/bold] for commands, [bold]exit[/bold] to quit.""" if self.source: welcome += f"\n[dim]Source filter: {self.source}[/dim]" console.print(Panel(welcome, border_style="cyan")) def _show_help(self): """Show help text.""" help_text = """[bold]Commands:[/bold] help Show this help exit/quit Leave interactive mode clear Clear conversation history status Show corpus statistics mode Set query mode: default - General Q&A emulate - Threat actor emulation playbook compare - Compare tools (use 'tool_a vs tool_b') brief - Intelligence briefing scenario - Adversary simulation scenario source Filter to specific source (or 'all') [bold]Tips:[/bold] - Ask specific questions for better results - Use source names (vault7, cablegate) to focus queries - Emulate mode: type an actor name to get their TTP playbook""" console.print(Panel(help_text, border_style="cyan")) def _show_status(self): """Show corpus statistics.""" stats = self.db.get_status_counts(self.source) tokens = self.db.get_total_tokens(self.source) console.print(f"\n[bold]Corpus Statistics[/bold]") console.print(f" Documents: {stats.get('total', 0)}") console.print(f" Parsed: {stats.get('parsing', {}).get('PARSED', 0)}") console.print(f" Tokens used: {tokens.get('total_tokens', 0):,}") if self.source: console.print(f" [dim]Filtered to: {self.source}[/dim]")