dapplepot-sdk 0.1.0__py3-none-any.whl

dapplepot_sdk/__init__.py

@@ -0,0 +1,243 @@
+import logging
+import random
+
+import requests
+
+from dapplepot_sdk._adapter import TraceAdapter
+from dapplepot_sdk._buffer import EventBuffer
+from dapplepot_sdk._interceptor import OnlineCheckInterceptor
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+logger = logging.getLogger(__name__)
+
+_ONLINE_CAPABLE_SUB_CHECKS: frozenset[str] = frozenset({
+    'PI-01a', 'PI-01b', 'PI-01c', 'PI-02a', 'PI-05a', 'PI-08a',
+    'SID-01a', 'SID-01c', 'SID-02a',
+    'IOH-01a',
+    'EA-01a', 'EA-02b',
+})
+
+
+class DapplePotBlockedError(Exception):
+    def __init__(self, signal: str, reason: str, session_id: str):
+        super().__init__(f'[{signal}] {reason}')
+        self.signal = signal
+        self.reason = reason
+        self.session_id = session_id
+
+
+class DapplePotSessionTerminatedError(Exception):
+    pass
+
+
+class DapplePot:
+    def __init__(
+        self,
+        sdk_key: str,
+        tenant_id: str,
+        agent_id: str,
+        ingest_url: str,
+        *,
+        sample_rate: float = 1.0,
+        pii_scrubber=None,
+        redact_keys: list = None,
+        flush_interval_ms: int = 500,
+        flush_batch_size: int = 100,
+    ):
+        self._sdk_key      = sdk_key
+        self._tenant_id    = tenant_id
+        self._agent_id     = agent_id
+        self._ingest_url   = ingest_url.rstrip('/')
+        self._sample_rate  = sample_rate
+        self._pii_scrubber = pii_scrubber
+        self._redact_keys  = set(redact_keys or [])
+        self._tool_allowlist = None
+
+        check_actions = self._fetch_check_actions()
+
+        self._buffer = EventBuffer(
+            ingest_url=self._ingest_url,
+            sdk_key=self._sdk_key,
+            flush_interval_ms=flush_interval_ms,
+            flush_batch_size=flush_batch_size,
+        )
+        self._interceptor = OnlineCheckInterceptor(
+            check_actions=check_actions,
+            buffer=self._buffer,
+            client=self,
+        )
+
+        tool_manifest, max_tool_calls = self._fetch_tool_manifest()
+        ea01a_action = check_actions.get('EA-01a', 'block_call')
+        ea02b_action = check_actions.get('EA-02b', 'alert')
+        self._interceptor.set_tool_manifest(
+            manifest=tool_manifest,
+            action=ea01a_action,
+            max_tool_calls=max_tool_calls,
+            ea02b_action=ea02b_action,
+        )
+
+        self._framework = 'unknown'
+
+    # ── startup ───────────────────────────────────────────────────────────────
+
+    def _fetch_check_actions(self) -> dict[str, str]:
+        """Pull per-subcheck online config from the API. Returns {sub_check_id: action}."""
+        url = f'{self._ingest_url}/v1/sdk/security/agents/{self._agent_id}/subcheck-config'
+        try:
+            resp = requests.get(
+                url,
+                headers={'Authorization': f'Bearer {self._sdk_key}'},
+                timeout=5,
+            )
+            resp.raise_for_status()
+            overrides: dict = resp.json().get('overrides', {})
+            check_actions = {}
+            for sub_check_id, cfg in overrides.items():
+                if cfg.get('online_detection') and sub_check_id in _ONLINE_CAPABLE_SUB_CHECKS:
+                    check_actions[sub_check_id] = cfg.get('action', 'alert')
+            logger.debug('Loaded %d online checks from API', len(check_actions))
+            return check_actions
+        except Exception as exc:
+            logger.warning(
+                'Could not fetch online check config from %s: %s — online checks disabled',
+                url, exc,
+            )
+            return {}
+
+    def _fetch_tool_manifest(self) -> tuple[list[str], int | None]:
+        """Fetch the tool manifest and max_tool_calls_per_session from the API."""
+        url = f'{self._ingest_url}/v1/sdk/security/agents/{self._agent_id}/tool-manifest'
+        try:
+            resp = requests.get(
+                url,
+                headers={'Authorization': f'Bearer {self._sdk_key}'},
+                timeout=5,
+            )
+            resp.raise_for_status()
+            data = resp.json()
+            manifest = data.get('tool_manifest') or []
+            max_calls = data.get('max_tool_calls_per_session')
+            logger.debug('Loaded tool manifest (%d tools, max_calls=%s)', len(manifest), max_calls)
+            return manifest, max_calls
+        except Exception as exc:
+            logger.warning(
+                'Could not fetch tool manifest from %s: %s — EA-01a/EA-02b disabled',
+                url, exc,
+            )
+            return [], None
+
+    # ── internal helpers ──────────────────────────────────────────────────────
+
+    def _adapter(self, framework: str) -> TraceAdapter:
+        return TraceAdapter(
+            tenant_id=self._tenant_id,
+            agent_id=self._agent_id,
+            framework=framework,
+        )
+
+    def _should_sample(self) -> bool:
+        return random.random() < self._sample_rate
+
+    def _scrub(self, event: dict) -> dict:
+        if not self._pii_scrubber and not self._redact_keys:
+            return event
+        payload = event.get('payload', {})
+        if self._pii_scrubber:
+            payload = self._pii_scrubber.scrub_value(payload)
+        if self._redact_keys:
+            payload = self._redact_keys_in(payload)
+        return {**event, 'payload': payload}
+
+    def _redact_keys_in(self, obj):
+        if isinstance(obj, dict):
+            return {k: '[REDACTED]' if k in self._redact_keys else self._redact_keys_in(v)
+                    for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [self._redact_keys_in(i) for i in obj]
+        return obj
+
+    def _process_event(self, event: dict) -> None:
+        event = self._interceptor.evaluate(event)
+        event = self._scrub(event)
+        self._buffer.push(event)
+
+    # ── public API ────────────────────────────────────────────────────────────
+
+    def instrument_anthropic(self) -> None:
+        """Patch the Anthropic SDK so all messages.create() calls are traced automatically.
+
+        Call once after DapplePot() is initialised, before creating your Anthropic client.
+        The standard anthropic package is unaffected — upgrade it freely at any time.
+
+        Usage::
+
+            import anthropic
+            from dapplepot_sdk import DapplePot
+
+            dp = DapplePot(...)
+            dp.instrument_anthropic()
+
+            client = anthropic.Anthropic(api_key="...")
+        """
+        from dapplepot_sdk import anthropic as _dp_anth
+        _dp_anth._patch(self)
+        self._framework = 'anthropic'
+
+    def instrument_openai(self) -> None:
+        """Patch the OpenAI SDK so all chat.completions.create() calls are traced automatically.
+
+        Call once after DapplePot() is initialised, before creating your OpenAI client.
+
+        Usage::
+
+            import openai
+            from dapplepot_sdk import DapplePot
+
+            dp = DapplePot(...)
+            dp.instrument_openai()
+
+            client = openai.OpenAI(api_key="...")
+        """
+        from dapplepot_sdk import openai as _dp_openai
+        _dp_openai._patch(self)
+        self._framework = 'openai'
+
+    def callback_handler(self, session_id: str = None, user_context_id: str = None,
+                         user_tenant_id: str = None):
+        """Return a fresh LangChain/LangGraph CallbackHandler for one session."""
+        from dapplepot_sdk._langchain import DapplePotCallbackHandler
+        self._framework = 'langchain'
+        return DapplePotCallbackHandler(self, session_id=session_id,
+                                        user_context_id=user_context_id,
+                                        user_tenant_id=user_tenant_id)
+
+    def session(self, session_id: str = None, user_context_id: str = None,
+                user_tenant_id: str = None):
+        """Context manager that wraps OpenAI / Anthropic calls in a DapplePot session."""
+        from dapplepot_sdk.session import SessionContext
+        return SessionContext(self, session_id=session_id, user_context_id=user_context_id,
+                              user_tenant_id=user_tenant_id)
+
+    def node(self, node_name: str, input=None):
+        """Context manager to trace a named step inside an active dp.session().
+
+        Emits node_start on enter and node_end / node_error on exit.
+        Use this to add structure to your agent loop — it is entirely optional.
+
+        Usage::
+
+            with dp.node("retrieval", input=query):
+                docs = vector_store.search(query)
+
+            with dp.node("call_model"):
+                response = client.messages.create(...)
+        """
+        from dapplepot_sdk._node_context import NodeContext
+        from dapplepot_sdk.session import get_current_session_id
+        return NodeContext(self, session_id=get_current_session_id(),
+                           node_name=node_name, input=input)
+
+    def shutdown(self, timeout_ms: int = 5000) -> None:
+        """Flush remaining events and stop background threads."""
+        self._buffer.shutdown(timeout_ms=timeout_ms)

dapplepot_sdk/_adapter.py

@@ -0,0 +1,168 @@
+import datetime
+import uuid
+
+
+def _now() -> str:
+    return datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3] + 'Z'
+
+
+def first_user_text(messages: list) -> "str | None":
+    """Extract the first user message text from a standard messages list."""
+    for msg in messages:
+        if not isinstance(msg, dict):
+            continue
+        if msg.get("role") in ("user", "human"):
+            content = msg.get("content", "")
+            if isinstance(content, str) and content:
+                return content
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        t = block.get("text", "")
+                        if t:
+                            return t
+    return None
+
+
+class TraceAdapter:
+    def __init__(self, tenant_id: str, agent_id: str, framework: str):
+        self._tenant_id = tenant_id
+        self._agent_id = agent_id
+        self._framework = framework
+
+    def _base(self, session_id: str, event_type: str) -> dict:
+        return {
+            'dp_tenant_id':      self._tenant_id,
+            'dp_agent_id':       self._agent_id,
+            'dp_session_id':     session_id,
+            'dp_event_type':     event_type,
+            'dp_schema_version': '2',
+            'dp_sampled':        True,
+            'dp_framework':      self._framework,
+            'ts':                _now(),
+            'event_id':          str(uuid.uuid4()),
+            'payload':           {},
+        }
+
+    def session_start(self, session_id: str, user_context_id: str = None, metadata=None,
+                      input=None, user_tenant_id: str = None) -> dict:
+        e = self._base(session_id, 'session_start')
+        if user_context_id:
+            e['user_context_id'] = user_context_id
+        if user_tenant_id:
+            e['user_tenant_id'] = user_tenant_id
+        e['payload'] = {'session_id': session_id, 'framework': self._framework, 'agent_id': self._agent_id}
+        if user_context_id:
+            e['payload']['user_context_id'] = user_context_id
+        if user_tenant_id:
+            e['payload']['user_tenant_id'] = user_tenant_id
+        if metadata:
+            e['payload']['metadata'] = metadata
+        if input is not None:
+            e['payload']['input'] = input
+        return e
+
+    def session_end(self, session_id: str, output=None, latency_ms=None, total_tokens=None) -> dict:
+        e = self._base(session_id, 'session_end')
+        e['payload'] = {}
+        if output is not None:
+            e['payload']['output'] = output
+        if latency_ms is not None:
+            e['payload']['latency_ms'] = latency_ms
+        if total_tokens is not None:
+            e['payload']['total_tokens'] = total_tokens
+        return e
+
+    def session_error(self, session_id: str, error_type: str, error_message: str,
+                      traceback: str = None, exit_reason: str = None) -> dict:
+        e = self._base(session_id, 'session_error')
+        e['payload'] = {'error_type': error_type, 'error_message': error_message}
+        if traceback:
+            e['payload']['traceback'] = traceback
+        if exit_reason:
+            e['payload']['exit_reason'] = exit_reason
+        return e
+
+    def node_start(self, session_id: str, node_name: str, parent_span_id=None, input=None) -> dict:
+        e = self._base(session_id, 'node_start')
+        e['payload'] = {'node_name': node_name}
+        if parent_span_id:
+            e['payload']['parent_span_id'] = parent_span_id
+        if input is not None:
+            e['payload']['input'] = input
+        return e
+
+    def node_end(self, session_id: str, node_name: str, output=None, latency_ms=None) -> dict:
+        e = self._base(session_id, 'node_end')
+        e['payload'] = {'node_name': node_name}
+        if output is not None:
+            e['payload']['output'] = output
+        if latency_ms is not None:
+            e['payload']['latency_ms'] = latency_ms
+        return e
+
+    def node_error(self, session_id: str, node_name: str, error_type: str, error_message: str, traceback: str = None) -> dict:
+        e = self._base(session_id, 'node_error')
+        e['payload'] = {'node_name': node_name, 'error_type': error_type, 'error_message': error_message}
+        if traceback:
+            e['payload']['traceback'] = traceback
+        return e
+
+    def llm_start(self, session_id: str, model: str, messages: list,
+                  temperature=None, max_tokens=None, tools=None) -> dict:
+        e = self._base(session_id, 'llm_start')
+        e['payload'] = {'model': model, 'messages': messages}
+        if temperature is not None:
+            e['payload']['temperature'] = temperature
+        if max_tokens is not None:
+            e['payload']['max_tokens'] = max_tokens
+        if tools:
+            e['payload']['tools'] = tools
+        return e
+
+    def llm_end(self, session_id: str, completion: str, model: str | None = None, finish_reason=None, usage=None, latency_ms=None) -> dict:
+        e = self._base(session_id, 'llm_end')
+        e['payload'] = {'completion': completion}
+        if model:
+            e['payload']['model'] = model
+        if finish_reason:
+            e['payload']['finish_reason'] = finish_reason
+        if usage:
+            e['payload']['usage'] = usage
+        if latency_ms is not None:
+            e['payload']['latency_ms'] = latency_ms
+        return e
+
+    def llm_error(self, session_id: str, model: str | None = None, error_type: str = None,
+                  error_message: str = None, latency_ms: int = None) -> dict:
+        e = self._base(session_id, 'llm_error')
+        e['payload'] = {}
+        if model:
+            e['payload']['model'] = model
+        if error_type:
+            e['payload']['error_type'] = error_type
+        if error_message:
+            e['payload']['error_message'] = error_message
+        if latency_ms is not None:
+            e['payload']['latency_ms'] = latency_ms
+        return e
+
+    def tool_start(self, session_id: str, tool_name: str, tool_input) -> dict:
+        e = self._base(session_id, 'tool_start')
+        e['payload'] = {'tool_name': tool_name, 'tool_input': tool_input}
+        return e
+
+    def tool_error(self, session_id: str, tool_name: str, error_message: str,
+                   error_type: str = "ToolError", tool_input=None) -> dict:
+        e = self._base(session_id, 'tool_error')
+        e['payload'] = {'tool_name': tool_name, 'error_type': error_type, 'error_message': error_message}
+        if tool_input is not None:
+            e['payload']['tool_input'] = tool_input
+        return e
+
+    def tool_end(self, session_id: str, tool_name: str, tool_output, latency_ms=None) -> dict:
+        e = self._base(session_id, 'tool_end')
+        e['payload'] = {'tool_name': tool_name, 'tool_output': tool_output}
+        if latency_ms is not None:
+            e['payload']['latency_ms'] = latency_ms
+        return e

