aisoc 0.1.0__py3-none-any.whl

aisoc/__init__.py

@@ -0,0 +1,108 @@
+"""aisoc — an AI SOC kernel.
+
+A SOC modeled as a *team*, not a classifier: a set of role-based LLM agents —
+triage, Tier 2, IR Lead, Threat Intel, Threat Hunter, Detection Engineer, and a
+SOC Manager over the top — that work a case together over a shared, event-sourced
+bus, with a human in the loop on every consequential action.
+
+aisoc owns the kernel: the event contract, the bus, the case-memory read models,
+and the agent framework. It owns none of your environment — you inject three
+seams (see :mod:`aisoc.seams`): an LLM, an alert source, and a tool registry.
+
+Quick start (zero infrastructure)::
+
+    from aisoc import InMemoryBus, AlertTriaged, STREAM_TRIAGE, parse_event
+
+    bus = InMemoryBus()
+    bus.publish(STREAM_TRIAGE, AlertTriaged(
+        correlation_id="TICKET-1", produced_by="triage", ticket_id="TICKET-1",
+        verdict="false_positive", confidence=0.91, summary="benign scanner",
+    ))
+
+    for raw in bus.replay():            # read the audit log back
+        print(parse_event(raw).event_type)
+"""
+
+from aisoc.bus import (
+    ALL_STREAMS,
+    STREAM_ALERTS,
+    STREAM_AUDIT,
+    STREAM_CASES,
+    STREAM_TRIAGE,
+    Bus,
+    InMemoryBus,
+)
+from aisoc.schemas import (
+    EVENT_TYPES,
+    VALID_SEVERITIES,
+    VALID_VERDICTS,
+    ActionDecision,
+    ActionProposed,
+    AlertReceived,
+    AlertTriaged,
+    BusEvent,
+    CampaignDetected,
+    CaseEscalated,
+    DetectionTuningReport,
+    HuntingReport,
+    IRPlan,
+    Severity,
+    ShiftSummary,
+    ThreatIntelReport,
+    Tier2Analysis,
+    Verdict,
+    parse_event,
+)
+from aisoc.seams import AlertSource, ChatModel, StructuredChatModel, ToolProvider
+
+# Case memory — the event-log read models + their sidecar stores. Imported so
+# they're part of the documented surface: `from aisoc import case_memory`.
+from aisoc import case_memory, hitl_store, verdict_store
+from aisoc.config import data_dir
+from aisoc.extract import extract_indicators, parse_event_ts
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Bus
+    "Bus",
+    "InMemoryBus",
+    "STREAM_ALERTS",
+    "STREAM_TRIAGE",
+    "STREAM_CASES",
+    "STREAM_AUDIT",
+    "ALL_STREAMS",
+    # Event contract
+    "BusEvent",
+    "AlertReceived",
+    "AlertTriaged",
+    "CaseEscalated",
+    "Tier2Analysis",
+    "IRPlan",
+    "ActionProposed",
+    "ActionDecision",
+    "ThreatIntelReport",
+    "DetectionTuningReport",
+    "HuntingReport",
+    "CampaignDetected",
+    "ShiftSummary",
+    "Verdict",
+    "VALID_VERDICTS",
+    "Severity",
+    "VALID_SEVERITIES",
+    "EVENT_TYPES",
+    "parse_event",
+    # Injection seams
+    "ChatModel",
+    "StructuredChatModel",
+    "AlertSource",
+    "ToolProvider",
+    # Case memory
+    "case_memory",
+    "verdict_store",
+    "hitl_store",
+    "data_dir",
+    "extract_indicators",
+    "parse_event_ts",
+    "__version__",
+]

aisoc/agents/__init__.py

