continuous-intelligence-layer 0.1.0__py3-none-any.whl

graph_builder/mongo_store.py

@@ -0,0 +1,588 @@
+"""
+mongo_store.py
+---------------
+Persists and queries an ExecutionGraph in MongoDB.
+
+Mongo Schema
+------------
+
+Collections:
+
+  executions   -- one lightweight header/summary document per execution_id.
+                  Holds status/latency/node_count/timestamps/snippets and the
+                  materialized `edges` array (edges are small, safe to embed).
+  nodes        -- one document per ExecutionNode, keyed by node_id. This is
+                  deliberately NOT embedded inside the execution header doc:
+                  embedding a `nodes` array in one document risks hitting
+                  MongoDB's 16MB per-document BSON limit on large executions
+                  (many nodes, or long LLM prompts/responses/retrieved_docs).
+                  One doc per node keeps each write bounded independently.
+  evaluations  -- one document per EvaluationResult, keyed by evaluation_id.
+                  `node_id` is a plain indexed field (replaces the Neo4j
+                  [:HAS_EVALUATION] edge).
+  rca_reports  -- one document per execution_id (1:1, same as Neo4j's
+                  :RCAReport node).
+
+Idempotent: every write is an upsert keyed on a stable natural id (`_id`),
+mirroring Neo4j's `MERGE ON node_id`.
+
+Useful pymongo queries
+-----------------------
+  # All LLM calls in a session with token costs
+  db.nodes.find({"session_id": "s1", "node_type": "LLM"})
+
+  # All nodes for one execution
+  db.nodes.find({"execution_id": "abc"})
+
+  # All tool failures across sessions
+  db.nodes.find({"node_type": "Tool", "status": "ERROR"})
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+from pymongo import MongoClient
+from pymongo.server_api import ServerApi
+
+from .models import ExecutionEdge, ExecutionGraph, ExecutionNode
+
+
+class MongoStore:
+    """Persists and queries execution graphs in MongoDB."""
+
+    def __init__(self, uri: str | None = None, db_name: str | None = None):
+        self.uri = uri or os.getenv("MONGODB_URI")
+        if not self.uri:
+            raise RuntimeError("MONGODB_URI environment variable is not set.")
+        self._client = MongoClient(self.uri, server_api=ServerApi("1"))
+        self._db = self._client[db_name or os.getenv("MONGODB_DB", "agentops")]
+        self._indexes_ensured = False
+
+    def close(self) -> None:
+        self._client.close()
+
+    # ── Schema ─────────────────────────────────────────────────────────────
+
+    def create_indexes(self) -> None:
+        """Create indexes. Idempotent — safe to call repeatedly, but cheap to
+        call once (e.g. at API startup) rather than on every write."""
+        if self._indexes_ensured:
+            return
+        self._db.executions.create_index("project_id")
+        self._db.executions.create_index("session_id")
+
+        self._db.nodes.create_index("execution_id")
+        self._db.nodes.create_index("project_id")
+        self._db.nodes.create_index("session_id")
+        self._db.nodes.create_index([("project_id", 1), ("execution_id", 1)])
+        self._db.nodes.create_index([("session_id", 1), ("node_type", 1)])
+        self._db.nodes.create_index([("project_id", 1), ("eval_overall_status", 1)])
+        self._db.nodes.create_index("status")
+
+        self._db.evaluations.create_index("node_id")
+        self._db.evaluations.create_index([("execution_id", 1), ("project_id", 1)])
+
+        self._db.rca_reports.create_index([("execution_id", 1), ("project_id", 1)])
+        self._indexes_ensured = True
+
+    # ── Write ──────────────────────────────────────────────────────────────
+
+    def save_graph(self, graph: ExecutionGraph) -> None:
+        """Persist an entire ExecutionGraph to MongoDB (idempotent)."""
+        self.create_indexes()
+        for node in graph.nodes:
+            self._upsert_node(node)
+        self._upsert_execution_header(graph)
+
+    def _upsert_node(self, node: ExecutionNode) -> None:
+        props = {
+            "node_id":          node.node_id,
+            "parent_id":        node.parent_id,
+            "execution_id":     node.execution_id,
+            "session_id":       node.session_id,
+            "project_id":       node.project_id,
+            "trace_id":         node.trace_id,
+            "node_type":        node.node_type,
+            "name":             node.name,
+            # LLM-specific
+            "model":            node.model,
+            # Tool-specific
+            "tool_name":        node.tool_name,
+            "tool_desc":        node.tool_desc,
+            "retrieved_docs":   _safe(node.retrieved_docs),
+            "available_tools":  node.available_tools,
+            "tool_calls":       node.tool_calls,
+            # I/O
+            "input":            _safe(node.input),
+            "output":           _safe(node.output),
+            "prompt":           _safe(node.prompt),
+            "response":         _safe(node.response),
+            # Metrics
+            "latency_ms":           node.latency_ms,
+            "tokens_prompt":        node.tokens.prompt,
+            "tokens_completion":    node.tokens.completion,
+            "tokens_total":         node.tokens.total,
+            # Status
+            "error":            node.error,
+            "status":           node.status,
+            # Temporal
+            "timestamp":        node.timestamp,
+            "end_timestamp":    node.end_timestamp,
+            # Extras
+            "metadata":         json.dumps(node.metadata, default=str)[:2000],
+        }
+        self._db.nodes.update_one(
+            {"_id": node.node_id}, {"$set": props}, upsert=True
+        )
+
+    def _upsert_execution_header(self, graph: ExecutionGraph) -> None:
+        """Write/refresh the lightweight `executions` summary document —
+        avoids the N+1 pattern of re-fetching full node bodies just to
+        render a trace-list row."""
+        nodes = graph.nodes
+        edges = [
+            {"src": e.source_id, "rel": e.edge_type, "tgt": e.target_id}
+            for e in graph.edges
+        ]
+        total_latency = sum(n.latency_ms for n in nodes)
+        status = "Success"
+        if any(n.status == "ERROR" or n.error for n in nodes):
+            status = "Failure"
+        timestamps = [n.timestamp for n in nodes if n.timestamp]
+        end_timestamps = [n.end_timestamp for n in nodes if n.end_timestamp]
+
+        targets = {e.target_id for e in graph.edges}
+        root_nodes = [n for n in nodes if n.node_id not in targets]
+        root = root_nodes[0] if root_nodes else (nodes[0] if nodes else None)
+
+        # ExecutionGraph.project_id/session_id are often left unset by the
+        # builder (it only stamps project_id onto individual nodes) — fall
+        # back to the nodes' own values so the header doc is still filterable.
+        project_id = graph.project_id or next(
+            (n.project_id for n in nodes if n.project_id), None
+        )
+        session_id = graph.session_id or next(
+            (n.session_id for n in nodes if n.session_id), None
+        )
+
+        header = {
+            "project_id":     project_id,
+            "session_id":     session_id,
+            "status":         status,
+            "latency_ms":     round(total_latency, 2),
+            "node_count":     len(nodes),
+            "timestamp":      min(timestamps) if timestamps else None,
+            "end_timestamp":  max(end_timestamps) if end_timestamps else None,
+            "input_snippet":  _safe(root.input) if root else None,
+            "output_snippet": _safe(root.output) if root else None,
+            "edges":          edges,
+        }
+        self._db.executions.update_one(
+            {"_id": graph.execution_id}, {"$set": header}, upsert=True
+        )
+
+    # ── Read ───────────────────────────────────────────────────────────────
+
+    def get_graph(self, execution_id: str, project_id: str | None = None) -> dict:
+        """Return all nodes and edges for a given execution_id."""
+        node_filter: dict = {"execution_id": execution_id}
+        if project_id:
+            node_filter["project_id"] = project_id
+        nodes = list(self._db.nodes.find(node_filter))
+
+        header = self._db.executions.find_one({"_id": execution_id}) or {}
+        edges = header.get("edges", [])
+
+        return {"nodes": nodes, "edges": edges}
+
+    def get_session_graph(self, session_id: str) -> dict:
+        """
+        Return all nodes and edges for an entire session (all executions).
+        A session = one user conversation thread, spanning multiple executions.
+        """
+        nodes = list(self._db.nodes.find({"session_id": session_id}))
+        exec_ids = {n["execution_id"] for n in nodes if n.get("execution_id")}
+        edges: list[dict] = []
+        for eid in exec_ids:
+            header = self._db.executions.find_one({"_id": eid}) or {}
+            edges.extend(header.get("edges", []))
+        return {"nodes": nodes, "edges": edges}
+
+    def get_llm_spans(self, session_id: str) -> list[dict]:
+        """Return all LLM spans for a session — for cost/token analytics."""
+        cursor = self._db.nodes.find(
+            {"session_id": session_id, "node_type": "LLM"},
+            {
+                "name": 1, "model": 1, "tokens_prompt": 1,
+                "tokens_completion": 1, "tokens_total": 1,
+                "latency_ms": 1, "execution_id": 1,
+            },
+        ).sort("timestamp", 1)
+        return list(cursor)
+
+    def get_tool_spans(self, session_id: str) -> list[dict]:
+        """Return all tool call spans for a session — for tool usage analytics."""
+        cursor = self._db.nodes.find(
+            {"session_id": session_id, "node_type": "Tool"},
+            {
+                "tool_name": 1, "input": 1, "output": 1, "latency_ms": 1,
+                "status": 1, "error": 1, "execution_id": 1,
+            },
+        ).sort("timestamp", 1)
+        return list(cursor)
+
+    def get_failed_nodes(self, session_id: str | None = None) -> list[dict]:
+        """Return all ERROR nodes, optionally scoped to a session."""
+        query: dict = {"status": "ERROR"}
+        if session_id:
+            query["session_id"] = session_id
+        cursor = self._db.nodes.find(
+            query,
+            {"name": 1, "node_type": 1, "error": 1, "execution_id": 1, "session_id": 1},
+        )
+        return list(cursor)
+
+    def get_all_execution_ids(self, project_id: str | None = None) -> list[str]:
+        query: dict = {}
+        if project_id:
+            query["project_id"] = project_id
+        return self._db.nodes.distinct("execution_id", query)
+
+    def get_all_session_ids(self) -> list[str]:
+        return self._db.nodes.distinct("session_id", {"session_id": {"$ne": None}})
+
+    # ── Trace list/detail (fixes the traces.py N+1 pattern) ────────────────
+
+    def get_trace_summaries(self, project_id: str) -> list[dict]:
+        """One row per execution, read straight off the `executions` header
+        doc — no per-row node fetch (the N+1 the Neo4j-era code had)."""
+        cursor = self._db.executions.find({"project_id": project_id}).sort(
+            "timestamp", -1
+        )
+        return [
+            {
+                "execution_id":    h["_id"],
+                "status":          h.get("status", "Success"),
+                "latency":         h.get("latency_ms", 0),
+                "has_evaluation":  h.get("has_evaluation", False),
+                "nodes_count":     h.get("node_count", 0),
+                "timestamp":       h.get("timestamp"),
+                "input":           h.get("input_snippet"),
+                "output":          h.get("output_snippet"),
+            }
+            for h in cursor
+        ]
+
+    def get_trace_detail(self, execution_id: str, project_id: str | None = None) -> dict:
+        """Graph + evaluations + RCA report for one execution, in one call."""
+        graph = self.get_graph(execution_id, project_id=project_id)
+
+        eval_filter: dict = {"execution_id": execution_id}
+        if project_id:
+            eval_filter["project_id"] = project_id
+        evals = [
+            {"node_id": e["node_id"], "ev": e}
+            for e in self._db.evaluations.find(eval_filter)
+        ]
+
+        rca_filter: dict = {"_id": execution_id}
+        if project_id:
+            rca_filter["project_id"] = project_id
+        rca = self._db.rca_reports.find_one(rca_filter)
+
+        return {"graph": graph, "evaluations": evals, "rca": rca}
+
+    def count_evaluations(self, execution_id: str, project_id: str | None = None) -> int:
+        query: dict = {"execution_id": execution_id}
+        if project_id:
+            query["project_id"] = project_id
+        return self._db.evaluations.count_documents(query)
+
+    def set_eval_status(self, execution_id: str, status: str, error: str | None = None) -> None:
+        """Set run-level evaluation status ('running'/'done'/'error') on the
+        execution header doc, so the API can report progress independent of
+        whichever request/background task is actually doing the work."""
+        self._db.executions.update_one(
+            {"_id": execution_id},
+            {"$set": {"eval_status": status, "eval_error": error}},
+        )
+
+    def get_eval_status(self, execution_id: str, project_id: str | None = None) -> dict:
+        query: dict = {"_id": execution_id}
+        if project_id:
+            query["project_id"] = project_id
+        doc = self._db.executions.find_one(query, {"eval_status": 1, "eval_error": 1}) or {}
+        return {
+            "status": doc.get("eval_status", "not_started"),
+            "error": doc.get("eval_error"),
+        }
+
+    # ── KPI aggregation ───────────────────────────────────────────────────────
+    # All queries here are index-backed aggregation pipelines over the
+    # `nodes` collection's flat evaluation props — never a full graph fetch.
+
+    def _kpi_aggregate(self, match: dict) -> dict:
+        nodes = self._db.nodes
+
+        total = len(nodes.distinct("execution_id", match))
+
+        evaluated = len(
+            nodes.distinct(
+                "execution_id", {**match, "eval_overall_status": {"$ne": None}}
+            )
+        )
+
+        breakdown_rows = list(
+            nodes.aggregate(
+                [
+                    {"$match": {**match, "eval_overall_status": {"$ne": None}}},
+                    {
+                        "$group": {
+                            "_id": "$execution_id",
+                            "statuses": {"$addToSet": "$eval_overall_status"},
+                        }
+                    },
+                    {
+                        "$project": {
+                            "trace_status": {
+                                "$switch": {
+                                    "branches": [
+                                        {
+                                            "case": {"$in": ["FAIL", "$statuses"]},
+                                            "then": "FAIL",
+                                        },
+                                        {
+                                            "case": {"$in": ["WARNING", "$statuses"]},
+                                            "then": "WARNING",
+                                        },
+                                    ],
+                                    "default": "PASS",
+                                }
+                            }
+                        }
+                    },
+                    {"$group": {"_id": "$trace_status", "c": {"$sum": 1}}},
+                ]
+            )
+        )
+        breakdown = {"PASS": 0, "WARNING": 0, "FAIL": 0}
+        for row in breakdown_rows:
+            breakdown[row["_id"]] = row["c"]
+
+        tool_row = next(
+            nodes.aggregate(
+                [
+                    {"$match": {**match, "tool_selection_correct": {"$ne": None}}},
+                    {
+                        "$group": {
+                            "_id": None,
+                            "total": {"$sum": 1},
+                            "correct": {
+                                "$sum": {"$cond": ["$tool_selection_correct", 1, 0]}
+                            },
+                        }
+                    },
+                ]
+            ),
+            {"total": 0, "correct": 0},
+        )
+
+        hallucination = len(
+            nodes.distinct("execution_id", {**match, "has_hallucination": True})
+        )
+        injection = len(
+            nodes.distinct("execution_id", {**match, "has_injection": True})
+        )
+        critical = len(
+            nodes.distinct("execution_id", {**match, "eval_severity": "CRITICAL"})
+        )
+
+        latency_row = next(
+            nodes.aggregate(
+                [
+                    {"$match": match},
+                    {
+                        "$group": {
+                            "_id": "$execution_id",
+                            "total_latency": {"$sum": "$latency_ms"},
+                        }
+                    },
+                    {"$group": {"_id": None, "avg_latency": {"$avg": "$total_latency"}}},
+                ]
+            ),
+            None,
+        )
+        latency = latency_row["avg_latency"] if latency_row else None
+
+        tool_total = tool_row["total"] or 0
+        tool_correct = tool_row["correct"] or 0
+
+        return {
+            "total_traces": total,
+            "evaluated_traces": evaluated,
+            "unevaluated_traces": total - evaluated,
+            "trace_status_breakdown": breakdown,
+            "tool_selection_accuracy_pct": (
+                round(100 * tool_correct / tool_total, 1) if tool_total else None
+            ),
+            "tool_selection_judged_count": tool_total,
+            "hallucination_count": hallucination,
+            "injection_count": injection,
+            "critical_failure_count": critical,
+            "avg_latency_ms": round(latency, 1) if latency is not None else None,
+        }
+
+    def get_project_kpis(self, project_id: str) -> dict:
+        """Aggregate KPIs for a single project."""
+        return self._kpi_aggregate({"project_id": project_id})
+
+    def get_global_kpis(self, project_ids: list[str]) -> dict:
+        """Aggregate KPIs across all of the given projects (e.g. one user's projects)."""
+        return self._kpi_aggregate({"project_id": {"$in": project_ids}})
+
+    # ── Validator Engine integration ────────────────────────────────────────
+
+    def update_node_validation(self, node_id: str, validation_results: list[dict]) -> None:
+        """Attach validator results to a node after validation pass (legacy)."""
+        self._db.nodes.update_one(
+            {"_id": node_id},
+            {"$set": {"validation_results": json.dumps(validation_results, default=str)}},
+        )
+
+    def write_evaluation_result(self, result: object, project_id: str | None = None) -> None:
+        """
+        Persist a single EvaluationResult document, keyed by evaluation_id,
+        with `node_id` as a plain indexed field (replaces Neo4j's
+        [:HAS_EVALUATION] edge). `project_id` is accepted explicitly since
+        EvaluationResult itself doesn't carry one — it's needed so
+        get_trace_detail/count_evaluations can scope by tenant without a
+        join back to the node document.
+        """
+        props = {
+            "evaluation_id": result.evaluation_id,
+            "evaluator":     result.evaluator,
+            "node_id":       result.node_id,
+            "node_name":     result.node_name,
+            "node_type":     result.node_type,
+            "execution_id":  result.execution_id,
+            "project_id":    project_id,
+            "session_id":    result.session_id,
+            "status":        str(result.status.value if hasattr(result.status, 'value') else result.status),
+            "severity":      str(result.severity.value if result.severity and hasattr(result.severity, 'value') else result.severity) if result.severity else None,
+            "confidence":    result.confidence,
+            "reason":        result.reason[:2000] if result.reason else "",
+            "suggestion":    result.suggestion,
+            "checks":        json.dumps(result.checks, default=str)[:4000],
+            "timestamp":     result.timestamp,
+            "metadata":      json.dumps(result.metadata or {}, default=str)[:2000],
+        }
+        self._db.evaluations.update_one(
+            {"_id": result.evaluation_id}, {"$set": props}, upsert=True
+        )
+
+    def update_node_evaluation(
+        self,
+        node_id: str,
+        flat_props: dict,
+        all_results: list[dict],
+    ) -> None:
+        """
+        Write flat boolean evaluation properties directly onto the node
+        document for fast filtering, and store the full results JSON blob.
+        Also denormalizes `has_evaluation`/`status` onto the execution header
+        doc (escalating Success -> Warning -> Failure, never downgrading) so
+        `get_trace_summaries` never needs a node fetch to render the list.
+        """
+        update = dict(flat_props)
+        update["evaluation_results"] = _safe_truncate_json_list(all_results)
+        node = self._db.nodes.find_one_and_update(
+            {"_id": node_id}, {"$set": update}
+        )
+        if node and node.get("execution_id"):
+            eid = node["execution_id"]
+            overall = flat_props.get("eval_overall_status")
+            header = self._db.executions.find_one({"_id": eid}) or {}
+            priority = {"Success": 0, "Warning": 1, "Failure": 2}
+            new_status = header.get("status", "Success")
+            if overall == "FAIL":
+                new_status = "Failure"
+            elif overall == "WARNING" and priority.get(new_status, 0) < priority["Warning"]:
+                new_status = "Warning"
+            self._db.executions.update_one(
+                {"_id": eid}, {"$set": {"has_evaluation": True, "status": new_status}}
+            )
+
+    def write_rca_result(self, rca: dict) -> None:
+        """Persist the RCA Result to MongoDB as an `rca_reports` document."""
+        if not rca.get("execution_id"):
+            return
+
+        props = {
+            "execution_id": rca["execution_id"],
+            "project_id": rca.get("project_id"),
+            "overall_status": rca.get("overall_status", "UNKNOWN"),
+            "root_cause": rca.get("root_cause", ""),
+            "confidence": rca.get("confidence", 0.0),
+            "propagation_chain": json.dumps(rca.get("propagation_chain", [])),
+            "contributing_factors": json.dumps(rca.get("contributing_factors", [])),
+            "evidence": json.dumps(rca.get("evidence", []), default=str)[:8000],
+            "recommendations": json.dumps(rca.get("recommendations", [])),
+            "raw_llm_analysis": rca.get("raw_llm_analysis", "")[:4000],
+        }
+        self._db.rca_reports.update_one(
+            {"_id": rca["execution_id"]}, {"$set": props}, upsert=True
+        )
+
+    # ── Maintenance ────────────────────────────────────────────────────────
+
+    def clear_execution(self, execution_id: str) -> None:
+        """Delete all nodes, evaluations, RCA report and the header doc for
+        a given execution_id."""
+        with self._client.start_session() as session:
+            def _do_delete(s):
+                self._db.nodes.delete_many({"execution_id": execution_id}, session=s)
+                self._db.evaluations.delete_many({"execution_id": execution_id}, session=s)
+                self._db.rca_reports.delete_many({"_id": execution_id}, session=s)
+                self._db.executions.delete_many({"_id": execution_id}, session=s)
+            session.with_transaction(_do_delete)
+
+    def clear_session(self, session_id: str) -> None:
+        """Delete all nodes across all executions in a session."""
+        exec_ids = self._db.nodes.distinct("execution_id", {"session_id": session_id})
+        with self._client.start_session() as session:
+            def _do_delete(s):
+                self._db.nodes.delete_many({"session_id": session_id}, session=s)
+                self._db.evaluations.delete_many({"session_id": session_id}, session=s)
+                if exec_ids:
+                    self._db.rca_reports.delete_many({"_id": {"$in": exec_ids}}, session=s)
+                    self._db.executions.delete_many({"_id": {"$in": exec_ids}}, session=s)
+            session.with_transaction(_do_delete)
+
+
+def _safe_truncate_json_list(items: list[dict], max_len: int = 16000) -> str:
+    """
+    Serialize a list of dicts to JSON, dropping trailing items (never slicing
+    the string mid-object) so the result always parses as valid JSON on the
+    frontend, even if the full list would exceed max_len.
+    """
+    remaining = list(items)
+    while remaining:
+        s = json.dumps(remaining, default=str)
+        if len(s) <= max_len:
+            return s
+        remaining = remaining[:-1]
+    return "[]"
+
+
+def _safe(val: Any, max_len: int = 128000) -> str | None:
+    if val is None:
+        return None
+    if isinstance(val, str):
+        return val[:max_len]
+    try:
+        return json.dumps(val, default=str)[:max_len]
+    except Exception:
+        return str(val)[:max_len]

llm_router/__init__.py

@@ -0,0 +1,3 @@
+from .router import LLMRouter
+
+__all__ = ["LLMRouter"]

llm_router/router.py

@@ -0,0 +1,71 @@
+"""
+router.py
+---------
+Dispatches evaluator/RCA LLM calls to whichever provider (openai, anthropic,
+gemini) the user supplied a key for at ``init()`` time. Every evaluator and
+the RCA engine talk to the LLM exclusively through this class -- nothing
+else in this codebase should build an OpenAI/Anthropic/Gemini client of its
+own for evaluation purposes.
+"""
+
+from __future__ import annotations
+
+SUPPORTED_PROVIDERS = ("openai", "anthropic", "gemini")
+
+
+class LLMRouter:
+    """Wraps a single (provider, model, api_key) triple and exposes one
+    uniform ``call(prompt)`` method."""
+
+    def __init__(self, provider: str, model: str, api_key: str):
+        if not api_key:
+            raise ValueError("LLMRouter requires a non-empty api_key.")
+        provider = (provider or "openai").lower()
+        if provider not in SUPPORTED_PROVIDERS:
+            raise ValueError(
+                f"Unsupported LLM provider: {provider!r}. "
+                f"Supported providers: {', '.join(SUPPORTED_PROVIDERS)}."
+            )
+        self.provider = provider
+        self.model = model
+        self.api_key = api_key
+
+    def call(self, prompt: str, json_mode: bool = False) -> str:
+        if self.provider == "openai":
+            return self._call_openai(prompt, json_mode)
+        if self.provider == "anthropic":
+            return self._call_anthropic(prompt, json_mode)
+        return self._call_gemini(prompt, json_mode)
+
+    def _call_openai(self, prompt: str, json_mode: bool) -> str:
+        from openai import OpenAI
+
+        client = OpenAI(api_key=self.api_key)
+        kwargs = {"response_format": {"type": "json_object"}} if json_mode else {}
+        resp = client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            **kwargs,
+        )
+        return resp.choices[0].message.content
+
+    def _call_anthropic(self, prompt: str, json_mode: bool) -> str:
+        import anthropic
+
+        client = anthropic.Anthropic(api_key=self.api_key)
+        if json_mode:
+            prompt = f"{prompt}\n\nRespond with valid JSON only, no markdown fences."
+        resp = client.messages.create(
+            model=self.model,
+            max_tokens=4096,
+            messages=[{"role": "user", "content": prompt}],
+        )
+        return resp.content[0].text
+
+    def _call_gemini(self, prompt: str, json_mode: bool) -> str:
+        from google import genai
+
+        client = genai.Client(api_key=self.api_key)
+        config = {"response_mime_type": "application/json"} if json_mode else None
+        resp = client.models.generate_content(model=self.model, contents=prompt, config=config)
+        return resp.text

rca_engine/__init__.py

@@ -0,0 +1,5 @@
+"""RCA Engine package."""
+from .rca_engine import RCAEngine, RCAResult
+from .incident_report import IncidentReportGenerator
+
+__all__ = ["RCAEngine", "RCAResult", "IncidentReportGenerator"]