Kai Client

A Python client library for interacting with the Keboola AI Assistant Backend API. This library provides async support, SSE streaming, and comprehensive type safety through Pydantic models.

Features

• Command-line interface for quick interactions without writing code
• Async/await support using httpx
• Server-Sent Events (SSE) streaming for real-time chat responses
• Type-safe models with Pydantic v2
• Comprehensive error handling with custom exception classes
• Session management for chat conversations
• Full API coverage including chat, history, and voting endpoints

Installation

Using uv (recommended)

uv add kai-client

Using pip

pip install kai-client

From source

git clone https://github.com/keboola/kai-client.git
cd kai-client
uv sync

Quick Start

import asyncio
from kai_client import KaiClient

async def main():
    # Production: Auto-discover the kai-assistant URL from your Keboola stack
    client = await KaiClient.from_storage_api(
        storage_api_token="your-keboola-token",
        storage_api_url="https://connection.keboola.com"  # Your stack URL
    )

    async with client:
        # Check server health
        ping = await client.ping()
        print(f"Server time: {ping.timestamp}")

        # Start a new chat
        chat_id = client.new_chat_id()

        # Send a message and stream the response
        async for event in client.send_message(chat_id, "What can you help me with?"):
            if event.type == "text":
                print(event.text, end="", flush=True)
            elif event.type == "tool-call":
                print(f"\n[Calling tool: {event.tool_name}]")
            elif event.type == "finish":
                print(f"\n[Finished: {event.finish_reason}]")

asyncio.run(main())

Command-Line Interface

The package includes a kai CLI for quick interactions without writing code.

Setup

Set your credentials as environment variables:

export STORAGE_API_TOKEN="your-keboola-token"
export STORAGE_API_URL="https://connection.keboola.com"

The CLI also auto-loads a .env.local file from the current directory if present, so you can put your credentials there instead:

# .env.local
STORAGE_API_TOKEN=your-keboola-token
STORAGE_API_URL=https://connection.keboola.com

Commands

Health & Info

kai ping              # Check if the server is alive
kai info              # Show server version, uptime, connected MCP servers
kai --version         # Show CLI version

Chat

# Start an interactive chat session
kai chat

# Send a single message (non-interactive)
kai chat -m "What tables do I have?"

# Continue an existing conversation
kai chat --chat-id <chat-id> -m "Tell me more about that"

# Auto-approve tool calls (skips interactive confirmation prompts)
kai chat --auto-approve -m "Create a bucket called test-bucket"

# Output raw JSON events (useful for scripting and piping)
kai chat --json-output -m "List my tables"

In interactive mode, type your messages and press Enter. Type exit, quit, or press Ctrl+C to end.

Tool approval: When Kai calls a write tool (e.g., update_descriptions, run_job, create_config), the CLI pauses and asks you to approve or deny. Use --auto-approve to skip this prompt.

History & Chat Management

# View recent chat history (default: 10)
kai history

# Show more chats
kai history -n 25

# Output history as JSON
kai history --json-output

# Get full details of a specific chat (including messages)
kai get-chat <chat-id>
kai get-chat <chat-id> --json-output

# Delete a chat (prompts for confirmation)
kai delete-chat <chat-id>

# Delete without confirmation
kai delete-chat <chat-id> -y

Voting

# Vote on a message
kai vote <chat-id> <message-id> up
kai vote <chat-id> <message-id> down

# Get votes for a chat
kai get-votes <chat-id>
kai get-votes <chat-id> --json-output

Global Options

These options apply to all commands:

# Pass credentials directly (instead of env vars)
kai --token "your-token" --url "https://connection.keboola.com" ping

# Use a custom base URL for local development
kai --base-url http://localhost:3000 chat -m "Hello"

Help

kai --help              # General help
kai chat --help         # Command-specific help
kai history --help

Local Development vs Production

# Local development (explicit base_url)
client = KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com",
    base_url="http://localhost:3000"
)

# Production (auto-discovers kai-assistant URL)
client = await KaiClient.from_storage_api(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
)

