Add analysis engine: batch pipeline, interactive Q&A, context builder
This commit is contained in:
n0mad1k
@@ -0,0 +1,3 @@
|
||||
from .claude_batch import BatchAnalyzer
|
||||
from .claude_interactive import InteractiveAnalyzer
|
||||
from .context_builder import ContextBuilder
|
||||
|
||||
@@ -0,0 +1,258 @@
|
||||
"""Batch document analysis pipeline for Mosaic."""
|
||||
import logging
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
|
||||
from rich.console import Console
|
||||
|
||||
from storage.db import MosaicDB
|
||||
from storage.models import Document
|
||||
from storage.jsonl_store import JSONLStore
|
||||
from utils.api_limiter import APIRateLimiter, CircuitBreakerOpen
|
||||
from utils.text_utils import estimate_tokens
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
console = Console()
|
||||
|
||||
EXTRACTOR_DOMAINS = {
|
||||
'tools': 'tool',
|
||||
'ttps': 'ttp',
|
||||
'tradecraft': 'tradecraft',
|
||||
'surveillance': 'surveillance',
|
||||
'infrastructure': 'infrastructure',
|
||||
'entities': 'entity',
|
||||
}
|
||||
|
||||
|
||||
class BatchAnalyzer:
|
||||
"""Run batch extraction pipeline across documents."""
|
||||
|
||||
def __init__(self, db: MosaicDB, api_key: str,
|
||||
model: str = "claude-sonnet-4-20250514",
|
||||
rate_limiter: Optional[APIRateLimiter] = None):
|
||||
self.db = db
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rate_limiter = rate_limiter or APIRateLimiter()
|
||||
self._extractors = {}
|
||||
|
||||
def _get_extractor(self, domain: str):
|
||||
"""Lazy-load extractor by domain name."""
|
||||
if domain not in self._extractors:
|
||||
from extractors.ttp_extractor import TTPExtractor
|
||||
from extractors.tool_extractor import ToolExtractor
|
||||
from extractors.tradecraft import TradecraftExtractor
|
||||
from extractors.surveillance import SurveillanceExtractor
|
||||
from extractors.infrastructure import InfrastructureExtractor
|
||||
from extractors.entity_extractor import EntityExtractor
|
||||
|
||||
extractor_classes = {
|
||||
'ttp': TTPExtractor,
|
||||
'tool': ToolExtractor,
|
||||
'tradecraft': TradecraftExtractor,
|
||||
'surveillance': SurveillanceExtractor,
|
||||
'infrastructure': InfrastructureExtractor,
|
||||
'entity': EntityExtractor,
|
||||
}
|
||||
|
||||
cls = extractor_classes.get(domain)
|
||||
if cls:
|
||||
if domain == 'infrastructure':
|
||||
self._extractors[domain] = cls(
|
||||
db=self.db, api_key=self.api_key,
|
||||
model=self.model, rate_limiter=self.rate_limiter,
|
||||
collection="", doc_title="", doc_date=None,
|
||||
)
|
||||
else:
|
||||
self._extractors[domain] = cls(
|
||||
db=self.db, api_key=self.api_key,
|
||||
model=self.model, rate_limiter=self.rate_limiter,
|
||||
)
|
||||
return self._extractors.get(domain)
|
||||
|
||||
def estimate(self, source: Optional[str] = None,
|
||||
domains: Optional[list[str]] = None,
|
||||
since: Optional[str] = None,
|
||||
until: Optional[str] = None,
|
||||
since_last_run: bool = False) -> dict:
|
||||
"""Estimate token costs without running extraction (dry-run).
|
||||
|
||||
Returns dict with document count, chunk count, estimated tokens.
|
||||
"""
|
||||
domains = domains or list(EXTRACTOR_DOMAINS.values())
|
||||
docs = self._get_target_documents(source, since, until, since_last_run)
|
||||
|
||||
total_chunks = 0
|
||||
total_input_tokens = 0
|
||||
|
||||
for doc in docs:
|
||||
for domain in domains:
|
||||
extractor = self._get_extractor(domain)
|
||||
if extractor:
|
||||
est = extractor.estimate_tokens(doc.text, doc.doc_type)
|
||||
total_chunks += est['chunks']
|
||||
total_input_tokens += est['estimated_input_tokens']
|
||||
|
||||
return {
|
||||
'document_count': len(docs),
|
||||
'domain_count': len(domains),
|
||||
'total_chunks': total_chunks,
|
||||
'estimated_input_tokens': total_input_tokens,
|
||||
'estimated_output_tokens': total_chunks * 2000,
|
||||
'estimated_total_tokens': total_input_tokens + (total_chunks * 2000),
|
||||
}
|
||||
|
||||
def run(self, source: Optional[str] = None,
|
||||
domains: Optional[list[str]] = None,
|
||||
since: Optional[str] = None,
|
||||
until: Optional[str] = None,
|
||||
since_last_run: bool = False,
|
||||
retry_failed: bool = False) -> dict:
|
||||
"""Run batch extraction pipeline.
|
||||
|
||||
Returns dict with counts of extracted items per domain.
|
||||
"""
|
||||
domains = domains or list(EXTRACTOR_DOMAINS.values())
|
||||
docs = self._get_target_documents(source, since, until, since_last_run)
|
||||
|
||||
if not docs:
|
||||
console.print("[yellow]No documents to process[/yellow]")
|
||||
return {}
|
||||
|
||||
results = {d: 0 for d in domains}
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
) as progress:
|
||||
for domain in domains:
|
||||
extractor = self._get_extractor(domain)
|
||||
if not extractor:
|
||||
console.print(f"[yellow]Unknown domain: {domain}[/yellow]")
|
||||
continue
|
||||
|
||||
task = progress.add_task(f"Extracting {domain}", total=len(docs))
|
||||
|
||||
for doc in docs:
|
||||
# Check if already done
|
||||
status = self.db.get_extraction_status(doc.id, extractor.EXTRACTOR_NAME)
|
||||
if status and status.status == 'DONE' and not retry_failed:
|
||||
progress.advance(task)
|
||||
continue
|
||||
if status and status.status == 'FAILED' and not retry_failed:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
try:
|
||||
# Update infrastructure extractor context
|
||||
if domain == 'infrastructure' and hasattr(extractor, 'collection'):
|
||||
extractor.collection = doc.source
|
||||
extractor.doc_title = doc.title
|
||||
extractor.doc_date = doc.date
|
||||
|
||||
count = extractor.extract_from_document(
|
||||
doc_id=doc.id,
|
||||
text=doc.text,
|
||||
doc_type=doc.doc_type,
|
||||
source=doc.source,
|
||||
metadata=doc.metadata,
|
||||
)
|
||||
results[domain] += count
|
||||
|
||||
except CircuitBreakerOpen:
|
||||
console.print(f"[red]Circuit breaker open — stopping {domain} extraction[/red]")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error("Extraction error for %s on doc %s: %s",
|
||||
domain, doc.id[:12], e)
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
return results
|
||||
|
||||
def run_full_pipeline(self, source: str, limit: Optional[int] = None) -> dict:
|
||||
"""Run full pipeline: collect → parse → extract → report.
|
||||
|
||||
Returns summary dict.
|
||||
"""
|
||||
from collectors import load_profile, CustomSource
|
||||
from parsers import ParserManager
|
||||
from utils.cache import DownloadCache
|
||||
from utils.http_limiter import HTTPRateLimiter
|
||||
|
||||
summary = {'collect': {}, 'parse': {}, 'extract': {}, 'export': []}
|
||||
|
||||
# Step 1: Collect
|
||||
console.print(f"\n[bold cyan]Step 1/4: Collecting from {source}[/bold cyan]")
|
||||
try:
|
||||
profile = load_profile(source)
|
||||
cache = DownloadCache()
|
||||
collector = CustomSource(profile, self.db, cache)
|
||||
summary['collect'] = collector.collect(limit=limit)
|
||||
except Exception as e:
|
||||
console.print(f"[red]Collection error: {e}[/red]")
|
||||
return summary
|
||||
|
||||
# Step 2: Parse
|
||||
console.print(f"\n[bold cyan]Step 2/4: Parsing documents[/bold cyan]")
|
||||
parser = ParserManager()
|
||||
docs = self.db.get_documents_by_source(source, status='CACHED')
|
||||
parsed = 0
|
||||
failed = 0
|
||||
|
||||
with Progress(SpinnerColumn(), TextColumn("{task.description}"),
|
||||
BarColumn(), TaskProgressColumn()) as progress:
|
||||
task = progress.add_task("Parsing", total=len(docs))
|
||||
for doc in docs:
|
||||
if doc.parse_status in ('PARSED', 'NEEDS_OCR'):
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
self.db.update_document_status(doc.id, parse_status='IN_PROGRESS')
|
||||
result = parser.parse(doc.raw_path, doc.doc_type, doc.content_hash)
|
||||
|
||||
if result.success:
|
||||
doc.text = result.text
|
||||
doc.title = result.title or doc.title
|
||||
doc.char_count = result.char_count
|
||||
doc.metadata.update(result.metadata)
|
||||
status = 'NEEDS_OCR' if result.needs_ocr else 'PARSED'
|
||||
doc.parse_status = status
|
||||
self.db.insert_document(doc)
|
||||
self.db.update_document_status(doc.id, parse_status=status)
|
||||
parsed += 1
|
||||
else:
|
||||
self.db.update_document_status(doc.id, parse_status='PARSE_FAILED')
|
||||
failed += 1
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
summary['parse'] = {'parsed': parsed, 'failed': failed}
|
||||
|
||||
# Step 3: Extract
|
||||
console.print(f"\n[bold cyan]Step 3/4: Running extraction[/bold cyan]")
|
||||
summary['extract'] = self.run(source=source)
|
||||
|
||||
# Step 4: Export
|
||||
console.print(f"\n[bold cyan]Step 4/4: Exporting results[/bold cyan]")
|
||||
store = JSONLStore(self.db)
|
||||
summary['export'] = store.export()
|
||||
|
||||
return summary
|
||||
|
||||
def _get_target_documents(self, source: Optional[str] = None,
|
||||
since: Optional[str] = None,
|
||||
until: Optional[str] = None,
|
||||
since_last_run: bool = False) -> list[Document]:
|
||||
"""Get documents to process based on filters."""
|
||||
if since or until:
|
||||
return self.db.get_documents_by_date_range(source=source, since=since, until=until)
|
||||
elif since_last_run:
|
||||
return self.db.get_unextracted_documents(source=source, since_last_run=True)
|
||||
elif source:
|
||||
return self.db.get_documents_by_source(source, parse_status='PARSED')
|
||||
else:
|
||||
return self.db.get_unextracted_documents()
|
||||
@@ -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]")
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Context builder for assembling relevant document chunks for Claude Q&A."""
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from storage.db import MosaicDB
|
||||
from utils.text_utils import estimate_tokens
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_TOKEN_BUDGET = 30000
|
||||
DEFAULT_MAX_RESULTS = 20
|
||||
|
||||
|
||||
class ContextBuilder:
|
||||
"""Build context from FTS5 search results for Claude queries.
|
||||
|
||||
Uses BM25 ranking from FTS5, assembles chunks within a token budget,
|
||||
and tracks source citations.
|
||||
"""
|
||||
|
||||
def __init__(self, db: MosaicDB, token_budget: int = DEFAULT_TOKEN_BUDGET,
|
||||
max_results: int = DEFAULT_MAX_RESULTS):
|
||||
self.db = db
|
||||
self.token_budget = token_budget
|
||||
self.max_results = max_results
|
||||
|
||||
def build_context(self, query: str, source: Optional[str] = None) -> dict:
|
||||
"""Search documents and assemble context within token budget.
|
||||
|
||||
Args:
|
||||
query: Search query string
|
||||
source: Optional source filter
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
'context_text': assembled context string with citations
|
||||
'documents': list of matched document metadata
|
||||
'total_tokens': estimated token count
|
||||
'truncated': whether results were truncated
|
||||
"""
|
||||
# Search with FTS5
|
||||
docs = self.db.search_documents(query, source=source, limit=self.max_results)
|
||||
|
||||
if not docs:
|
||||
return {
|
||||
'context_text': '',
|
||||
'documents': [],
|
||||
'total_tokens': 0,
|
||||
'truncated': False,
|
||||
}
|
||||
|
||||
# Assemble context within budget
|
||||
context_parts = []
|
||||
used_tokens = 0
|
||||
included_docs = []
|
||||
truncated = False
|
||||
|
||||
for doc in docs:
|
||||
# Build document excerpt
|
||||
excerpt = self._format_document_excerpt(doc)
|
||||
excerpt_tokens = estimate_tokens(excerpt)
|
||||
|
||||
if used_tokens + excerpt_tokens > self.token_budget:
|
||||
# Try to fit a truncated version
|
||||
remaining_tokens = self.token_budget - used_tokens
|
||||
if remaining_tokens > 500: # Worth including a partial
|
||||
truncated_excerpt = self._truncate_to_tokens(excerpt, remaining_tokens)
|
||||
context_parts.append(truncated_excerpt)
|
||||
used_tokens += remaining_tokens
|
||||
included_docs.append(self._doc_metadata(doc))
|
||||
truncated = True
|
||||
break
|
||||
|
||||
context_parts.append(excerpt)
|
||||
used_tokens += excerpt_tokens
|
||||
included_docs.append(self._doc_metadata(doc))
|
||||
|
||||
context_text = '\n\n---\n\n'.join(context_parts)
|
||||
|
||||
return {
|
||||
'context_text': context_text,
|
||||
'documents': included_docs,
|
||||
'total_tokens': used_tokens,
|
||||
'truncated': truncated,
|
||||
}
|
||||
|
||||
def build_entity_context(self, entity_name: str) -> str:
|
||||
"""Build context focused on a specific entity across all sources."""
|
||||
docs = self.db.search_documents(entity_name, limit=self.max_results)
|
||||
|
||||
parts = []
|
||||
used_tokens = 0
|
||||
for doc in docs:
|
||||
excerpt = self._format_document_excerpt(doc)
|
||||
tokens = estimate_tokens(excerpt)
|
||||
if used_tokens + tokens > self.token_budget:
|
||||
break
|
||||
parts.append(excerpt)
|
||||
used_tokens += tokens
|
||||
|
||||
return '\n\n---\n\n'.join(parts)
|
||||
|
||||
def build_tool_context(self, tool_name: str) -> str:
|
||||
"""Build context focused on a specific tool/implant."""
|
||||
# Search both FTS and extracted tools
|
||||
docs = self.db.search_documents(tool_name, limit=self.max_results)
|
||||
|
||||
parts = []
|
||||
used_tokens = 0
|
||||
|
||||
# Add extracted tool data if available
|
||||
tools = self.db.get_all_tools()
|
||||
for tool in tools:
|
||||
if tool_name.lower() in tool.name.lower() or \
|
||||
any(tool_name.lower() in a.lower() for a in tool.aliases):
|
||||
tool_text = f"[Extracted Tool: {tool.name}]\n"
|
||||
tool_text += f"Capability: {tool.capability}\n"
|
||||
tool_text += f"Description: {tool.description}\n"
|
||||
tool_text += f"Platforms: {', '.join(tool.target_platforms)}\n"
|
||||
tool_text += f"CVEs: {', '.join(tool.cves)}\n"
|
||||
tool_text += f"Source context: {tool.source_context}\n"
|
||||
parts.append(tool_text)
|
||||
used_tokens += estimate_tokens(tool_text)
|
||||
|
||||
# Add document excerpts
|
||||
for doc in docs:
|
||||
excerpt = self._format_document_excerpt(doc)
|
||||
tokens = estimate_tokens(excerpt)
|
||||
if used_tokens + tokens > self.token_budget:
|
||||
break
|
||||
parts.append(excerpt)
|
||||
used_tokens += tokens
|
||||
|
||||
return '\n\n---\n\n'.join(parts)
|
||||
|
||||
def _format_document_excerpt(self, doc) -> str:
|
||||
"""Format a document for context inclusion with citation info."""
|
||||
header = f"[Source: {doc.source} | Title: {doc.title}"
|
||||
if doc.date:
|
||||
header += f" | Date: {doc.date}"
|
||||
if doc.classification:
|
||||
header += f" | Classification: {doc.classification}"
|
||||
header += f" | ID: {doc.id[:12]}]"
|
||||
|
||||
# Truncate text if needed
|
||||
text = doc.text
|
||||
max_chars = (self.token_budget // len(self.db.search_documents("", limit=1) or [doc])) * 4
|
||||
if len(text) > max_chars:
|
||||
text = text[:max_chars] + "\n[...truncated...]"
|
||||
|
||||
return f"{header}\n\n{text}"
|
||||
|
||||
def _truncate_to_tokens(self, text: str, max_tokens: int) -> str:
|
||||
"""Truncate text to approximately max_tokens."""
|
||||
max_chars = max_tokens * 4 # Rough chars-per-token estimate
|
||||
if len(text) <= max_chars:
|
||||
return text
|
||||
return text[:max_chars] + "\n[...truncated...]"
|
||||
|
||||
def _doc_metadata(self, doc) -> dict:
|
||||
"""Extract metadata dict from a document for citation tracking."""
|
||||
return {
|
||||
'id': doc.id,
|
||||
'title': doc.title,
|
||||
'source': doc.source,
|
||||
'date': doc.date,
|
||||
'doc_type': doc.doc_type,
|
||||
}
|
||||
Reference in New Issue
Block a user