Ai 1343 Add Clients for DataScience, Encryption Services And Refactor

integtests/test_client.py → integtests/clients/test_client.py

@@ -3,7 +3,8 @@
 import pytest
 
 from integtests.conftest import ProjectDef, TableDef
-from keboola_mcp_server.client import AsyncStorageClient, GlobalSearchResponse, KeboolaClient
+from keboola_mcp_server.clients.client import KeboolaClient
+from keboola_mcp_server.clients.storage import AsyncStorageClient, GlobalSearchResponse
 
 LOG = logging.getLogger(__name__)
 
@@ -25,6 +26,7 @@ async def test_global_search(self, storage_client: AsyncStorageClient):
         assert ret.by_project == {}
 
     @pytest.mark.asyncio
+    @pytest.mark.skip(reason='Writing to the global search index takes too long, making the test unstable.')
     async def test_global_search_with_results(self, storage_client: AsyncStorageClient, tables: list[TableDef]):
         search_for_name = 'test'
         is_global_search_enabled = await storage_client.is_enabled('global-search')

integtests/clients/test_data_science.py

@@ -0,0 +1,105 @@
+import logging
+
+import pytest
+
+from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient
+from keboola_mcp_server.clients.data_science import DataAppResponse, DataScienceClient
+
+LOG = logging.getLogger(__name__)
+
+
+def _minimal_parameters(slug: str) -> dict[str, object]:
+    """Build minimal valid parameters for a code-based Streamlit data app."""
+    return {
+        'size': 'tiny',
+        'autoSuspendAfterSeconds': 600,
+        'dataApp': {
+            'slug': slug,
+            'streamlit': {
+                'config.toml': '[theme]\nbase = "light"',
+            },
+        },
+        'script': [
+            'import streamlit as st',
+            "st.write('Hello from integration test')",
+        ],
+    }
+
+
+def _public_access_authorization() -> dict[str, object]:
+    """Allow public access to all paths; no providers required."""
+    return {
+        'app_proxy': {
+            'auth_providers': [],
+            'auth_rules': [
+                {'type': 'pathPrefix', 'value': '/', 'auth_required': False},
+            ],
+        }
+    }
+
+
+@pytest.fixture
+def ds_client(keboola_client: KeboolaClient) -> DataScienceClient:
+    return keboola_client.data_science_client
+
+
+@pytest.mark.asyncio
+async def test_create_and_fetch_data_app(
+    ds_client: DataScienceClient, unique_id: str, keboola_client: KeboolaClient
+) -> None:
+    """Test creating a data app and fetching it from detail and list endpoints"""
+    slug = f'test-app-{unique_id}'
+    created: DataAppResponse = await ds_client.create_data_app(
+        name=f'IntegTest {slug}',
+        description='Created by integration tests',
+        parameters=_minimal_parameters(slug),
+        authorization=_public_access_authorization(),
+    )
+
+    try:
+        # Check if the created data app is valid
+        assert isinstance(created, DataAppResponse)
+        assert created.id
+        assert created.type == 'streamlit'
+        assert created.component_id == DATA_APP_COMPONENT_ID
+
+        # Deploy the data app
+        response = await ds_client.deploy_data_app(created.id, created.config_version)
+        assert response.id == created.id
+
+        # Fetch the data app from data science
+        fethced_ds = await ds_client.get_data_app(created.id)
+        assert fethced_ds.id == created.id
+        assert fethced_ds.type == created.type
+        assert fethced_ds.component_id == created.component_id
+        assert fethced_ds.project_id == created.project_id
+        assert fethced_ds.config_id == created.config_id
+        assert fethced_ds.config_version == created.config_version
+
+        # Fetch the data app config from storage
+        fetched_s = await keboola_client.storage_client.configuration_detail(
+            component_id=DATA_APP_COMPONENT_ID,
+            configuration_id=created.config_id,
+        )
+
+        # check if the data app ids are the same (data app from data science and config from storage)
+        assert 'configuration' in fetched_s
+        assert isinstance(fetched_s['configuration'], dict)
+        assert 'parameters' in fetched_s['configuration']
+        assert isinstance(fetched_s['configuration']['parameters'], dict)
+        assert 'id' in fetched_s['configuration']['parameters']
+        assert fethced_ds.id == fetched_s['configuration']['parameters']['id']
+
+        # Fetch the all data apps and check if the created data app is in the list
+        data_apps = await ds_client.list_data_apps()
+        assert isinstance(data_apps, list)
+        assert len(data_apps) > 0
+        assert any(app.id == created.id for app in data_apps)
+
+    finally:
+        for _ in range(2):  # Delete configuration 2 times (from storage and then from temporal bin)
+            try:
+                await ds_client.delete_data_app(created.id)
+            except Exception as e:
+                LOG.exception(f'Failed to delete data app: {e}')
+                pass

