trace-agent-server 0.0.1__py3-none-any.whl

trace_agent_server/services/analytics.py

@@ -0,0 +1,305 @@
+from __future__ import annotations
+
+import re
+from collections import defaultdict
+from typing import Any
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from trace_agent_server.config import Settings
+from trace_agent_server.models import Event, Run
+from trace_agent_server.services.run_summary import build_tool_chain
+
+
+TOKEN_PATTERN = re.compile(r"[a-z0-9]+", re.IGNORECASE)
+STOP_WORDS = {
+    "the",
+    "and",
+    "for",
+    "with",
+    "that",
+    "this",
+    "from",
+    "into",
+    "your",
+    "have",
+    "need",
+    "plan",
+    "build",
+    "using",
+    "use",
+    "agent",
+    "tool",
+    "tools",
+}
+
+
+def _goal_tokens(goal: str) -> set[str]:
+    tokens = {token.lower() for token in TOKEN_PATTERN.findall(goal)}
+    return {token for token in tokens if len(token) > 2 and token not in STOP_WORDS}
+
+
+def _duration_ms(run: Run) -> int | None:
+    if run.ended_at is None:
+        return None
+    return max(0, int((run.ended_at - run.started_at).total_seconds() * 1000))
+
+
+def _extract_usage(event: Event) -> dict[str, int]:
+    usage = event.event_metadata.get("usage", {})
+    if not isinstance(usage, dict):
+        return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+    return {
+        "prompt_tokens": int(usage.get("prompt_tokens") or 0),
+        "completion_tokens": int(usage.get("completion_tokens") or 0),
+        "total_tokens": int(usage.get("total_tokens") or 0),
+    }
+
+
+def _pricing_key(provider: str | None, model: str | None) -> list[str]:
+    keys: list[str] = []
+    if provider and model:
+        keys.append(f"{provider}:{model}")
+    if model:
+        keys.append(model)
+    return keys
+
+
+def build_tool_graph(tool_chain: list[dict[str, Any]]) -> dict[str, Any]:
+    nodes: list[dict[str, Any]] = []
+    edges: list[dict[str, Any]] = []
+    step_to_node: dict[str, str] = {}
+
+    for index, item in enumerate(tool_chain):
+        node_id = f"tool-chain-{item['step_id']}"
+        step_to_node[item["step_id"]] = node_id
+        nodes.append(
+            {
+                "id": node_id,
+                "type": "tool_chain",
+                "label": item["tool_name"],
+                "position": {"x": float(index * 220), "y": 80.0 if index % 2 == 0 else 180.0},
+                "data": {
+                    "status": item["status"],
+                    "durationMs": item["duration_ms"],
+                    "callId": item["call_id"],
+                    "summary": item["error_summary"] or item["output_summary"] or item["args_summary"],
+                    "evidenceIds": item["event_ids"],
+                    "artifactCount": item["artifact_count"],
+                },
+            }
+        )
+
+    for index, item in enumerate(tool_chain):
+        if index == 0:
+            continue
+        source_step = item["parent_step_id"] or tool_chain[index - 1]["step_id"]
+        source_node = step_to_node.get(source_step)
+        target_node = step_to_node[item["step_id"]]
+        if source_node is None:
+            source_node = step_to_node[tool_chain[index - 1]["step_id"]]
+        edges.append(
+            {
+                "id": f"tool-chain-edge-{source_node}-{target_node}",
+                "source": source_node,
+                "target": target_node,
+                "label": item["status"],
+                "animated": item["status"] == "failed",
+            }
+        )
+    return {"nodes": nodes, "edges": edges}
+
+
+def build_tool_metrics(events: list[Event], tool_chain: list[dict[str, Any]], settings: Settings) -> list[dict[str, Any]]:
+    attribution: dict[str, dict[str, int]] = defaultdict(lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
+    for index, event in enumerate(events):
+        if event.type != "model.responded":
+            continue
+        tool_call_count = int(event.event_metadata.get("tool_call_count") or 0)
+        if tool_call_count <= 0:
+            continue
+
+        requested_events: list[Event] = []
+        cursor = index + 1
+        while cursor < len(events) and events[cursor].type == "tool.requested":
+            requested_events.append(events[cursor])
+            cursor += 1
+        if not requested_events:
+            continue
+
+        usage = _extract_usage(event)
+        divisor = max(1, len(requested_events))
+        # Split one model turn across the tool calls it produced.
+        for requested_event in requested_events:
+            key = requested_event.event_metadata.get("call_id") if isinstance(requested_event.event_metadata.get("call_id"), str) else requested_event.step_id
+            if not key:
+                continue
+            attribution[key]["prompt_tokens"] += usage["prompt_tokens"] // divisor
+            attribution[key]["completion_tokens"] += usage["completion_tokens"] // divisor
+            attribution[key]["total_tokens"] += usage["total_tokens"] // divisor
+
+    grouped: dict[str, dict[str, Any]] = {}
+    for item in tool_chain:
+        tool_name = item["tool_name"]
+        metric = grouped.setdefault(
+            tool_name,
+            {
+                "tool_name": tool_name,
+                "call_count": 0,
+                "success_count": 0,
+                "failure_count": 0,
+                "total_duration_ms": 0,
+                "max_duration_ms": 0,
+                "artifact_count": 0,
+                "attributed_prompt_tokens": 0,
+                "attributed_completion_tokens": 0,
+                "attributed_total_tokens": 0,
+                "estimated_cost_usd": 0.0,
+                "providers": set(),
+                "models": set(),
+                "_cost_available": False,
+            },
+        )
+        metric["call_count"] += 1
+        if item["status"] == "succeeded":
+            metric["success_count"] += 1
+        if item["status"] == "failed":
+            metric["failure_count"] += 1
+
+        duration = int(item["duration_ms"] or 0)
+        metric["total_duration_ms"] += duration
+        metric["max_duration_ms"] = max(metric["max_duration_ms"], duration)
+        metric["artifact_count"] += int(item["artifact_count"] or 0)
+        if item["provider"]:
+            metric["providers"].add(item["provider"])
+        if item["model"]:
+            metric["models"].add(item["model"])
+
+        key = item["call_id"] or item["step_id"]
+        usage = attribution.get(key, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
+        metric["attributed_prompt_tokens"] += usage["prompt_tokens"]
+        metric["attributed_completion_tokens"] += usage["completion_tokens"]
+        metric["attributed_total_tokens"] += usage["total_tokens"]
+
+        pricing = None
+        for pricing_key in _pricing_key(item["provider"], item["model"]):
+            if pricing_key in settings.model_pricing:
+                pricing = settings.model_pricing[pricing_key]
+                break
+        if pricing:
+            prompt_rate = float(pricing.get("prompt_per_1k", 0.0))
+            completion_rate = float(pricing.get("completion_per_1k", 0.0))
+            metric["estimated_cost_usd"] += (
+                usage["prompt_tokens"] / 1000.0 * prompt_rate
+                + usage["completion_tokens"] / 1000.0 * completion_rate
+            )
+            metric["_cost_available"] = True
+
+    results: list[dict[str, Any]] = []
+    for metric in grouped.values():
+        call_count = max(1, metric["call_count"])
+        results.append(
+            {
+                "tool_name": metric["tool_name"],
+                "call_count": metric["call_count"],
+                "success_count": metric["success_count"],
+                "failure_count": metric["failure_count"],
+                "total_duration_ms": metric["total_duration_ms"],
+                "avg_duration_ms": int(metric["total_duration_ms"] / call_count),
+                "max_duration_ms": metric["max_duration_ms"],
+                "artifact_count": metric["artifact_count"],
+                "attributed_prompt_tokens": metric["attributed_prompt_tokens"],
+                "attributed_completion_tokens": metric["attributed_completion_tokens"],
+                "attributed_total_tokens": metric["attributed_total_tokens"],
+                "estimated_cost_usd": round(metric["estimated_cost_usd"], 4) if metric["_cost_available"] else None,
+                "providers": sorted(metric["providers"]),
+                "models": sorted(metric["models"]),
+            }
+        )
+    results.sort(key=lambda item: (-item["attributed_total_tokens"], -item["total_duration_ms"], item["tool_name"]))
+    return results
+
+
+def build_similar_runs(
+    session: Session,
+    run: Run,
+    current_tool_chain: list[dict[str, Any]],
+    *,
+    limit: int = 5,
+    same_agent_only: bool = False,
+    status_filter: str | None = None,
+    shared_tools_only: bool = False,
+    min_score: float = 0.15,
+) -> list[dict[str, Any]]:
+    candidates = list(
+        session.scalars(
+            select(Run)
+            .where(Run.id != run.id)
+            .order_by(Run.started_at.desc())
+            .limit(40)
+        )
+    )
+    current_goal_tokens = _goal_tokens(run.goal)
+    current_tools = {item["tool_name"] for item in current_tool_chain}
+    current_scenario = run.run_metadata.get("scenario")
+    results: list[dict[str, Any]] = []
+
+    for candidate in candidates:
+        if same_agent_only and candidate.agent_name != run.agent_name:
+            continue
+        if status_filter and candidate.status != status_filter:
+            continue
+
+        candidate_events = list(session.scalars(select(Event).where(Event.run_id == candidate.id).order_by(Event.seq)))
+        candidate_tool_chain = build_tool_chain(candidate_events)
+        candidate_tools = {item["tool_name"] for item in candidate_tool_chain}
+        candidate_goal_tokens = _goal_tokens(candidate.goal)
+
+        goal_union = current_goal_tokens | candidate_goal_tokens
+        tool_union = current_tools | candidate_tools
+        goal_similarity = len(current_goal_tokens & candidate_goal_tokens) / len(goal_union) if goal_union else 0.0
+        tool_similarity = len(current_tools & candidate_tools) / len(tool_union) if tool_union else 0.0
+        score = 0.45 * goal_similarity + 0.2 * tool_similarity
+
+        reason_parts: list[str] = []
+        if candidate.agent_name == run.agent_name:
+            score += 0.2
+            reason_parts.append("same agent")
+        if current_scenario and candidate.run_metadata.get("scenario") == current_scenario:
+            score += 0.15
+            reason_parts.append("same scenario")
+        shared_tools = sorted(current_tools & candidate_tools)
+        if shared_tools_only and not shared_tools:
+            continue
+        if shared_tools:
+            reason_parts.append(f"shared tools: {', '.join(shared_tools[:3])}")
+        if goal_similarity > 0:
+            reason_parts.append("goal overlap")
+
+        # Keep the score simple and explainable: goal overlap, shared tools and same execution context.
+        if score < min_score:
+            continue
+
+        results.append(
+            {
+                "run_id": candidate.id,
+                "agent_name": candidate.agent_name,
+                "goal": candidate.goal,
+                "status": candidate.status,
+                "started_at": candidate.started_at,
+                "ended_at": candidate.ended_at,
+                "duration_ms": _duration_ms(candidate),
+                "tool_count": candidate.tool_count,
+                "error_count": candidate.error_count,
+                "retry_count": candidate.retry_count,
+                "provider": candidate.provider,
+                "model": candidate.model,
+                "similarity_score": round(min(score, 1.0), 3),
+                "shared_tools": shared_tools,
+                "reason": ", ".join(reason_parts) if reason_parts else "similar recent run",
+            }
+        )
+
+    results.sort(key=lambda item: (item["similarity_score"], item["started_at"]), reverse=True)
+    return results[:limit]

trace_agent_server/services/audit_guardrails.py

@@ -0,0 +1,303 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+from trace_agent_server.config import Settings
+
+
+EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
+PHONE_RE = re.compile(r"\+?[0-9][0-9\s().-]{7,}[0-9]")
+CREDIT_CARD_RE = re.compile(r"(?:\d[ -]*?){13,19}")
+SECRET_TOKEN_RE = re.compile(r"\b(?:sk-[A-Za-z0-9]{20,}|AKIA[0-9A-Z]{16}|ghp_[A-Za-z0-9]{20,})\b")
+
+PROMPT_INJECTION_PATTERNS = (
+    "ignore previous instructions",
+    "ignore all previous instructions",
+    "disregard previous instructions",
+    "reveal the system prompt",
+    "show hidden instructions",
+    "developer message",
+    "system prompt",
+    "act as",
+    "jailbreak",
+    "do anything now",
+)
+
+
+@dataclass(slots=True)
+class GuardrailCheckResult:
+    content: str
+    findings: list[dict[str, Any]]
+    blocked: bool = False
+    block_reason: str | None = None
+
+
+class AuditGuardrails:
+    """Applies best-effort guardrails for prompt injection and PII leakage.
+
+    All integrations are optional. If third-party libraries are not available,
+    lightweight regex heuristics are used instead.
+    """
+
+    def __init__(self, settings: Settings, observability: Any | None = None) -> None:
+        self.enabled = settings.audit_enable_guardrails
+        self.fail_on_prompt_injection = settings.audit_fail_on_prompt_injection
+        self.enable_pii_anonymization = settings.audit_enable_pii_anonymization
+        self.prompt_injection_threshold = max(0.0, min(1.0, settings.audit_prompt_injection_threshold))
+        self.pii_entities = settings.audit_pii_entities
+        self.observability = observability
+
+        self._analyzer = None
+        self._anonymizer = None
+        self._llm_guard_scanner = None
+        self._initialize_optional_engines()
+
+    def sanitize_messages(
+        self,
+        messages: list[dict[str, Any]],
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], bool, str | None]:
+        if not self.enabled:
+            return messages, [], False, None
+
+        sanitized_messages: list[dict[str, Any]] = []
+        findings: list[dict[str, Any]] = []
+        blocked = False
+        block_reason: str | None = None
+
+        for index, message in enumerate(messages):
+            updated = dict(message)
+            content = message.get("content")
+            if isinstance(content, str):
+                pii_check = self._apply_pii_filters(content)
+                injection_findings, injection_score = self._scan_prompt_injection(pii_check.content)
+
+                for finding in pii_check.findings:
+                    findings.append({**finding, "scope": "input", "message_index": index})
+                for finding in injection_findings:
+                    findings.append({**finding, "scope": "input", "message_index": index})
+
+                if self.fail_on_prompt_injection and injection_score >= self.prompt_injection_threshold:
+                    blocked = True
+                    block_reason = "prompt_injection"
+
+                updated["content"] = pii_check.content
+
+            sanitized_messages.append(updated)
+
+        self._record_findings(findings, blocked=blocked, block_reason=block_reason)
+        return sanitized_messages, findings, blocked, block_reason
+
+    def sanitize_assistant_message(
+        self,
+        message: dict[str, Any] | None,
+    ) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
+        if not self.enabled or message is None:
+            return message, []
+
+        updated = dict(message)
+        findings: list[dict[str, Any]] = []
+
+        content = message.get("content")
+        if isinstance(content, str):
+            pii_check = self._apply_pii_filters(content)
+            secret_findings, scrubbed_content = self._scrub_secrets(pii_check.content)
+            updated["content"] = scrubbed_content
+
+            findings.extend({**finding, "scope": "output"} for finding in pii_check.findings)
+            findings.extend({**finding, "scope": "output"} for finding in secret_findings)
+
+        self._record_findings(findings, blocked=False, block_reason=None)
+        return updated, findings
+
+    def _initialize_optional_engines(self) -> None:
+        if self.enable_pii_anonymization:
+            try:
+                from presidio_analyzer import AnalyzerEngine  # type: ignore
+                from presidio_anonymizer import AnonymizerEngine  # type: ignore
+
+                self._analyzer = AnalyzerEngine()
+                self._anonymizer = AnonymizerEngine()
+            except Exception:
+                self._analyzer = None
+                self._anonymizer = None
+
+        if self.enabled:
+            try:
+                from llm_guard.input_scanners import PromptInjection  # type: ignore
+
+                self._llm_guard_scanner = PromptInjection(threshold=self.prompt_injection_threshold)
+            except Exception:
+                self._llm_guard_scanner = None
+
+    def _apply_pii_filters(self, content: str) -> GuardrailCheckResult:
+        if not self.enable_pii_anonymization:
+            return GuardrailCheckResult(content=content, findings=[])
+
+        if self._analyzer is not None and self._anonymizer is not None:
+            return self._apply_presidio(content)
+
+        return self._apply_regex_pii_fallback(content)
+
+    def _apply_presidio(self, content: str) -> GuardrailCheckResult:
+        try:
+            analyzer_results = self._analyzer.analyze(
+                text=content,
+                entities=self.pii_entities or None,
+                language="en",
+            )
+            if not analyzer_results:
+                return GuardrailCheckResult(content=content, findings=[])
+
+            anonymized = self._anonymizer.anonymize(
+                text=content,
+                analyzer_results=analyzer_results,
+            )
+            findings = [
+                {
+                    "kind": "pii",
+                    "severity": "medium",
+                    "source": "presidio",
+                    "entity": result.entity_type,
+                    "score": round(float(result.score), 4),
+                    "start": int(result.start),
+                    "end": int(result.end),
+                    "message": f"Detected {result.entity_type}.",
+                }
+                for result in analyzer_results
+            ]
+            return GuardrailCheckResult(content=anonymized.text, findings=findings)
+        except Exception:
+            return self._apply_regex_pii_fallback(content)
+
+    def _apply_regex_pii_fallback(self, content: str) -> GuardrailCheckResult:
+        findings: list[dict[str, Any]] = []
+        scrubbed = content
+
+        for pattern, replacement, entity in (
+            (EMAIL_RE, "[redacted:email]", "EMAIL_ADDRESS"),
+            (PHONE_RE, "[redacted:phone]", "PHONE_NUMBER"),
+            (CREDIT_CARD_RE, "[redacted:card]", "CREDIT_CARD"),
+        ):
+            matches = list(pattern.finditer(scrubbed))
+            if not matches:
+                continue
+            scrubbed = pattern.sub(replacement, scrubbed)
+            findings.append(
+                {
+                    "kind": "pii",
+                    "severity": "medium",
+                    "source": "regex",
+                    "entity": entity,
+                    "count": len(matches),
+                    "message": f"Detected {len(matches)} {entity} match(es).",
+                }
+            )
+
+        return GuardrailCheckResult(content=scrubbed, findings=findings)
+
+    def _scan_prompt_injection(self, content: str) -> tuple[list[dict[str, Any]], float]:
+        findings: list[dict[str, Any]] = []
+        max_score = 0.0
+
+        normalized = content.lower()
+        matched = [pattern for pattern in PROMPT_INJECTION_PATTERNS if pattern in normalized]
+        if matched:
+            score = min(1.0, 0.35 + (0.15 * len(matched)))
+            max_score = max(max_score, score)
+            findings.append(
+                {
+                    "kind": "prompt_injection",
+                    "severity": "high" if score >= self.prompt_injection_threshold else "medium",
+                    "source": "heuristic",
+                    "score": round(score, 4),
+                    "patterns": matched,
+                    "message": "Detected patterns commonly associated with prompt injection.",
+                }
+            )
+
+        if self._llm_guard_scanner is not None:
+            llm_guard_findings, llm_guard_score = self._scan_with_llm_guard(content)
+            findings.extend(llm_guard_findings)
+            max_score = max(max_score, llm_guard_score)
+
+        return findings, max_score
+
+    def _scan_with_llm_guard(self, content: str) -> tuple[list[dict[str, Any]], float]:
+        try:
+            scan_result = self._llm_guard_scanner.scan(content)
+        except Exception:
+            return [], 0.0
+
+        score = 0.0
+        verdict = "unknown"
+        if isinstance(scan_result, tuple):
+            for value in scan_result:
+                if isinstance(value, bool):
+                    verdict = "safe" if value else "unsafe"
+                if isinstance(value, (int, float)):
+                    score = max(score, float(value))
+                if isinstance(value, dict):
+                    score = max(score, self._extract_score_from_dict(value))
+        elif isinstance(scan_result, dict):
+            score = self._extract_score_from_dict(scan_result)
+            unsafe = scan_result.get("is_valid") is False
+            verdict = "unsafe" if unsafe else "safe"
+
+        if score <= 0:
+            return [], 0.0
+
+        finding = {
+            "kind": "prompt_injection",
+            "severity": "high" if score >= self.prompt_injection_threshold else "medium",
+            "source": "llm_guard",
+            "score": round(score, 4),
+            "message": f"LLM Guard classified prompt as {verdict}.",
+        }
+        return [finding], score
+
+    @staticmethod
+    def _extract_score_from_dict(value: dict[str, Any]) -> float:
+        candidates = (
+            value.get("risk"),
+            value.get("risk_score"),
+            value.get("score"),
+            value.get("confidence"),
+        )
+        for candidate in candidates:
+            if isinstance(candidate, (int, float)):
+                return float(candidate)
+        return 0.0
+
+    def _scrub_secrets(self, content: str) -> tuple[list[dict[str, Any]], str]:
+        matches = list(SECRET_TOKEN_RE.finditer(content))
+        if not matches:
+            return [], content
+
+        scrubbed = SECRET_TOKEN_RE.sub("[redacted:secret]", content)
+        findings = [
+            {
+                "kind": "secret",
+                "severity": "high",
+                "source": "regex",
+                "count": len(matches),
+                "message": "Detected possible API tokens in assistant output.",
+            }
+        ]
+        return findings, scrubbed
+
+    def _record_findings(self, findings: list[dict[str, Any]], *, blocked: bool, block_reason: str | None) -> None:
+        if self.observability is None:
+            return
+
+        try:
+            for finding in findings:
+                self.observability.record_guardrail_finding(
+                    kind=str(finding.get("kind") or "unknown"),
+                    severity=str(finding.get("severity") or "unknown"),
+                )
+            if blocked:
+                self.observability.record_guardrail_block(block_reason or "blocked")
+        except Exception:
+            return