kai-client 0.9.0

Python client library for the Keboola AI Assistant Backend API

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"

Basic Commands

# Check server health
kai ping

# Get server info
kai info

# Start an interactive chat
kai chat

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

# View chat history
kai history

# Get details of a specific chat
kai get-chat <chat-id>

# Delete a chat
kai delete-chat <chat-id>

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

Chat Options

# Auto-approve tool calls (for automation)
kai chat --auto-approve -m "Create a bucket called test-bucket"

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

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

Local Development

For local development, specify a custom base URL:

kai --base-url http://localhost:3000 chat -m "Hello"

Help

# General help
kai --help

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

Local Development vs Production

Setting	Local Dev	Production
Base URL	http://localhost:3000	Auto-discovered
Setup	Manual base_url parameter	Use from_storage_api()

# 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 create_config, run_job, create_flow) require explicit approval before execution:

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 = 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":
                # Tool is waiting for approval
                print(f"\nTool {event.tool_name} needs approval")
                pending_approval = event
            elif event.state == "output-available":
                print(f"\nTool {event.tool_name} completed")

    # Approve the pending tool call
    if pending_approval:
        async for event in client.confirm_tool(
            chat_id=chat_id,
            tool_call_id=pending_approval.tool_call_id,
            tool_name=pending_approval.tool_name,
        ):
            if event.type == "text":
                print(event.text, end="")

    # Or deny it
    # async for event in client.deny_tool(chat_id, tool_call_id, tool_name):
    #     ...

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

Method	Description
from_storage_api(...)	[Class method] Create client with auto-discovered URL
ping()	Check server health
info()	Get server information
send_message(chat_id, text, ...)	Send a message and stream response
send_tool_result(chat_id, tool_call_id, ...)	Send tool approval/denial result
confirm_tool(chat_id, tool_call_id, ...)	Approve a pending tool call
deny_tool(chat_id, tool_call_id, ...)	Deny a pending tool call
chat(text, ...)	Simple non-streaming chat
get_chat(chat_id)	Get chat details with messages
get_history(limit, ...)	Get chat history
get_all_history()	Iterate through all history
delete_chat(chat_id)	Delete a chat
vote(chat_id, message_id, type)	Vote on a message
upvote(chat_id, message_id)	Upvote a message
downvote(chat_id, message_id)	Downvote a message
get_votes(chat_id)	Get votes for a chat

SSE Event Types

Event Type	Description	Fields
text	Text content	text, state
step-start	Processing step started	-
tool-call	Tool being called	tool_call_id, tool_name, state, input, output
finish	Stream completed	finish_reason
error	Error occurred	message, code

Exceptions

Exception	Error Code	Description
KaiError	-	Base exception
KaiAuthenticationError	unauthorized:chat	Invalid credentials
KaiForbiddenError	forbidden:chat	Access denied
KaiNotFoundError	not_found:chat	Resource not found
KaiRateLimitError	rate_limit:chat	Rate limit exceeded
KaiBadRequestError	bad_request:api	Invalid request
KaiStreamError	-	SSE stream error
KaiConnectionError	-	Connection failed
KaiTimeoutError	-	Request timed out

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/jordanrburger/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/jordanrburger/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/jordanrburger/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.