Usage Examples

Simple Chat (Non-Streaming)

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    # Simple one-shot conversation
    chat_id, response = await client.chat("What is 2 + 2?")
    print(response)

Continuing a Conversation

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    # Create a chat session
    chat_id = client.new_chat_id()

    # First message
    async for event in client.send_message(chat_id, "Hello!"):
        if event.type == "text":
            print(event.text, end="")
    print()

    # Continue the conversation (reuse same chat_id)
    async for event in client.send_message(chat_id, "What did I just say?"):
        if event.type == "text":
            print(event.text, end="")
    print()

Handling Tool Calls

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    chat_id = client.new_chat_id()

    async for event in client.send_message(chat_id, "List my Keboola tables"):
        match event.type:
            case "text":
                print(event.text, end="")
            case "step-start":
                print("\n--- New step ---")
            case "tool-call":
                if event.state == "input-available":
                    print(f"\n[Calling {event.tool_name} with {event.input}]")
                elif event.state == "output-available":
                    print(f"\n[{event.tool_name} returned: {event.output}]")
            case "finish":
                print(f"\n[Done: {event.finish_reason}]")

Chat History

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    # Get recent chats
    history = await client.get_history(limit=20)
    for chat in history.chats:
        print(f"Chat {chat.id}: {chat.title}")

    # Iterate through all history
    async for chat in client.get_all_history():
        print(f"Chat: {chat.title}")

    # Get full chat details with messages
    chat_detail = await client.get_chat(chat_id="some-chat-id")
    for message in chat_detail.messages:
        print(f"{message.role}: {message.parts}")

Voting on Messages

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    # Upvote a helpful response
    await client.upvote(chat_id="chat-uuid", message_id="message-uuid")

    # Or downvote
    await client.downvote(chat_id="chat-uuid", message_id="message-uuid")

    # Get all votes for a chat
    votes = await client.get_votes(chat_id="chat-uuid")

Tool Approval for Write Operations

Some tools (like update_descriptions, run_job, create_config) require explicit approval before execution. The server sends a tool-approval-request event with an approval_id that you use to approve or reject.

from kai_client import KaiClient, ToolApprovalRequestEvent

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    chat_id = client.new_chat_id()
    pending_approval_id = None

    async for event in client.send_message(chat_id, "Create a new bucket"):
        if event.type == "text":
            print(event.text, end="")
        elif event.type == "tool-call":
            if event.state == "input-available":
                print(f"\nTool {event.tool_name} needs approval")
            elif event.state == "output-available":
                print(f"\nTool {event.tool_name} completed")
        elif event.type == "tool-approval-request":
            pending_approval_id = event.approval_id

    # Approve the pending tool
    if pending_approval_id:
        async for event in client.approve_tool(
            chat_id=chat_id,
            approval_id=pending_approval_id,
        ):
            if event.type == "text":
                print(event.text, end="")

    # Or reject it
    # async for event in client.reject_tool(
    #     chat_id=chat_id,
    #     approval_id=pending_approval_id,
    #     reason="Not right now",
    # ):
    #     ...

Using SSE Stream Parser

from kai_client import KaiClient, SSEStreamParser

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    parser = SSEStreamParser()
    chat_id = client.new_chat_id()

    async for event in client.send_message(chat_id, "Hello!"):
        parser.process_event(event)

    # Access accumulated data
    print(f"Full response: {parser.text}")
    print(f"Tool calls: {parser.tool_calls}")
    print(f"Finished: {parser.finished}")

Error Handling

from kai_client import (
    KaiClient,
    KaiError,
    KaiAuthenticationError,
    KaiRateLimitError,
    KaiNotFoundError,
)

async with KaiClient(
    storage_api_token="your-token",
    storage_api_url="https://connection.keboola.com"
) as client:
    try:
        async for event in client.send_message("chat-id", "Hello"):
            print(event)
    except KaiAuthenticationError as e:
        print(f"Authentication failed: {e}")
    except KaiRateLimitError as e:
        print(f"Rate limited, try again later: {e}")
    except KaiNotFoundError as e:
        print(f"Chat not found: {e}")
    except KaiError as e:
        print(f"API error: {e.code} - {e.message}")