dapplepot_sdk/_buffer.py

@@ -0,0 +1,115 @@
+import atexit
+import json
+import queue
+import threading
+import time
+import logging
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+class _SafeEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if hasattr(obj, 'model_dump'):
+            return obj.model_dump()
+        if hasattr(obj, 'dict'):
+            return obj.dict()
+        if hasattr(obj, '__dict__'):
+            return obj.__dict__
+        return repr(obj)
+
+
+class EventBuffer:
+    def __init__(self, ingest_url: str, sdk_key: str,
+                 flush_interval_ms: int = 500, flush_batch_size: int = 100):
+        self._url = ingest_url.rstrip('/') + '/v1/ingest/events'
+        self._sdk_key = sdk_key
+        self._interval = flush_interval_ms / 1000.0
+        self._batch_size = flush_batch_size
+        self._queue: queue.Queue = queue.Queue()
+        self._session_samples: dict = {}
+        self._session_seqs: dict = {}
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._loop, daemon=True, name='telemetry')
+        self._thread.start()
+        # Flush remaining events on normal process exit so the daemon thread
+        # dying at interpreter shutdown doesn't silently drop buffered events.
+        atexit.register(self.shutdown)
+
+    # ── sampling ──────────────────────────────────────────────────────────────
+
+    def set_sampled(self, session_id: str, sampled: bool) -> None:
+        self._session_samples[session_id] = sampled
+
+    def is_sampled(self, session_id: str) -> bool:
+        return self._session_samples.get(session_id, True)
+
+    # ── push ──────────────────────────────────────────────────────────────────
+
+    def push(self, event: dict) -> None:
+        sid = event.get('dp_session_id')
+        if sid and not self.is_sampled(sid):
+            return
+        # Stamp a monotonic per-session sequence_index for SDK paths that don't
+        # set one (openai, anthropic, session.py). LangChain sets it in _emit()
+        # before calling here, so we leave those events unchanged.
+        if 'sequence_index' not in event and sid:
+            n = self._session_seqs.get(sid, 0)
+            self._session_seqs[sid] = n + 1
+            event = {**event, 'sequence_index': n}
+        self._queue.put(event)
+
+    def push_sync(self, event: dict) -> None:
+        """Flush a single event immediately (used before raising a blocked error)."""
+        self._send([event])
+
+    def flush_sync(self) -> None:
+        """Drain all queued events immediately on the calling thread.
+        Safe to call after shutdown() — push() always enqueues regardless of stop state."""
+        self._drain()
+
+    # ── flush loop ────────────────────────────────────────────────────────────
+
+    def _loop(self) -> None:
+        while not self._stop.is_set():
+            time.sleep(self._interval)
+            self._drain()
+
+    def _drain(self) -> None:
+        batch = []
+        try:
+            while len(batch) < self._batch_size:
+                batch.append(self._queue.get_nowait())
+        except queue.Empty:
+            pass
+        if batch:
+            self._send(batch)
+
+    def _send(self, batch: list, retries: int = 3) -> None:
+        for attempt in range(retries):
+            try:
+                resp = requests.post(
+                    self._url,
+                    data=json.dumps({'events': batch}, cls=_SafeEncoder),
+                    headers={
+                        'Authorization': f'Bearer {self._sdk_key}',
+                        'Content-Type': 'application/json',
+                    },
+                    timeout=5,
+                )
+                resp.raise_for_status()
+                return
+            except Exception as exc:
+                if attempt == retries - 1:
+                    logger.error('event flush failed after %d retries: %s', retries, exc)
+                else:
+                    time.sleep(0.1 * (2 ** attempt))
+
+    # ── shutdown ──────────────────────────────────────────────────────────────
+
+    def shutdown(self, timeout_ms: int = 5000) -> None:
+        self._stop.set()
+        self._thread.join(timeout=timeout_ms / 1000.0)
+        self._drain()