@@ -0,0 +1,60 @@
+"""aisoc role agents — the bus-consuming team over injected seams.
+
+Part of the ``agent`` extra (``pip install "aisoc[agent]"``). Two shapes:
+
+**Per-ticket roles** subclass :class:`~aisoc.agents.base.Agent` and are
+constructed with an injected bus, chat model, and tool provider. They consume
+events and publish follow-ups:
+
+    from aisoc import InMemoryBus
+    from aisoc.agents import TriageAgent, Tier2Agent, IRLeadAgent, ThreatIntelAgent
+
+    TriageAgent(bus=bus, model=model, tools=tools).run()   # alerts -> verdicts
+
+The chain is: ``TriageAgent`` (alert → first-pass verdict) → ``Tier2Agent``
+(deeper look, escalate?) → ``IRLeadAgent`` (containment plan + a human-gated
+``ActionProposed``) → ``ThreatIntelAgent`` (actor attribution).
+
+**Windowed roles** are scheduled ``run_once(*, bus, model=...)`` functions that
+replay the audit log over a time window and publish a report — call them on a
+timer rather than running a consume loop:
+
+    from aisoc.agents import detection_eng, soc_manager, threat_hunter, campaign_detector
+
+    detection_eng.run_once(bus=bus, model=model, window_hours=24)
+    soc_manager.run_once(bus=bus, model=model, window_hours=8)
+    threat_hunter.run_once(bus=bus, model=model, window_hours=24)
+    campaign_detector.run_once(bus=bus, model=model, window_days=14)
+
+Everything is injected — no role touches your environment directly. The same
+code runs in a unit test against a stub, in the demo, or in production.
+"""
+
+from aisoc.agents import (
+    campaign_detector,
+    detection_eng,
+    soc_manager,
+    threat_hunter,
+)
+from aisoc.agents.base import Agent, BudgetExceeded, parse_json_block
+from aisoc.agents.ir_lead import IRLeadAgent
+from aisoc.agents.threat_intel import ThreatIntelAgent
+from aisoc.agents.tier2 import Tier2Agent
+from aisoc.agents.triage import TriageAgent
+
+__all__ = [
+    # framework
+    "Agent",
+    "BudgetExceeded",
+    "parse_json_block",
+    # per-ticket roles
+    "TriageAgent",
+    "Tier2Agent",
+    "IRLeadAgent",
+    "ThreatIntelAgent",
+    # windowed roles (call <module>.run_once)
+    "detection_eng",
+    "soc_manager",
+    "threat_hunter",
+    "campaign_detector",
+]

aisoc/agents/base.py

