WIP: BYOK via MCP — knowledge-type MCP servers

ols/app/models/config.py

@@ -631,6 +631,10 @@ class MCPServerConfig(BaseModel):
     MCP (Model Context Protocol) servers provide tools and capabilities to the
     AI agents. These are configured by this structure. Only MCP servers
     defined in the olsconfig.yaml configuration are available to the agents.
+
+    Servers with ``type: knowledge`` are queried on every request (forced
+    retrieval) and their results are injected as additional RAG context.
+    Servers with ``type: tool`` (the default) expose tools the LLM may call.
     """
 
     name: str = Field(
@@ -643,6 +647,68 @@ class MCPServerConfig(BaseModel):
         description="URL of the MCP server",
     )
 
+    type: str = Field(
+        default="tool",
+        title="MCP server type",
+        description=(
+            "Type of MCP server. "
+            "'tool' (default) — tools are passed to the LLM which decides when to call them. "
+            "'knowledge' — a knowledge source queried according to retrieval_mode."
+        ),
+    )
+
+    retrieval_mode: str = Field(
+        default="always",
+        title="Retrieval mode",
+        description=(
+            "How the knowledge source is queried. Only used when type is 'knowledge'. "
+            "'always' — search_tool is called on every query and results are injected "
+            "as RAG context (forced retrieval). "
+            "'on_demand' — the search_tool is exposed as an LLM tool, the model decides "
+            "when to call it. Requires 'description' to tell the LLM what the source contains."
+        ),
+    )
+
+    knowledge_description: Optional[str] = Field(
+        default=None,
+        title="Knowledge source description",
+        description=(
+            "Human-readable description of what this knowledge source contains. "
+            "Used when retrieval_mode is 'on_demand' to help the LLM decide when to "
+            "query this source. Example: 'Internal incident response runbooks and "
+            "security playbooks. Search when the user asks about incident response, "
+            "ransomware, phishing, or security procedures.'"
+        ),
+    )
+
+    search_tool: Optional[str] = Field(
+        default=None,
+        title="Search tool name",
+        description=(
+            "Name of the MCP tool to call for knowledge retrieval. "
+            "Required when type is 'knowledge'. "
+            "The tool must accept a 'query' parameter."
+        ),
+    )
+
+    max_chunks: int = Field(
+        default=10,
+        title="Max chunks",
+        description=(
+            "Maximum number of chunks to retrieve from this knowledge source. "
+            "Only used when type is 'knowledge'."
+        ),
+    )
+
+    priority: int = Field(
+        default=1,
+        title="Priority",
+        description=(
+            "Priority for ordering results when multiple knowledge sources return results. "
+            "Lower numbers = higher priority. Only used when type is 'knowledge'."
+        ),
+    )
+
     timeout: Optional[int] = Field(
         default=None,
         title="Request timeout",
@@ -676,6 +742,36 @@ def resolved_headers(self) -> dict[str, str]:
         """Resolved headers (computed from headers)."""
         return self._resolved_headers
 
+    @model_validator(mode="after")
+    def validate_knowledge_config(self) -> "MCPServerConfig":
+        """Validate that knowledge-type servers have search_tool configured."""
+        if self.type not in ("tool", "knowledge"):
+            raise ValueError(
+                f"Invalid MCP server type '{self.type}' for server '{self.name}'. "
+                "Must be 'tool' or 'knowledge'."
+            )
+        if self.type == "knowledge" and not self.search_tool:
+            raise ValueError(
+                f"MCP server '{self.name}' has type 'knowledge' but no search_tool configured. "
+                "Set search_tool to the name of the MCP tool to call for retrieval."
+            )
+        if self.retrieval_mode not in ("always", "on_demand"):
+            raise ValueError(
+                f"Invalid retrieval_mode '{self.retrieval_mode}' for server '{self.name}'. "
+                "Must be 'always' or 'on_demand'."
+            )
+        if (
+            self.type == "knowledge"
+            and self.retrieval_mode == "on_demand"
+            and not self.knowledge_description
+        ):
+            raise ValueError(
+                f"MCP server '{self.name}' has retrieval_mode 'on_demand' but no "
+                "knowledge_description. The LLM needs a description to decide when "
+                "to query this source."
+            )
+        return self
+
 
 class ToolFilteringConfig(BaseModel):
     """Configuration for tool filtering using hybrid RAG retrieval.

ols/src/query_helpers/docs_summarizer.py

@@ -29,7 +29,12 @@
 from ols.src.query_helpers.query_helper import QueryHelper
 from ols.src.skills.skills_rag import create_skill_support_tool
 from ols.src.tools.offloaded_content import OffloadManager