integtests/clients/test_encryption.py

@@ -0,0 +1,44 @@
+from typing import Any
+
+import pytest
+
+from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient
+
+
+@pytest.mark.asyncio
+async def test_encrypt_without_project_id(keboola_client: KeboolaClient) -> None:
+    """Check that the encryption client does not send any authorization headers."""
+    assert 'Authorization' not in keboola_client.encryption_client.raw_client.headers
+    assert 'X-StorageAPI-Token' not in keboola_client.encryption_client.raw_client.headers
+
+
+@pytest.mark.asyncio
+async def test_encrypt_string_not_equal(keboola_client: KeboolaClient) -> None:
+    project_id = await keboola_client.storage_client.project_id()
+    plaintext = 'my-plain-text'
+    encrypted = await keboola_client.encryption_client.encrypt(
+        value=plaintext,
+        project_id=str(project_id),
+        component_id=DATA_APP_COMPONENT_ID,
+    )
+    assert isinstance(encrypted, str)
+    assert encrypted != plaintext
+
+
+@pytest.mark.asyncio
+async def test_encrypt_dict_hash_keys_only(keboola_client: KeboolaClient) -> None:
+    project_id = await keboola_client.storage_client.project_id()
+    payload: dict[str, Any] = {
+        '#secret': 'sensitive-value',
+        'public': 'visible-value',
+    }
+    result = await keboola_client.encryption_client.encrypt(
+        value=payload,
+        project_id=str(project_id),
+        component_id=DATA_APP_COMPONENT_ID,
+    )
+    assert isinstance(result, dict)
+    # Values under keys beginning with '#' should be encrypted (changed)
+    assert result['#secret'] != payload['#secret']
+    # Non-secret values should remain the same
+    assert result['public'] == payload['public']

integtests/conftest.py

@@ -20,7 +20,7 @@
 from mcp.server.session import ServerSession
 from mcp.shared.context import RequestContext
 
-from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.clients.client import KeboolaClient
 from keboola_mcp_server.config import Config
 from keboola_mcp_server.mcp import ServerState
 from keboola_mcp_server.workspace import WorkspaceManager

integtests/test_errors.py

@@ -11,7 +11,7 @@
 from fastmcp import Context
 from mcp.types import ClientCapabilities, Implementation, InitializeRequestParams
 
-from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.clients.client import KeboolaClient
 from keboola_mcp_server.errors import tool_errors
 from keboola_mcp_server.tools.doc import docs_query
 from keboola_mcp_server.tools.jobs import get_job

integtests/test_mcp_server.py

@@ -6,7 +6,7 @@
 from mcp.types import TextContent
 
 from integtests.conftest import AsyncContextClientRunner, AsyncContextServerRemoteRunner, ConfigDef
-from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.clients.client import KeboolaClient
 from keboola_mcp_server.config import Config
 from keboola_mcp_server.server import create_server
 from keboola_mcp_server.tools.components.model import Configuration

integtests/test_validate.py

@@ -16,7 +16,9 @@
 import jsonschema
 import pytest
 
-from keboola_mcp_server.client import ComponentAPIResponse, JsonDict, KeboolaClient
+from keboola_mcp_server.clients.base import JsonDict
+from keboola_mcp_server.clients.client import KeboolaClient
+from keboola_mcp_server.clients.storage import ComponentAPIResponse
 from keboola_mcp_server.tools.components.model import Component
 from keboola_mcp_server.tools.validation import KeboolaParametersValidator
 

integtests/test_workspace.py

@@ -5,7 +5,7 @@
 import requests
 from kbcstorage.client import Client as SyncStorageClient
 
-from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.clients.client import KeboolaClient
 from keboola_mcp_server.workspace import WorkspaceManager
 
 LOG = logging.getLogger(__name__)

integtests/tools/components/test_tools.py

@@ -7,7 +7,7 @@
 from mcp.server.fastmcp import Context
 
 from integtests.conftest import ConfigDef, ProjectDef
-from keboola_mcp_server.client import KeboolaClient, get_metadata_property
+from keboola_mcp_server.clients.client import KeboolaClient, get_metadata_property
 from keboola_mcp_server.config import Config, MetadataField
 from keboola_mcp_server.links import Link
 from keboola_mcp_server.server import create_server

integtests/tools/flow/test_tools.py

@@ -8,7 +8,7 @@
 from pydantic import ValidationError
 
 from integtests.conftest import ConfigDef, ProjectDef
