mirror of
https://github.com/mblanke/ThreatHunt.git
synced 2026-03-01 14:00:20 -05:00
feat: interactive network map, IOC highlighting, AUP hunt selector, type filters
- NetworkMap: hunt-scoped force-directed graph with click-to-inspect popover - NetworkMap: zoom/pan (wheel, drag, buttons), viewport transform - NetworkMap: clickable IP/Host/Domain/URL legend chips to filter node types - NetworkMap: brighter colors, 20% smaller nodes - DatasetViewer: IOC columns highlighted with colored headers + cell tinting - AUPScanner: hunt dropdown replacing dataset checkboxes, auto-select all - Rename 'Social Media (Personal)' theme to 'Social Media' with DB migration - Fix /api/hunts timeout: Dataset.rows lazy='noload' (was selectin cascade) - Add OS column mapping to normalizer - Full backend services, DB models, alembic migrations, new routes - New components: Dashboard, HuntManager, FileUpload, NetworkMap, etc. - Docker Compose deployment with nginx reverse proxy
This commit is contained in:
161
backend/app/agents/registry.py
Normal file
161
backend/app/agents/registry.py
Normal file
@@ -0,0 +1,161 @@
|
||||
"""Model registry — inventory of all Ollama models across Wile and Roadrunner.
|
||||
|
||||
Each model is tagged with capabilities (chat, code, vision, embedding) and
|
||||
performance tier (fast, medium, heavy) for the TaskRouter.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class Capability(str, Enum):
|
||||
CHAT = "chat"
|
||||
CODE = "code"
|
||||
VISION = "vision"
|
||||
EMBEDDING = "embedding"
|
||||
|
||||
|
||||
class Tier(str, Enum):
|
||||
FAST = "fast" # < 15B params — quick responses
|
||||
MEDIUM = "medium" # 15–40B params — balanced
|
||||
HEAVY = "heavy" # 40B+ params — deep analysis
|
||||
|
||||
|
||||
class Node(str, Enum):
|
||||
WILE = "wile"
|
||||
ROADRUNNER = "roadrunner"
|
||||
CLUSTER = "cluster" # Open WebUI balances across both
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelEntry:
|
||||
name: str
|
||||
node: Node
|
||||
capabilities: list[Capability]
|
||||
tier: Tier
|
||||
param_size: str = "" # e.g. "7b", "70b"
|
||||
notes: str = ""
|
||||
|
||||
|
||||
# ── Roadrunner (100.110.190.11) ──────────────────────────────────────
|
||||
|
||||
ROADRUNNER_MODELS: list[ModelEntry] = [
|
||||
# General / chat
|
||||
ModelEntry("llama3.1:latest", Node.ROADRUNNER, [Capability.CHAT], Tier.FAST, "8b"),
|
||||
ModelEntry("qwen2.5:14b-instruct", Node.ROADRUNNER, [Capability.CHAT], Tier.FAST, "14b"),
|
||||
ModelEntry("mistral:7b-instruct", Node.ROADRUNNER, [Capability.CHAT], Tier.FAST, "7b"),
|
||||
ModelEntry("mistral:7b", Node.ROADRUNNER, [Capability.CHAT], Tier.FAST, "7b"),
|
||||
ModelEntry("qwen2.5:7b", Node.ROADRUNNER, [Capability.CHAT], Tier.FAST, "7b"),
|
||||
ModelEntry("phi3:medium", Node.ROADRUNNER, [Capability.CHAT], Tier.MEDIUM, "14b"),
|
||||
# Code
|
||||
ModelEntry("qwen2.5-coder:7b", Node.ROADRUNNER, [Capability.CODE], Tier.FAST, "7b"),
|
||||
ModelEntry("qwen2.5-coder:latest", Node.ROADRUNNER, [Capability.CODE], Tier.FAST, "7b"),
|
||||
ModelEntry("codestral:latest", Node.ROADRUNNER, [Capability.CODE], Tier.MEDIUM, "22b"),
|
||||
ModelEntry("codellama:13b", Node.ROADRUNNER, [Capability.CODE], Tier.FAST, "13b"),
|
||||
# Vision
|
||||
ModelEntry("llama3.2-vision:11b", Node.ROADRUNNER, [Capability.VISION], Tier.FAST, "11b"),
|
||||
ModelEntry("minicpm-v:latest", Node.ROADRUNNER, [Capability.VISION], Tier.FAST, "8b"),
|
||||
ModelEntry("llava:13b", Node.ROADRUNNER, [Capability.VISION], Tier.FAST, "13b"),
|
||||
# Embeddings
|
||||
ModelEntry("bge-m3:latest", Node.ROADRUNNER, [Capability.EMBEDDING], Tier.FAST, "0.6b"),
|
||||
ModelEntry("nomic-embed-text:latest", Node.ROADRUNNER, [Capability.EMBEDDING], Tier.FAST, "0.1b"),
|
||||
# Heavy
|
||||
ModelEntry("llama3.1:70b-instruct-q4_K_M", Node.ROADRUNNER, [Capability.CHAT], Tier.HEAVY, "70b"),
|
||||
]
|
||||
|
||||
# ── Wile (100.110.190.12) ────────────────────────────────────────────
|
||||
|
||||
WILE_MODELS: list[ModelEntry] = [
|
||||
# General / chat
|
||||
ModelEntry("llama3.1:latest", Node.WILE, [Capability.CHAT], Tier.FAST, "8b"),
|
||||
ModelEntry("llama3:latest", Node.WILE, [Capability.CHAT], Tier.FAST, "8b"),
|
||||
ModelEntry("gemma2:27b", Node.WILE, [Capability.CHAT], Tier.MEDIUM, "27b"),
|
||||
# Code
|
||||
ModelEntry("qwen2.5-coder:7b", Node.WILE, [Capability.CODE], Tier.FAST, "7b"),
|
||||
ModelEntry("qwen2.5-coder:latest", Node.WILE, [Capability.CODE], Tier.FAST, "7b"),
|
||||
ModelEntry("qwen2.5-coder:32b", Node.WILE, [Capability.CODE], Tier.MEDIUM, "32b"),
|
||||
ModelEntry("deepseek-coder:33b", Node.WILE, [Capability.CODE], Tier.MEDIUM, "33b"),
|
||||
ModelEntry("codestral:latest", Node.WILE, [Capability.CODE], Tier.MEDIUM, "22b"),
|
||||
# Vision
|
||||
ModelEntry("llava:13b", Node.WILE, [Capability.VISION], Tier.FAST, "13b"),
|
||||
# Embeddings
|
||||
ModelEntry("bge-m3:latest", Node.WILE, [Capability.EMBEDDING], Tier.FAST, "0.6b"),
|
||||
# Heavy
|
||||
ModelEntry("llama3.1:70b", Node.WILE, [Capability.CHAT], Tier.HEAVY, "70b"),
|
||||
ModelEntry("llama3.1:70b-instruct-q4_K_M", Node.WILE, [Capability.CHAT], Tier.HEAVY, "70b"),
|
||||
ModelEntry("llama3.1:70b-instruct-q5_K_M", Node.WILE, [Capability.CHAT], Tier.HEAVY, "70b"),
|
||||
ModelEntry("mixtral:8x22b-instruct", Node.WILE, [Capability.CHAT], Tier.HEAVY, "141b"),
|
||||
ModelEntry("qwen2:72b-instruct", Node.WILE, [Capability.CHAT], Tier.HEAVY, "72b"),
|
||||
]
|
||||
|
||||
ALL_MODELS = ROADRUNNER_MODELS + WILE_MODELS
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""Registry of all available models and their capabilities."""
|
||||
|
||||
def __init__(self, models: list[ModelEntry] | None = None):
|
||||
self.models = models or ALL_MODELS
|
||||
self._by_name: dict[str, list[ModelEntry]] = {}
|
||||
self._by_capability: dict[Capability, list[ModelEntry]] = {}
|
||||
self._by_node: dict[Node, list[ModelEntry]] = {}
|
||||
self._index()
|
||||
|
||||
def _index(self):
|
||||
for m in self.models:
|
||||
self._by_name.setdefault(m.name, []).append(m)
|
||||
for cap in m.capabilities:
|
||||
self._by_capability.setdefault(cap, []).append(m)
|
||||
self._by_node.setdefault(m.node, []).append(m)
|
||||
|
||||
def find(
|
||||
self,
|
||||
capability: Capability | None = None,
|
||||
tier: Tier | None = None,
|
||||
node: Node | None = None,
|
||||
) -> list[ModelEntry]:
|
||||
"""Find models matching all given criteria."""
|
||||
results = list(self.models)
|
||||
if capability:
|
||||
results = [m for m in results if capability in m.capabilities]
|
||||
if tier:
|
||||
results = [m for m in results if m.tier == tier]
|
||||
if node:
|
||||
results = [m for m in results if m.node == node]
|
||||
return results
|
||||
|
||||
def get_best(
|
||||
self,
|
||||
capability: Capability,
|
||||
prefer_tier: Tier | None = None,
|
||||
prefer_node: Node | None = None,
|
||||
) -> ModelEntry | None:
|
||||
"""Get the best model for a capability, with optional preference."""
|
||||
candidates = self.find(capability=capability, tier=prefer_tier, node=prefer_node)
|
||||
if not candidates:
|
||||
candidates = self.find(capability=capability, tier=prefer_tier)
|
||||
if not candidates:
|
||||
candidates = self.find(capability=capability)
|
||||
return candidates[0] if candidates else None
|
||||
|
||||
def list_nodes(self) -> list[Node]:
|
||||
return list(self._by_node.keys())
|
||||
|
||||
def list_models_on_node(self, node: Node) -> list[ModelEntry]:
|
||||
return self._by_node.get(node, [])
|
||||
|
||||
def to_dict(self) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"name": m.name,
|
||||
"node": m.node.value,
|
||||
"capabilities": [c.value for c in m.capabilities],
|
||||
"tier": m.tier.value,
|
||||
"param_size": m.param_size,
|
||||
}
|
||||
for m in self.models
|
||||
]
|
||||
|
||||
|
||||
# Singleton
|
||||
registry = ModelRegistry()
|
||||
Reference in New Issue
Block a user