atlas-chat 0.1.0__py3-none-any.whl

atlas/domain/rag_mcp_service.py

@@ -0,0 +1,664 @@
+"""RAG MCP Aggregator Service (Phase 1: Discovery)
+
+Aggregates discovery of RAG resources from authorized MCP servers that expose
+the `rag_discover_resources` tool. Returns a flat list of data source IDs for
+backward-compatible UI, with server-qualified IDs to avoid collisions.
+
+Future phases will add search/synthesis and richer shapes.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from atlas.core.compliance import get_compliance_manager
+from atlas.core.log_sanitizer import sanitize_for_logging
+from atlas.core.prompt_risk import calculate_prompt_injection_risk, log_high_risk_event
+
+logger = logging.getLogger(__name__)
+
+
+class RAGMCPService:
+    """Aggregator for RAG over MCP servers."""
+
+    def __init__(self, mcp_manager, config_manager, auth_check_func) -> None:
+        self.mcp_manager = mcp_manager
+        self.config_manager = config_manager
+        self.auth_check_func = auth_check_func
+
+    async def _get_authorized_rag_servers(self, username: str, rag_servers: dict) -> List[str]:
+        """Get list of RAG servers the user is authorized to access.
+
+        This checks authorization directly against rag_mcp_config servers,
+        independent of mcp_manager.servers_config (which excludes RAG servers
+        to keep them separate from the tools panel).
+        """
+        authorized = []
+        for server_name, server_config in rag_servers.items():
+            if not server_config.enabled:
+                continue
+
+            required_groups = server_config.groups or []
+            if not required_groups:
+                # No group restriction - available to all
+                authorized.append(server_name)
+                continue
+
+            # Check if user is in any of the required groups
+            group_checks = [
+                await self.auth_check_func(username, group)
+                for group in required_groups
+            ]
+            if any(group_checks):
+                authorized.append(server_name)
+
+        return authorized
+
+    async def discover_data_sources(self, username: str, user_compliance_level: Optional[str] = None) -> List[str]:
+        """Discover data sources across authorized MCP RAG servers.
+
+        Phase 1 returns a flat list of strings for backward compatibility.
+        Uses server-qualified IDs: "{server}:{resource_id}" to avoid collisions.
+        """
+        # Ensure RAG servers are initialized from rag_mcp_config, without polluting tool inventory
+        try:
+            rag_servers = self.config_manager.rag_mcp_config.servers
+            # If these servers aren't in mcp_manager.clients, initialize just these
+            missing = [name for name in rag_servers.keys() if name not in getattr(self.mcp_manager, "clients", {})]
+            if missing:
+                # Temporarily extend servers_config with rag servers and initialize them
+                original = dict(getattr(self.mcp_manager, "servers_config", {}))
+                try:
+                    self.mcp_manager.servers_config.update({name: cfg.model_dump() for name, cfg in rag_servers.items()})
+                    await self.mcp_manager.initialize_clients()
+                    await self.mcp_manager.discover_tools()
+                finally:
+                    # Restore original list for general tools panel separation
+                    self.mcp_manager.servers_config = original
+        except Exception:
+            # If anything goes wrong, fallback silently to existing clients
+            pass
+        try:
+            # Determine RAG servers current user can see
+            # Use rag_mcp_config directly since servers_config was restored above
+            rag_servers = self.config_manager.rag_mcp_config.servers
+            authorized_servers: List[str] = await self._get_authorized_rag_servers(
+                username, rag_servers
+            )
+
+            if not authorized_servers:
+                logger.info("No authorized MCP servers for user %s", sanitize_for_logging(username))
+                return []
+
+            # --- Compliance Filtering (Step 2) ---
+            if user_compliance_level:
+                compliance_mgr = get_compliance_manager()
+                filtered_servers = []
+                for server in authorized_servers:
+                    cfg = (self.mcp_manager.available_tools.get(server) or {}).get("config", {})
+                    server_compliance_level = cfg.get("compliance_level")
+                    if compliance_mgr.is_accessible(
+                        user_level=user_compliance_level, resource_level=server_compliance_level
+                    ):
+                        filtered_servers.append(server)
+                    else:
+                        logger.info(
+                            "Skipping RAG server %s due to compliance level mismatch (user: %s, server: %s)",
+                            sanitize_for_logging(server),
+                            sanitize_for_logging(user_compliance_level),
+                            sanitize_for_logging(server_compliance_level),
+                        )
+                authorized_servers = filtered_servers
+                if not authorized_servers:
+                    logger.info("No authorized MCP servers remain after compliance filtering for user %s", sanitize_for_logging(username))
+                    return []
+            # -------------------------------------
+
+            # Filter to servers that advertise the discovery tool
+            servers_with_discovery: List[str] = []
+            for server in authorized_servers:
+                server_data = self.mcp_manager.available_tools.get(server)
+                tool_list = (server_data or {}).get("tools", [])
+                if any(getattr(t, "name", None) == "rag_discover_resources" for t in tool_list):
+                    servers_with_discovery.append(server)
+
+            if not servers_with_discovery:
+                logger.info("No servers implement rag_discover_resources for user %s", sanitize_for_logging(username))
+                return []
+
+            # Fan out discovery calls
+            sources: List[str] = []
+            for server in servers_with_discovery:
+                try:
+                    raw = await self.mcp_manager.call_tool(
+                        server_name=server,
+                        tool_name="rag_discover_resources",
+                        arguments={"username": username},
+                    )
+
+                    structured = self._extract_structured_result(raw)
+                    resources = self._extract_resources(structured)
+                    for r in resources:
+                        rid = r.get("id") or r.get("name")
+                        if not isinstance(rid, str):
+                            continue
+                        # Qualify with server to avoid collisions across providers
+                        sources.append(f"{server}:{rid}")
+                except Exception as e:
+                    logger.warning(
+                        "Discovery failed on server %s for user %s: %s",
+                        sanitize_for_logging(server),
+                        sanitize_for_logging(username),
+                        e,
+                    )
+
+            # De-dupe while preserving order
+            seen = set()
+            deduped = []
+            for s in sources:
+                if s not in seen:
+                    seen.add(s)
+                    deduped.append(s)
+            return deduped
+
+        except Exception as e:
+            logger.error("Error during RAG MCP discovery: %s", e, exc_info=True)
+            return []
+
+    async def discover_servers(self, username: str, user_compliance_level: Optional[str] = None) -> List[Dict[str, Any]]:
+        """Return richer per-server discovery structure for UI (rag_servers).
+
+        Shape:
+        [
+          {
+            "server": "docsRag",
+            "displayName": "docsRag",
+            "icon": <optional>,
+            "sources": [
+               {"id": "handbook", "name": "Employee Handbook", "authRequired": False, "selected": False}
+            ]
+          }
+        ]
+        """
+        # Ensure RAG servers are initialized from rag_mcp_config, without polluting tool inventory
+        try:
+            rag_cfg_servers = self.config_manager.rag_mcp_config.servers
+            missing = [name for name in rag_cfg_servers.keys() if name not in getattr(self.mcp_manager, "clients", {})]
+            if missing:
+                original = dict(getattr(self.mcp_manager, "servers_config", {}))
+                try:
+                    self.mcp_manager.servers_config.update({name: cfg.model_dump() for name, cfg in rag_cfg_servers.items()})
+                    await self.mcp_manager.initialize_clients()
+                    await self.mcp_manager.discover_tools()
+                finally:
+                    self.mcp_manager.servers_config = original
+        except Exception:
+            # Fallback silently if RAG config init fails; we'll just return empty set
+            pass
+
+        rag_servers: List[Dict[str, Any]] = []
+        try:
+            compliance_mgr = get_compliance_manager() if user_compliance_level else None
+
+            # Use rag_mcp_config directly since servers_config was restored above
+            rag_cfg_servers = self.config_manager.rag_mcp_config.servers
+            authorized_servers: List[str] = await self._get_authorized_rag_servers(
+                username, rag_cfg_servers
+            )
+
+            # --- Compliance Filtering (Step 2) ---
+            if compliance_mgr:
+                filtered_servers = []
+                for server in authorized_servers:
+                    cfg = (self.mcp_manager.available_tools.get(server) or {}).get("config", {})
+                    server_compliance_level = cfg.get("compliance_level")
+                    if compliance_mgr.is_accessible(
+                        user_level=user_compliance_level, resource_level=server_compliance_level
+                    ):
+                        filtered_servers.append(server)
+                    else:
+                        logger.info(
+                            "Skipping RAG server %s due to compliance level mismatch (user: %s, server: %s)",
+                            sanitize_for_logging(server),
+                            sanitize_for_logging(user_compliance_level),
+                            sanitize_for_logging(server_compliance_level),
+                        )
+                authorized_servers = filtered_servers
+            # -------------------------------------
+
+            for server in authorized_servers:
+                server_data = self.mcp_manager.available_tools.get(server)
+                tools = (server_data or {}).get("tools", [])
+                if not any(getattr(t, "name", None) == "rag_discover_resources" for t in tools):
+                    continue
+
+                # Call discovery
+                try:
+                    raw = await self.mcp_manager.call_tool(
+                        server_name=server,
+                        tool_name="rag_discover_resources",
+                        arguments={"username": username},
+                    )
+                    structured = self._extract_structured_result(raw)
+                    resources = self._extract_resources(structured)
+                except Exception as e:
+                    logger.warning("Discovery failed for server %s: %s", server, e)
+                    resources = []
+
+                # Build UI sources array
+                ui_sources: List[Dict[str, Any]] = []
+                for r in resources:
+                    rid = r.get("id") or r.get("name")
+                    if not isinstance(rid, str):
+                        continue
+
+                    # --- Compliance Filtering (Step 3) ---
+                    # Check for both camelCase (MCP standard) and snake_case (RAG mock standard)
+                    resource_compliance_level = r.get("complianceLevel") or r.get("compliance_level")
+                    if compliance_mgr and not compliance_mgr.is_accessible(
+                        user_level=user_compliance_level, resource_level=resource_compliance_level
+                    ):
+                        logger.info(
+                            "Skipping RAG resource %s:%s due to compliance level mismatch (user: %s, resource: %s)",
+                            sanitize_for_logging(server),
+                            sanitize_for_logging(rid),
+                            sanitize_for_logging(user_compliance_level),
+                            sanitize_for_logging(resource_compliance_level),
+                        )
+                        continue
+                    # -------------------------------------
+
+                    ui_sources.append({
+                        "id": rid,
+                        "name": r.get("name") or rid,
+                        # New contract: authRequired expected true; pass-through in case of legacy servers
+                        "authRequired": bool(r.get("authRequired", True)),
+                        # New: include per-resource groups when provided
+                        "groups": list(r.get("groups", [])) if isinstance(r.get("groups"), list) else None,
+                        "selected": bool(r.get("defaultSelected", False)),
+                        # Include compliance_level from resource or inherit from server
+                        "complianceLevel": resource_compliance_level if resource_compliance_level else None,
+                    })
+
+                # Optional config-driven icon/name and compliance level
+                cfg = (self.mcp_manager.available_tools.get(server) or {}).get("config", {})
+                display_name = cfg.get("displayName") or server
+                icon = (cfg.get("ui") or {}).get("icon") if isinstance(cfg.get("ui"), dict) else None
+                compliance_level = cfg.get("compliance_level")
+
+                rag_servers.append({
+                    "server": server,
+                    "displayName": display_name,
+                    "icon": icon,
+                    "complianceLevel": compliance_level,
+                    "sources": ui_sources,
+                })
+        except Exception as e:
+            logger.error("discover_servers error: %s", e, exc_info=True)
+
+        return rag_servers
+
+    async def search_raw(
+        self,
+        username: str,
+        query: str,
+        sources: List[str],
+        top_k: int = 8,
+        filters: Optional[Dict[str, Any]] = None,
+        ranking: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Call rag_get_raw_results across servers and merge results.
+
+        sources are server-qualified (server:id). We group by server.
+        """
+        logger.debug(
+            "[MCP-RAG] search_raw called: user=%s, query_preview=%s..., sources=%s, top_k=%d",
+            sanitize_for_logging(username),
+            sanitize_for_logging(query[:100]) if query else "(empty)",
+            sources,
+            top_k,
+        )
+
+        filters = filters or {}
+        ranking = ranking or {}
+        by_server: Dict[str, List[str]] = {}
+        for s in sources or []:
+            if isinstance(s, str) and ":" in s:
+                srv, rid = s.split(":", 1)
+                by_server.setdefault(srv, []).append(rid)
+
+        logger.debug("[MCP-RAG] search_raw sources grouped by server: %s", by_server)
+
+        all_hits: List[Dict[str, Any]] = []
+        meta: Dict[str, Any] = {"providers": {}, "top_k": top_k}
+
+        for server, rids in by_server.items():
+            logger.debug("[MCP-RAG] search_raw processing server=%s, resource_ids=%s", server, rids)
+            try:
+                # Check tool availability
+                server_data = self.mcp_manager.available_tools.get(server) or {}
+                tool_list = server_data.get("tools", [])
+                if not any(getattr(t, "name", None) == "rag_get_raw_results" for t in tool_list):
+                    logger.debug("[MCP-RAG] Server %s lacks rag_get_raw_results tool, skipping", server)
+                    continue
+
+                logger.debug("[MCP-RAG] Calling rag_get_raw_results on server %s", server)
+                raw = await self.mcp_manager.call_tool(
+                    server_name=server,
+                    tool_name="rag_get_raw_results",
+                    arguments={
+                        "username": username,
+                        "query": query,
+                        "sources": rids,
+                        "top_k": top_k,
+                        "filters": filters,
+                        "ranking": ranking,
+                    },
+                )
+                logger.debug("[MCP-RAG] Server %s raw response type: %s", server, type(raw).__name__)
+
+                payload = self._extract_structured_result(raw) or {}
+                results = payload.get("results") or {}
+                hits = results.get("hits") or []
+                logger.debug("[MCP-RAG] Server %s returned %d hits", server, len(hits))
+
+                # Annotate with server for provenance
+                for h in hits:
+                    if isinstance(h, dict):
+                        h.setdefault("server", server)
+                all_hits.extend([h for h in hits if isinstance(h, dict)])
+                meta["providers"][server] = {
+                    "returned": len(hits),
+                    "error": None,
+                }
+            except Exception as e:
+                logger.error("[MCP-RAG] Server %s search_raw error: %s", server, e, exc_info=True)
+                meta["providers"][server] = {"returned": 0, "error": str(e)}
+
+        # Merge + rerank (simple): sort by score desc if present
+        def score_of(h: Dict[str, Any]) -> float:
+            try:
+                return float(h.get("score", 0.0))
+            except Exception:
+                return 0.0
+
+        all_hits.sort(key=score_of, reverse=True)
+        merged = all_hits[: top_k or len(all_hits)]
+
+        logger.info(
+            "[MCP-RAG] search_raw complete: total_hits=%d, merged_count=%d, providers=%s",
+            len(all_hits),
+            len(merged),
+            list(meta["providers"].keys()),
+        )
+
+        # Prompt-injection risk check on retrieved snippets (observe + log)
+        try:
+            for h in merged:
+                if not isinstance(h, dict):
+                    continue
+                text = h.get("snippet") or h.get("chunk") or h.get("text") or ""
+                if not isinstance(text, str) or not text.strip():
+                    continue
+                pi = calculate_prompt_injection_risk(text, mode="general")
+                if pi.get("risk_level") in ("medium", "high"):
+                    log_high_risk_event(
+                        source="rag_chunk",
+                        user=username,
+                        content=text,
+                        score=int(pi.get("score", 0)),
+                        risk_level=str(pi.get("risk_level")),
+                        triggers=list(pi.get("triggers", [])),
+                        extra={
+                            "server": h.get("server"),
+                            "resourceId": h.get("resourceId"),
+                        },
+                    )
+        except Exception:
+            logger.debug("Prompt risk check failed (RAG results)", exc_info=True)
+
+        return {
+            "results": {
+                "hits": merged,
+                "stats": {
+                    "total_found": len(all_hits),
+                    "top_k": top_k,
+                },
+            },
+            "meta_data": meta,
+        }
+
+    async def synthesize(
+        self,
+        username: str,
+        query: str,
+        sources: List[str],
+        top_k: Optional[int] = None,
+        synthesis_params: Optional[Dict[str, Any]] = None,
+        provided_context: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Call rag_get_synthesized_results across servers when available.
+
+        If not available, fall back to raw search and concatenate snippets.
+        """
+        logger.debug(
+            "[MCP-RAG] synthesize called: user=%s, query_preview=%s..., sources=%s, top_k=%s",
+            sanitize_for_logging(username),
+            sanitize_for_logging(query[:100]) if query else "(empty)",
+            sources,
+            top_k,
+        )
+
+        synthesis_params = synthesis_params or {}
+        provided_context = provided_context or {}
+
+        by_server: Dict[str, List[str]] = {}
+        for s in sources or []:
+            if isinstance(s, str) and ":" in s:
+                srv, rid = s.split(":", 1)
+                by_server.setdefault(srv, []).append(rid)
+
+        logger.debug("[MCP-RAG] Sources grouped by server: %s", by_server)
+
+        answers: List[str] = []
+        citations: List[Dict[str, Any]] = []
+        meta: Dict[str, Any] = {"providers": {}}
+        used_fallback = False
+
+        for server, rids in by_server.items():
+            logger.debug(
+                "[MCP-RAG] Processing server=%s, resource_ids=%s",
+                server,
+                rids,
+            )
+            try:
+                server_data = self.mcp_manager.available_tools.get(server) or {}
+                tool_list = server_data.get("tools", [])
+                tool_names = [getattr(t, "name", None) for t in tool_list]
+                logger.debug("[MCP-RAG] Server %s available tools: %s", server, tool_names)
+
+                has_synth = any(getattr(t, "name", None) == "rag_get_synthesized_results" for t in tool_list)
+                if has_synth:
+                    logger.debug("[MCP-RAG] Server %s has rag_get_synthesized_results, calling...", server)
+                    raw = await self.mcp_manager.call_tool(
+                        server_name=server,
+                        tool_name="rag_get_synthesized_results",
+                        arguments={
+                            "username": username,
+                            "query": query,
+                            "sources": rids,
+                            **({"top_k": top_k} if top_k is not None else {}),
+                            "synthesis_params": synthesis_params,
+                            "provided_context": provided_context,
+                        },
+                    )
+                    logger.debug("[MCP-RAG] Server %s raw response type: %s", server, type(raw).__name__)
+
+                    payload = self._extract_structured_result(raw) or {}
+                    logger.debug("[MCP-RAG] Server %s extracted payload keys: %s", server, list(payload.keys()))
+                    logger.debug("[MCP-RAG] Server %s full payload: %s", server, sanitize_for_logging(str(payload)[:1000]))
+
+                    results = payload.get("results") or {}
+                    logger.debug("[MCP-RAG] Server %s results type: %s, results keys: %s", server, type(results).__name__, list(results.keys()) if isinstance(results, dict) else "N/A")
+                    logger.debug("[MCP-RAG] Server %s results content: %s", server, sanitize_for_logging(str(results)[:500]))
+
+                    ans = results.get("answer")
+                    if isinstance(ans, str) and ans:
+                        logger.debug(
+                            "[MCP-RAG] Server %s answer length=%d, preview=%s...",
+                            server,
+                            len(ans),
+                            sanitize_for_logging(ans[:200]),
+                        )
+                        answers.append(ans)
+                    cits = results.get("citations") or []
+                    if isinstance(cits, list):
+                        for c in cits:
+                            if isinstance(c, dict):
+                                c.setdefault("server", server)
+                        citations.extend([c for c in cits if isinstance(c, dict)])
+                    meta["providers"][server] = {"used_synth": True, "error": None}
+                    logger.info("[MCP-RAG] Server %s synthesis complete: answer_length=%d", server, len(ans) if ans else 0)
+                else:
+                    logger.debug("[MCP-RAG] Server %s lacks rag_get_synthesized_results, using fallback search_raw", server)
+                    used_fallback = True
+                    raw_payload = await self.search_raw(username, query, [f"{server}:{rid}" for rid in rids], top_k=top_k or 8)
+                    # Build a rudimentary answer from snippets
+                    hits = ((raw_payload.get("results") or {}).get("hits") or [])
+                    logger.debug("[MCP-RAG] Server %s fallback search returned %d hits", server, len(hits))
+                    snippet_texts = [h.get("snippet") or h.get("chunk") or "" for h in hits if isinstance(h, dict)]
+                    if snippet_texts:
+                        answers.append("\n\n".join(snippet_texts[:3]))
+                    meta["providers"][server] = {"used_synth": False, "error": None}
+            except Exception as e:
+                logger.error("[MCP-RAG] Server %s synthesis error: %s", server, e, exc_info=True)
+                meta["providers"][server] = {"used_synth": False, "error": str(e)}
+
+        final_answer = "\n\n---\n\n".join([a for a in answers if a]) if answers else ""
+        logger.info(
+            "[MCP-RAG] synthesize complete: total_answers=%d, final_answer_length=%d, used_fallback=%s",
+            len(answers),
+            len(final_answer),
+            used_fallback,
+        )
+
+        return {
+            "results": {
+                "answer": final_answer,
+                "citations": citations or None,
+                "limits": {"truncated": False} if final_answer else None,
+            },
+            "meta_data": {**meta, "fallback_used": used_fallback},
+        }
+
+    # --- helpers ---------------------------------------------------------
+    def _extract_structured_result(self, raw: Any) -> Dict[str, Any]:
+        """Best-effort extraction of a structured payload from FastMCP result."""
+        import json
+
+        logger.debug("[MCP-RAG] _extract_structured_result: raw type=%s", type(raw).__name__)
+
+        try:
+            # Log available attributes for debugging
+            if hasattr(raw, "__dict__"):
+                logger.debug("[MCP-RAG] _extract_structured_result: raw attributes=%s", list(raw.__dict__.keys()) if hasattr(raw, "__dict__") else "N/A")
+
+            # If raw is already a dict, return it directly
+            if isinstance(raw, dict):
+                logger.debug("[MCP-RAG] _extract_structured_result: raw is already a dict with keys=%s", list(raw.keys()))
+                return raw
+
+            # Preferred attributes from fastmcp
+            if hasattr(raw, "structured_content") and raw.structured_content:
+                logger.debug("[MCP-RAG] _extract_structured_result: found structured_content")
+                if isinstance(raw.structured_content, dict):
+                    return raw.structured_content
+            if hasattr(raw, "data") and raw.data:
+                logger.debug("[MCP-RAG] _extract_structured_result: found data")
+                if isinstance(raw.data, dict):
+                    return raw.data
+
+            if hasattr(raw, "content") and raw.content:
+                contents = getattr(raw, "content")
+                logger.debug("[MCP-RAG] _extract_structured_result: found content, type=%s, len=%s", type(contents).__name__, len(contents) if hasattr(contents, '__len__') else "N/A")
+
+                # content is typically a list of segments with .text
+                if isinstance(contents, list) and contents:
+                    # Try all content items, not just the first
+                    for idx, item in enumerate(contents):
+                        logger.debug("[MCP-RAG] _extract_structured_result: content[%d] type=%s", idx, type(item).__name__)
+
+                        # Try .text attribute
+                        text = getattr(item, "text", None)
+                        if text is None and isinstance(item, dict):
+                            text = item.get("text")
+
+                        if text:
+                            logger.debug("[MCP-RAG] _extract_structured_result: text type=%s, preview=%s", type(text).__name__, sanitize_for_logging(str(text)[:300]))
+                            if isinstance(text, str) and text.strip():
+                                try:
+                                    obj = json.loads(text)
+                                    if isinstance(obj, dict):
+                                        logger.debug("[MCP-RAG] _extract_structured_result: parsed JSON with keys=%s", list(obj.keys()))
+                                        return obj
+                                except Exception as json_err:
+                                    logger.debug("[MCP-RAG] _extract_structured_result: JSON parse failed for content[%d]: %s", idx, json_err)
+
+                        # Try if item is itself a dict with results/meta_data
+                        if isinstance(item, dict):
+                            if "results" in item or "meta_data" in item:
+                                logger.debug("[MCP-RAG] _extract_structured_result: content[%d] is a dict with results/meta_data", idx)
+                                return item
+
+                # If content is a single string (not list), try to parse as JSON
+                elif isinstance(contents, str) and contents.strip():
+                    logger.debug("[MCP-RAG] _extract_structured_result: content is a string, trying JSON parse")
+                    try:
+                        obj = json.loads(contents)
+                        if isinstance(obj, dict):
+                            logger.debug("[MCP-RAG] _extract_structured_result: parsed content string as JSON with keys=%s", list(obj.keys()))
+                            return obj
+                    except Exception as json_err:
+                        logger.debug("[MCP-RAG] _extract_structured_result: JSON parse of content string failed: %s", json_err)
+
+            # Try to convert raw to string and parse as JSON (last resort)
+            if hasattr(raw, "__str__"):
+                raw_str = str(raw)
+                # Only try if it looks like JSON
+                if raw_str.strip().startswith("{"):
+                    logger.debug("[MCP-RAG] _extract_structured_result: trying __str__ as JSON: %s...", sanitize_for_logging(raw_str[:200]))
+                    try:
+                        obj = json.loads(raw_str)
+                        if isinstance(obj, dict):
+                            logger.debug("[MCP-RAG] _extract_structured_result: parsed __str__ as JSON with keys=%s", list(obj.keys()))
+                            return obj
+                    except Exception:
+                        pass
+
+        except Exception as parse_err:  # pragma: no cover - defensive
+            logger.debug("Non-fatal: failed to parse structured result: %s", parse_err)
+
+        logger.debug("[MCP-RAG] _extract_structured_result: returning empty dict")
+        return {}
+
+    def _extract_resources(self, payload: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """Extract list of resource dicts from a normalized tool result."""
+        if not isinstance(payload, dict):
+            return []
+        results = payload.get("results") if isinstance(payload.get("results"), dict) else payload
+        # Support both {results: {resources: [...]}} and {results: [...]}
+        # Also support the RAG mock format: {accessible_data_sources: [...]}
+        resources = (
+            (results.get("resources") if isinstance(results, dict) else None)
+            or payload.get("resources")
+            or payload.get("accessible_data_sources") # Added support for RAG mock format
+            or []
+        )
+        if isinstance(resources, list):
+            # ensure each entry is a dict
+            return [r for r in resources if isinstance(r, dict)]
+        return []
+
+
+__all__ = ["RAGMCPService"]

atlas/domain/sessions/__init__.py

@@ -0,0 +1,7 @@
+"""Domain models for sessions."""
+
+from .models import Session
+
+__all__ = [
+    "Session",
+]