API Reference

KaiClient

The main client class for interacting with the Kai API.

Factory Method (Recommended for Production)

client = await KaiClient.from_storage_api(
    storage_api_token: str,      # Keboola Storage API token
    storage_api_url: str,        # Keboola connection URL (e.g., https://connection.keboola.com)
    timeout: float = 300.0,      # Request timeout in seconds
    stream_timeout: float = 600.0  # Streaming timeout in seconds
)

This method auto-discovers the kai-assistant service URL from your Keboola stack.

Constructor (For Local Development)

KaiClient(
    storage_api_token: str,      # Keboola Storage API token
    storage_api_url: str,        # Keboola connection URL
    base_url: str = "http://localhost:3000",  # Kai API base URL
    timeout: float = 300.0,      # Request timeout in seconds
    stream_timeout: float = 600.0  # Streaming timeout in seconds
)

Methods

SSE Event Types

The tool-call event has these states: started, input-available (waiting for approval or auto-executing), output-available (completed).

Exceptions

Development

Setup

# Clone the repository
git clone https://github.com/keboola/kai-client.git
cd kai-client

# Install with dev dependencies
uv sync --dev

# Run tests
uv run pytest

# Run linting
uv run ruff check .

Running Tests

# All tests
uv run pytest

# With coverage
uv run pytest --cov=kai_client

# Specific test file
uv run pytest tests/test_client.py

Claude Code Plugin

This repository includes a Claude Code plugin that teaches Claude how to use the Kai CLI correctly.

Installation

Option 1: Download directly (no clone required)

Download the plugin to your Claude Code plugins directory:

# Create plugins directory if it doesn't exist
mkdir -p ~/.claude/plugins

# Download the plugin using curl
curl -L https://github.com/keboola/kai-client/archive/refs/heads/main.tar.gz | \
  tar -xz --strip-components=2 -C ~/.claude/plugins kai-client-main/plugins/kai-cli

Or using wget:

mkdir -p ~/.claude/plugins
wget -qO- https://github.com/keboola/kai-client/archive/refs/heads/main.tar.gz | \
  tar -xz --strip-components=2 -C ~/.claude/plugins kai-client-main/plugins/kai-cli

Option 2: Clone and link (for development)

# Clone the repository
git clone https://github.com/keboola/kai-client.git
cd kai-client

# Option A: Run Claude Code with the plugin directory
claude --plugin-dir plugins/kai-cli

# Option B: Symlink to your plugins directory for persistent access
ln -s "$(pwd)/plugins/kai-cli" ~/.claude/plugins/kai-cli

Verify Installation

After installation, the plugin should be available in Claude Code. Ask Claude to "use kai" or "help me with kai cli" to trigger the skill.

What the Plugin Provides

The plugin includes a skill that activates when you ask Claude to:

• "use kai" or "run kai command"
• "chat with Keboola AI" or "query Keboola"
• "list tables", "check kai history"
• "interact with Keboola assistant"

It teaches Claude about:

• Environment setup - Setting STORAGE_API_TOKEN and STORAGE_API_URL
• Core commands - ping, info, chat, history, get-chat, delete-chat, vote
• Tool approval - Interactive prompts vs --auto-approve for write operations
• Scripting - Using --json-output for automation

Plugin Structure

plugins/kai-cli/
├── .claude-plugin/
│   └── plugin.json              # Plugin manifest
└── skills/
    └── kai-cli/
        ├── SKILL.md             # Main skill guide
        ├── references/
        │   ├── api-details.md   # Python API documentation
        │   └── sse-events.md    # SSE event types reference
        └── examples/
            ├── basic-chat.sh    # Basic usage examples
            └── workflow-automation.sh

License

MIT License - see LICENSE for details.