Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 131 additions & 0 deletions airbyte/mcp/kapa.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,131 @@
# Copyright (c) 2026 Airbyte, Inc., all rights reserved.
"""Kapa knowledge source MCP operations."""

from __future__ import annotations

from typing import TYPE_CHECKING, Annotated

import requests
from fastmcp_extensions import mcp_tool, register_mcp_tools
from pydantic import Field

from airbyte.secrets.util import try_get_secret


if TYPE_CHECKING:
from fastmcp import FastMCP


_KAPA_RETRIEVAL_URL_TEMPLATE = "https://api.kapa.ai/query/v1/projects/{project_id}/retrieval/"
_KAPA_TIMEOUT_SECONDS = 30.0
_KAPA_API_KEY_ENV_VAR = "KAPA_API_KEY"
_KAPA_DOCS_MCP_BEARER_TOKEN_ENV_VAR = "KAPA_DOCS_MCP_BEARER_TOKEN"
_KAPA_BEARER_TOKEN_ENV_VAR = "KAPA_BEARER_TOKEN"
_KAPA_CREDENTIAL_ENV_VARS = (
_KAPA_API_KEY_ENV_VAR,
_KAPA_DOCS_MCP_BEARER_TOKEN_ENV_VAR,
_KAPA_BEARER_TOKEN_ENV_VAR,
)
_KAPA_PROJECT_ID_ENV_VAR = "KAPA_PROJECT_ID"
_KAPA_INTEGRATION_ID_ENV_VAR = "KAPA_INTEGRATION_ID"


def _kapa_config_value(name: str) -> str:
value = try_get_secret(name)
if value is None:
return ""

return str(value).strip()


def _kapa_credentials_configured() -> bool:
return any(_kapa_config_value(name) for name in _KAPA_CREDENTIAL_ENV_VARS)


def _kapa_project_configured() -> bool:
return bool(_kapa_config_value(_KAPA_PROJECT_ID_ENV_VAR))


def _kapa_auth_headers() -> dict[str, str]:
api_key = _kapa_config_value(_KAPA_API_KEY_ENV_VAR)
if api_key:
return {"X-API-KEY": api_key}

bearer_token = _kapa_config_value(_KAPA_DOCS_MCP_BEARER_TOKEN_ENV_VAR) or _kapa_config_value(
_KAPA_BEARER_TOKEN_ENV_VAR
)
if bearer_token:
return {"Authorization": f"Bearer {bearer_token}"}

raise ValueError(
"Kapa docs search is not configured. Set KAPA_API_KEY, "
"KAPA_DOCS_MCP_BEARER_TOKEN, or KAPA_BEARER_TOKEN."
)
Comment on lines +49 to +63

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Keep this DRY and call the other helper, or take its output.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated in commit 7ed7611. _kapa_auth_headers() now uses the same _kapa_config_value() helper as registration and payload construction, so the credential lookup and fallback order are centralized.


Devin session



def _kapa_retrieval_url() -> str:
project_id = _kapa_config_value(_KAPA_PROJECT_ID_ENV_VAR)
if not project_id:
raise ValueError("KAPA_PROJECT_ID must be set to use Kapa docs search.")

return _KAPA_RETRIEVAL_URL_TEMPLATE.format(project_id=project_id)


def _kapa_request_payload(query: str) -> dict[str, str]:
payload = {"query": query}
integration_id = _kapa_config_value(_KAPA_INTEGRATION_ID_ENV_VAR)
if integration_id:
payload["integration_id"] = integration_id
return payload


def _normalize_kapa_results(data: object) -> list[dict[str, str]]:
if not isinstance(data, list):
raise TypeError("Kapa retrieval API returned an unexpected response shape.")

results: list[dict[str, str]] = []
for item in data:
if not isinstance(item, dict):
raise TypeError("Kapa retrieval API returned an unexpected result item.")

source_url = item.get("source_url")
content = item.get("content")
if not isinstance(source_url, str) or not isinstance(content, str):
raise TypeError(
"Kapa retrieval API returned a result without source_url and content strings."
)

results.append({"source_url": source_url, "content": content})

return results


@mcp_tool(
read_only=True,
idempotent=True,
)
def search_airbyte_knowledge_sources(
query: Annotated[
str,
Field(
description=(
"A single, well-formed natural-language query. Must be a complete sentence."
),
),
],
) -> list[dict[str, str]]:
"""Search Airbyte knowledge sources."""
response = requests.post(
_kapa_retrieval_url(),
headers=_kapa_auth_headers(),
json=_kapa_request_payload(query),
timeout=_KAPA_TIMEOUT_SECONDS,
)
response.raise_for_status()
return _normalize_kapa_results(response.json())


