Add analysis engine: batch pipeline, interactive Q&A, context builder
This commit is contained in:
n0mad1k
@@ -0,0 +1,206 @@
|
||||
"""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 <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 <name> 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]")
|
||||
Reference in New Issue
Block a user