-from keboola_mcp_server.client import (
+from keboola_mcp_server.clients.client import (
     CONDITIONAL_FLOW_COMPONENT_ID,
     ORCHESTRATOR_COMPONENT_ID,
     FlowType,

integtests/tools/test_jobs.py

@@ -5,7 +5,7 @@
 from mcp.server.fastmcp import Context
 
 from integtests.conftest import ConfigDef, ProjectDef
-from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.clients.client import KeboolaClient
 from keboola_mcp_server.links import Link
 from keboola_mcp_server.tools.components import create_config
 from keboola_mcp_server.tools.jobs import JobDetail, ListJobsOutput, get_job, list_jobs, run_job

integtests/tools/test_search.py

@@ -4,7 +4,8 @@
 from fastmcp import Context
 
 from integtests.conftest import BucketDef, ConfigDef, TableDef
-from keboola_mcp_server.client import KeboolaClient, SuggestedComponent
+from keboola_mcp_server.clients.ai_service import SuggestedComponent
+from keboola_mcp_server.clients.client import KeboolaClient
 from keboola_mcp_server.tools.search import GlobalSearchOutput, find_component_id, search
 
 LOG = logging.getLogger(__name__)

integtests/tools/test_storage.py

@@ -5,7 +5,7 @@
 from fastmcp import Context
 
 from integtests.conftest import BucketDef, TableDef
-from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.clients.client import KeboolaClient
 from keboola_mcp_server.config import MetadataField
 from keboola_mcp_server.tools.storage import (
     BucketDetail,

src/keboola_mcp_server/cli.py

@@ -125,10 +125,7 @@ async def lifespan(app: Starlette):
                     async with sse_app.lifespan(app):
                         yield
 
-            app = Starlette(
-                middleware=[Middleware(ForwardSlashMiddleware)],
-                lifespan=lifespan
-            )
+            app = Starlette(middleware=[Middleware(ForwardSlashMiddleware)], lifespan=lifespan)
             app.mount('/mcp', http_app)
             app.mount('/sse', sse_app)  # serves /sse/ and /messages
             custom_routes.add_to_starlette(app)

src/keboola_mcp_server/clients/ai_service.py

@@ -0,0 +1,85 @@
+from typing import Any, Optional, cast
+
+from pydantic import BaseModel, Field
+
+from keboola_mcp_server.clients.base import JsonDict, KeboolaServiceClient, RawKeboolaClient
+
+
+class DocsQuestionResponse(BaseModel):
+    """
+    The AI service response to a request to `/docs/question` endpoint.
+    """
+
+    text: str = Field(description='Text of the answer to a documentation query.')
+    source_urls: list[str] = Field(
+        description='List of URLs to the sources of the answer.',
+        default_factory=list,
+        alias='sourceUrls',
+    )
+
+
+class SuggestedComponent(BaseModel):
+    """The AI service response to a /docs/suggest-component request."""
+
+    component_id: str = Field(description='The component ID.', alias='componentId')
+    score: float = Field(description='Score of the component suggestion.')
+    source: str = Field(description='Source of the component suggestion.')
+
+
+class ComponentSuggestionResponse(BaseModel):
+    """The AI service response to a /suggest/component request."""
+
+    components: list[SuggestedComponent] = Field(description='List of suggested components.', default_factory=list)
+
+
+class AIServiceClient(KeboolaServiceClient):
+    """Async client for Keboola AI Service."""
+
+    @classmethod
+    def create(cls, root_url: str, token: Optional[str], headers: dict[str, Any] | None = None) -> 'AIServiceClient':
+        """
+        Creates an AIServiceClient from a Keboola Storage API token.
+
+        :param root_url: The root URL of the AI service API.
+        :param token: The Keboola Storage API token. If None, the client will not send any authorization header.
+        :param headers: Additional headers for the requests.
+        :return: A new instance of AIServiceClient.
+        """
+        return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token, headers=headers))
+
+    async def get_component_detail(self, component_id: str) -> JsonDict:
+        """
+        Retrieves information about a given component.
+
+        :param component_id: The id of the component.
+        :return: Component details as dictionary.
+        """
+        return cast(JsonDict, await self.get(endpoint=f'docs/components/{component_id}'))
+
+    async def docs_question(self, query: str) -> DocsQuestionResponse:
+        """
+        Answers a question using the Keboola documentation as a source.
+        :param query: The query to answer.
+        :return: Response containing the answer and source URLs.
+        """
+        response = await self.raw_client.post(
+            endpoint='docs/question',
+            data={'query': query},
+            headers={'Accept': 'application/json'},
+        )
+
+        return DocsQuestionResponse.model_validate(response)
+
+    async def suggest_component(self, query: str) -> ComponentSuggestionResponse:
+        """
+        Provides list of component suggestions based on natural language query.
+        :param query: The query to answer.
+        :return: Response containing the list of suggested component IDs, their score and source.
+        """
+        response = await self.raw_client.post(
+            endpoint='suggest/component',
+            data={'prompt': query},
+            headers={'Accept': 'application/json'},
+        )
+
+        return ComponentSuggestionResponse.model_validate(response)