@@ -0,0 +1,328 @@
+"""Base class for aisoc role agents.
+
+Each role (triage, Tier 2, IR Lead, Threat Intel, Threat Hunter, Detection
+Engineer, SOC Manager) subclasses :class:`Agent` and implements:
+
+    - ``role``: short identifier — the consumer group *and* the key the
+      :class:`~aisoc.seams.ToolProvider` is asked for tools with
+    - ``streams_to_consume()``: which bus streams to subscribe to
+    - ``handle(stream, event)``: process one event, publish follow-up events
+
+Everything an agent needs is **injected** — there are no environment hooks in
+this layer:
+
+    - the **bus** (:class:`~aisoc.bus.Bus`) it reads from and publishes to,
+    - the **chat model** (:class:`~aisoc.seams.ChatModel`) it reasons with,
+    - the **tool provider** (:class:`~aisoc.seams.ToolProvider`) that hands each
+      role its allowed tools.
+
+The base owns the parts that are the same for every role: consumer-group
+registration, the bind-tools-and-iterate **tool-call loop**, per-event tool-call
+budget enforcement, and graceful SIGTERM shutdown. A subclass is then just "which
+events do I care about, and what do I publish in response."
+
+This module is part of the ``agent`` extra (``pip install "aisoc[agent]"``) — it
+imports ``langchain_core`` for its message/tool types. The kernel itself (bus,
+schemas, seams, case memory) never imports it.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+import signal
+import time
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
+from pydantic import BaseModel
+
+from aisoc.bus import Bus
+from aisoc.seams import ChatModel, ToolProvider
+
+logger = logging.getLogger(__name__)
+
+
+class BudgetExceeded(Exception):
+    """Raised when an agent burns through its per-event tool-call budget."""
+
+
+def parse_json_block(text: str) -> Optional[dict[str, Any]]:
+    """Pull the first JSON object out of an LLM response, tolerant of fences.
+
+    Role agents ask the model for a strict JSON decision block; models still
+    wrap it in prose or ```` ```json ```` fences often enough that every role
+    needs this. Returns ``None`` when no object can be recovered.
+    """
+    if not text:
+        return None
+    text = text.strip()
+    text = re.sub(r"^```(?:json)?\s*", "", text)
+    text = re.sub(r"\s*```$", "", text)
+    start = text.find("{")
+    end = text.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        return None
+    try:
+        result = json.loads(text[start:end + 1])
+    except json.JSONDecodeError as exc:
+        logger.warning("parse_json_block: JSON parse failed: %s", exc)
+        return None
+    return result if isinstance(result, dict) else None
+
+
+class Agent(ABC):
+    """Abstract bus-consuming agent over the three injection seams.
+
+    Subclasses set ``role`` and optionally override ``budget`` (the per-event
+    tool-call cap — enrichment-heavy roles want it higher, advisory roles
+    lower).
+    """
+
+    role: str = ""
+    budget: int = 20
+
+    def __init__(
+        self,
+        bus: Bus,
+        model: ChatModel,
+        tools: Optional[ToolProvider] = None,
+    ):
+        if not self.role:
+            raise ValueError(f"{type(self).__name__} must set role")
+
+        self.bus = bus
+        self.model = model
+        self.tools = tools
+        # Unique consumer id per process so multiple instances of one role can
+        # run in parallel and share the consumer group (the bus distributes
+        # events across them).
+        self.consumer = f"{self.role}-{int(time.time())}"
+        self._stop = False
+        self._tool_calls_used = 0
+
+        signal.signal(signal.SIGTERM, self._on_signal)
+        signal.signal(signal.SIGINT, self._on_signal)
+
+    # ----- subclass surface ------------------------------------------------
+
+    @abstractmethod
+    def streams_to_consume(self) -> list[str]:
+        """Bus stream names this agent reads from."""
+
+    @abstractmethod
+    def handle(self, stream: str, event: dict[str, Any]) -> None:
+        """Process one event. Call ``self.publish()`` to emit follow-ups."""
+
+    # ----- shared infrastructure ------------------------------------------
+
+    @property
+    def consumer_group(self) -> str:
+        return self.role
+
+    def tools_for_role(self) -> list[Any]:
+        """The tools this role may call, from the injected provider (or none)."""
+        if self.tools is None:
+            return []
+        return list(self.tools.tools_for(self.role))
+
+    def publish(self, stream: str, event: BaseModel) -> str:
+        """Publish an event, then fire the :meth:`notify` side-effect hook."""
+        msg_id = self.bus.publish(stream, event)
+        try:
+            self.notify(stream, event)
+        except Exception:  # a notifier side-effect must never break the case
+            logger.exception("agent.%s notify hook failed on %s", self.role, stream)
+        return msg_id
+
+    def notify(self, stream: str, event: BaseModel) -> None:
+        """Side-effect hook, fired after each successful publish. No-op by default.
+
+        The kernel notifies no one — every result is just an event on the bus,
+        and the intended consumer is another agent (or a dedicated notifier that
+        subscribes to the stream). But an integration often wants to *also* push
+        an out-of-band notification — a chat card, a page, a webhook — at the
+        moment a role publishes. Override this to do that. It runs only after the
+        event is durably on the bus, dispatch on ``event.event_type``, and any
+        exception it raises is logged and swallowed so a flaky notifier can never
+        break the investigation it is reporting on.
+        """
+
+    # ----- the tool-call loop ---------------------------------------------
+
+    def run_tool_loop(
+        self,
+        messages: list[BaseMessage],
+        tools: Optional[list[Any]] = None,
+        max_iterations: Optional[int] = None,
+    ) -> tuple[list[BaseMessage], list[dict[str, Any]]]:
+        """Bind ``tools`` to the model and iterate until it stops calling them.
+
+        The generic version of what every reasoning role does: invoke the model,
+        run any tool calls it asks for (each charged against the per-event
+        budget), feed the results back, and repeat until the model answers
+        without a tool call or the budget runs out. Returns the grown message
+        list (the final answer is the last message) plus a trace of the tool
+        calls made, for the verdict record.
+
+        Models that can't bind tools (a plain stub, a non-LangChain wrapper) are
+        invoked as-is — the loop then just does a single pass.
+        """
+        tools = tools if tools is not None else self.tools_for_role()
+        tool_map = {getattr(t, "name", str(i)): t for i, t in enumerate(tools)}
+        cap = max_iterations if max_iterations is not None else self.budget
+
+        bind = getattr(self.model, "bind_tools", None)
+        runnable: Any = bind(tools) if (tools and callable(bind)) else self.model
+
+        trace: list[dict[str, Any]] = []
+        for _ in range(cap + 1):  # +1 so the final no-tool answer is captured
+            resp = runnable.invoke(messages)
+            messages.append(resp)
+            tool_calls = getattr(resp, "tool_calls", None) or []
+            if not tool_calls:
+                break
+            for call in tool_calls:
+                name = call.get("name", "")
+                args = call.get("args", {}) or {}
+                call_id = call.get("id") or name
+                try:
+                    self.check_budget()
+                except BudgetExceeded:
+                    messages.append(ToolMessage(
+                        content="error: tool-call budget exhausted",
+                        tool_call_id=call_id,
+                    ))
+                    return messages, trace
+                result = self._invoke_tool(tool_map.get(name), name, args)
+                trace.append({"tool": name, "args": args})
+                messages.append(ToolMessage(content=str(result), tool_call_id=call_id))
+        return messages, trace
+
+    @staticmethod
+    def _invoke_tool(tool: Any, name: str, args: dict[str, Any]) -> Any:
+        """Run one tool call, duck-typing LangChain tools and plain callables."""
+        if tool is None:
+            return f"error: unknown tool {name!r}"
+        try:
+            invoke = getattr(tool, "invoke", None)
+            if callable(invoke):
+                return invoke(args)
+            return tool(**args)
+        except Exception as exc:  # a tool failure must not kill the case
+            logger.warning("tool %s raised: %s", name, exc)
+            return f"error: {exc}"
+
+    def final_text(self, messages: list[BaseMessage]) -> str:
+        """The text content of the last message, or '' if there is none."""
+        if not messages:
+            return ""
+        content = getattr(messages[-1], "content", "")
+        return content if isinstance(content, str) else str(content)
+
+    def force_json(
+        self, messages: list[BaseMessage], instruction: str | None = None,
+    ) -> dict[str, Any]:
+        """Coerce a JSON decision block out of the conversation.
+
+        Tries the last message first; if it carries no JSON, asks the model once
+        more for the block alone (no tools bound). Returns ``{}`` if even that
+        fails — callers coerce missing fields against their own defaults.
+        """
+        decision = parse_json_block(self.final_text(messages))
+        if decision is not None:
+            return decision
+        messages.append(HumanMessage(
+            content=instruction or "Emit ONLY the JSON decision block now. "
+                                   "No prose, no fences."
+        ))
+        resp = self.model.invoke(messages)
+        text = resp.content if hasattr(resp, "content") else str(resp)
+        return parse_json_block(text if isinstance(text, str) else str(text)) or {}
+
+    # ----- budget ----------------------------------------------------------
+
+    def check_budget(self) -> None:
+        """Call before each tool invocation. Raises at the per-event cap."""
+        if self._tool_calls_used >= self.budget:
+            raise BudgetExceeded(
+                f"{self.role} exceeded {self.budget} tool calls per event"
+            )
+        self._tool_calls_used += 1
+
+    def reset_budget(self) -> None:
+        self._tool_calls_used = 0
+
+    # ----- run loop --------------------------------------------------------
+
+    def _on_signal(self, signum: int, _frame: Any) -> None:
+        logger.info("agent.%s received signal %d — draining", self.role, signum)
+        self._stop = True
+
+    def _process_one(self, stream: str, msg_id: str, event: dict[str, Any]) -> None:
+        """Run the handler for one event and ack per the same policy as run()."""
+        try:
+            self.reset_budget()
+            self.handle(stream, event)
+            self.bus.ack(stream, self.consumer_group, msg_id)
+        except BudgetExceeded as exc:
+            logger.warning("agent.%s budget exhausted msg=%s: %s",
+                           self.role, msg_id, exc)
+            self.bus.ack(stream, self.consumer_group, msg_id)
+        except Exception:
+            logger.exception("agent.%s handler error msg=%s", self.role, msg_id)
+
+    def drain(self, max_events: Optional[int] = None) -> int:
+        """Process every event currently available, then return — no blocking.
+
+        The non-looping counterpart to :meth:`run`: it consumes what's on the
+        agent's streams right now and stops once caught up. This is what the
+        offline backtest and the zero-infra demo drive each role with — publish
+        a batch, drain triage, drain Tier 2, and so on down the chain. Returns
+        the number of events processed.
+        """
+        streams = self.streams_to_consume()
+        processed = 0
+        while not self._stop:
+            batch = self.bus.consume_batch(
+                streams, self.consumer_group, self.consumer,
+                batch_size=10, block_ms=0,
+            )
+            if not batch:
+                break
+            for stream, msg_id, event in batch:
+                self._process_one(stream, msg_id, event)
+                processed += 1
+                if max_events is not None and processed >= max_events:
+                    return processed
+        return processed
+
+    def run(self) -> None:
+        """Main loop. Returns on SIGTERM/SIGINT.
+
+        Each event is acked on success or on budget exhaustion (a normal-ish
+        outcome). A handler that raises leaves the event unacked so the bus
+        redelivers it — a crash mid-case never silently drops work.
+        """
+        streams = self.streams_to_consume()
+        logger.info(
+            "agent.%s starting consumer=%s group=%s streams=%s budget=%d",
+            self.role, self.consumer, self.consumer_group, streams, self.budget,
+        )
+
+        while not self._stop:
+            batch = self.bus.consume_batch(
+                streams, self.consumer_group, self.consumer,
+                batch_size=10, block_ms=1000,
+            )
+            for stream, msg_id, event in batch:
+                if self._stop:
+                    break
+                self._process_one(stream, msg_id, event)
+
+        logger.info("agent.%s shut down clean", self.role)
+
+
+__all__ = ["Agent", "BudgetExceeded", "parse_json_block"]