-from ols.utils.mcp_utils import ClientHeaders, build_mcp_config, get_mcp_tools
+from ols.utils.mcp_utils import (
+    ClientHeaders,
+    build_mcp_config,
+    get_mcp_tools,
+    retrieve_from_knowledge_mcps,
+)
 from ols.utils.token_handler import (
     PromptTooLongError,
     TokenBudgetTracker,
@@ -288,6 +293,28 @@ async def generate_response(  # noqa: C901  # pylint: disable=too-many-branches,
         """
         rag_chunks = await self._prepare_prompt_context(query, rag_retriever)
 
+        knowledge_results = await retrieve_from_knowledge_mcps(
+            query, self.user_token, self.client_headers
+        )
+        if knowledge_results:
+            for kr in knowledge_results:
+                byok_chunk = RagChunk(
+                    text=kr["text"],
+                    doc_url=kr.get("source", ""),
+                    doc_title=kr.get("title", "BYOK"),
+                )
+                rag_chunks.insert(0, byok_chunk)
+            byok_tokens = sum(
+                self._tracker.count_tokens(kr["text"]) for kr in knowledge_results
+            )
+            self._tracker.charge(TokenCategory.RAG, byok_tokens)
+            logger.info(
+                "Injected %d BYOK knowledge chunks (%d tokens) from %d MCP sources",
+                len(knowledge_results),
+                byok_tokens,
+                len(knowledge_results),
+            )
+
         skill_content: Optional[str] = None
         has_support_files = False
         skill = None

ols/utils/mcp_utils.py

@@ -206,13 +206,31 @@ async def gather_mcp_tools(
                     tool for tool in server_tools if tool.name in allowed_tool_names
                 ]
 
-            # Add MCP server name to each tool's metadata
+            server_cfg = config.mcp_servers_dict.get(server_name)
+            knowledge_desc = (
+                server_cfg.knowledge_description
+                if server_cfg and server_cfg.type == "knowledge"
+                and server_cfg.knowledge_description
+                else None
+            )
+
             for tool in server_tools:
                 _normalize_tool_schema(tool)
                 if not hasattr(tool, "metadata") or tool.metadata is None:
                     tool.metadata = {}
                 tool.metadata["mcp_server"] = server_name
 
+                if knowledge_desc and server_cfg and tool.name == server_cfg.search_tool:
+                    tool.description = (
+                        f"{knowledge_desc}\n\n"
+                        f"Original tool description: {tool.description}"
+                    )
+                    logger.info(
+                        "Augmented tool '%s' with knowledge_description from '%s'",
+                        tool.name,
+                        server_name,
+                    )
+
             all_tools.extend(server_tools)
             logger.info(
                 "Loaded %d tools from MCP server '%s'",
@@ -347,10 +365,15 @@ async def get_mcp_tools(
     Returns:
         List of all tools from MCP servers (filtered if tools_rag configured).
     """
+    tool_servers = [
+        s for s in config.mcp_servers.servers
+        if not (s.type == "knowledge" and s.retrieval_mode == "always")
+    ]
+
     # If tools_rag is not configured, return all tools
     if not config.tools_rag:
         mcp_servers_config, all_tools = await _gather_and_populate_tools(
-            config.mcp_servers.servers, user_token, client_headers, deduplicate=True
+            tool_servers, user_token, client_headers, deduplicate=True
         )
 
         if not mcp_servers_config:
@@ -432,6 +455,94 @@ async def get_mcp_tools(
     return []
 
 
+async def retrieve_from_knowledge_mcps(
+    query: str,
+    user_token: Optional[str] = None,
+    client_headers: ClientHeaders | None = None,
+) -> list[dict[str, str]]:
+    """Retrieve context from all knowledge-type MCP servers (forced retrieval).
+
+    Connects to each MCP server with type='knowledge', calls its configured
+    search_tool with the query, and returns the results as a list of dicts
+    with 'text', 'title', and 'source' keys.
+
+    Args:
+        query: The user query to search for.
+        user_token: Optional user authentication token.
+        client_headers: Optional client-provided MCP headers.
+
+    Returns:
+        List of dicts with 'text', 'title', and 'source' keys,
+        sorted by server priority (lower = higher priority).
+    """
+    knowledge_servers = [
+        s for s in config.mcp_servers.servers
+        if s.type == "knowledge" and s.retrieval_mode == "always"
+    ]
+    if not knowledge_servers:
+        return []
+
+    servers_config = build_mcp_config(knowledge_servers, user_token, client_headers)
+    if not servers_config:
+        return []
+
+    all_results: list[tuple[int, dict[str, str]]] = []
+
+    mcp_client = MultiServerMCPClient(servers_config)
+
+    for server_cfg in knowledge_servers:
+        if server_cfg.name not in servers_config:
+            continue
+        try:
+            tools = await mcp_client.get_tools(server_name=server_cfg.name)
+            search_tool = next(
+                (t for t in tools if t.name == server_cfg.search_tool), None
+            )
+            if search_tool is None:
+                logger.error(
+                    "Knowledge MCP '%s': search_tool '%s' not found. Available: %s",
+                    server_cfg.name,
+                    server_cfg.search_tool,
+                    [t.name for t in tools],
+                )
+                continue
+
+            result = await search_tool.ainvoke({"query": query})
+            result_text = str(result) if result else ""
+
+            if not result_text.strip():
+                logger.info(
+                    "Knowledge MCP '%s': no results for query", server_cfg.name
+                )
+                continue
+
+            all_results.append((
+                server_cfg.priority,
+                {
+                    "text": result_text[: 5000],
+                    "title": f"BYOK: {server_cfg.name}",
+                    "source": server_cfg.name,
+                },
+            ))
+
+            logger.info(
+                "Knowledge MCP '%s': retrieved %d chars via '%s'",
+                server_cfg.name,
+                len(result_text),
+                server_cfg.search_tool,
+            )
+        except Exception as e:
+            logger.error(
+                "Failed to retrieve from knowledge MCP '%s': %s: %s",
+                server_cfg.name,
+                type(e).__name__,
+                e,
+            )
+
+    all_results.sort(key=lambda x: x[0])
+    return [r[1] for r in all_results]
+
+
 def build_mcp_config(
     servers_list: list[MCPServerConfig],
     user_token: Optional[str],