def register_kapa_tools(app: FastMCP) -> None:
"""Register Kapa tools with the FastMCP app."""
if _kapa_credentials_configured() and _kapa_project_configured():
register_mcp_tools(app, mcp_module=__name__)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
2 changes: 2 additions & 0 deletions airbyte/mcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
airbyte_readonly_mode_filter,
)
from airbyte.mcp.cloud import register_cloud_tools
from airbyte.mcp.kapa import register_kapa_tools
from airbyte.mcp.local import register_local_tools
from airbyte.mcp.prompts import register_prompts
from airbyte.mcp.registry import register_registry_tools
Expand Down Expand Up @@ -93,6 +94,7 @@
register_cloud_tools(app)
register_local_tools(app)
register_registry_tools(app)
register_kapa_tools(app)
register_prompts(app)


Expand Down
107 changes: 107 additions & 0 deletions tests/unit_tests/test_mcp_kapa.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
# Copyright (c) 2026 Airbyte, Inc., all rights reserved.
"""Unit tests for Kapa MCP tools."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
import responses

from airbyte.mcp import kapa


@pytest.fixture(autouse=True)
def clear_kapa_env(monkeypatch: pytest.MonkeyPatch) -> None:
"""Clear Kapa environment variables before each test."""
monkeypatch.delenv("KAPA_API_KEY", raising=False)
monkeypatch.delenv("KAPA_DOCS_MCP_BEARER_TOKEN", raising=False)
monkeypatch.delenv("KAPA_BEARER_TOKEN", raising=False)
monkeypatch.delenv("KAPA_PROJECT_ID", raising=False)
monkeypatch.delenv("KAPA_INTEGRATION_ID", raising=False)


@pytest.mark.parametrize(
("env_name", "expected_headers"),
[
pytest.param("KAPA_API_KEY", {"X-API-KEY": "secret"}, id="api-key"),
pytest.param(
"KAPA_DOCS_MCP_BEARER_TOKEN",
{"Authorization": "Bearer secret"},
id="docs-mcp-bearer-token",
),
pytest.param(
"KAPA_BEARER_TOKEN", {"Authorization": "Bearer secret"}, id="bearer-token"
),
],
)
def test_kapa_auth_headers(
env_name: str,
expected_headers: dict[str, str],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Test Kapa auth header selection from supported env vars."""
monkeypatch.setenv(env_name, "secret")

assert kapa._kapa_auth_headers() == expected_headers # noqa: SLF001


@responses.activate
def test_search_airbyte_knowledge_sources_calls_kapa_retrieval_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Test Kapa search request construction and response normalization."""
monkeypatch.setenv("KAPA_API_KEY", "secret")
monkeypatch.setenv("KAPA_PROJECT_ID", "project-id")
monkeypatch.setenv("KAPA_INTEGRATION_ID", "integration-id")
responses.add(
responses.POST,
"https://api.kapa.ai/query/v1/projects/project-id/retrieval/",
json=[{"source_url": "https://docs.airbyte.com", "content": "Airbyte docs"}],
status=200,
)

result = kapa.search_airbyte_knowledge_sources("How do I set up a source?")

assert result == [
{"source_url": "https://docs.airbyte.com", "content": "Airbyte docs"}
]
assert len(responses.calls) == 1
request = responses.calls[0].request
assert request.headers["X-API-KEY"] == "secret"
assert (
request.body
== b'{"query": "How do I set up a source?", "integration_id": "integration-id"}'
)


def test_register_kapa_tools_skips_registration_without_credentials() -> None:
"""Test that Kapa tools are hidden when credentials are absent."""
app = MagicMock()

with patch("airbyte.mcp.kapa.register_mcp_tools") as register:
kapa.register_kapa_tools(app)

register.assert_not_called()


@pytest.mark.parametrize(
"env_name",
[
pytest.param(name, id=name.lower())
for name in kapa._KAPA_CREDENTIAL_ENV_VARS # noqa: SLF001
],
)
def test_register_kapa_tools_registers_when_credentials_are_configured(
monkeypatch: pytest.MonkeyPatch,
env_name: str,
) -> None:
"""Test that Kapa tools are visible when credentials are configured."""
app = MagicMock()
monkeypatch.setenv(env_name, "secret")
monkeypatch.setenv("KAPA_PROJECT_ID", "project-id")

with patch("airbyte.mcp.kapa.register_mcp_tools") as register:
kapa.register_kapa_tools(app)

register.assert_called_once_with(app, mcp_module="airbyte.mcp.kapa")
Loading