Source code for klea_utils.ui.tui.repl

#!/usr/bin/env python3
"""
Shared async REPL for Klea chat interfaces.

File: klea_utils/ui/tui/repl.py

Copyright 2026 Ankur Sinha
Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>
"""

import logging




async def run_repl(
    url: str,
    title: str,
    single_query: str = "",
    app_prefix: str = "klea",
    app_name: str = "klea-tui",
) -> None:
    """Run an interactive or single-query chat REPL.

    Connects to the API server's SSE ``/query/stream`` endpoint and
    displays progress labels and the final answer.

    In interactive mode each query is sent as the user types it;
    the REPL does not batch queries.

    :param url: Base URL of the API server (e.g. ``http://127.0.0.1:8005``)
    :param title: Application title displayed on start
    :param single_query: If set, run one query and exit instead of REPL loop
    :param app_prefix: Prefix for user/assistant labels (e.g. ``"klea"``)
    :param app_name: Log identity for this frontend, used as the log file
        name so each process keeps its own logs (e.g. ``"klea-rag-tui"``)
    """
    # Configure process-wide logging for this client process.  Lazy:
    # platformdirs / plogging imports are cheap, and everything else is
    # deferred below so --help stays fast.
    from platformdirs import PlatformDirs

    from klea_utils.plogging import setup_root_logger

    setup_root_logger(
        app_name,
        stderr_level=logging.INFO,
        log_dir=PlatformDirs(app_name).user_data_dir,
    )

    # Lazy: avoids importing yaspin (and its deps) at module level
    import coolname
    from yaspin import yaspin

    from klea_utils.api.sse import stream_events
    from klea_utils.api.utils import check_api_is_ready

    chat_id = coolname.generate_slug(2)

    with yaspin(text="Waiting for API..."):
        await check_api_is_ready(f"{url}/health/ready")

    async def _query_one(query: str) -> None:
        full_response = ""
        error_msg = ""
        print()

        with yaspin(text="Working ...", timer=True) as spinner:
            async for event in stream_events(query, chat_id, url):
                if event["type"] == "progress":
                    spinner.text = event["node"]
                elif event["type"] == "complete":
                    full_response = event.get("message_for_user", "")
                    spinner.ok("[OK]")
                elif event["type"] == "error":
                    error_msg = event.get("message", "Unknown server error")
                    spinner.fail("[ERROR]")
                    break

        output = error_msg or full_response
        label = "(ERROR)" if error_msg else "(AI)"
        print(f"{app_prefix} {label} >>> {output}")
        print("\n" + "-" * 40 + "\n")

    if single_query:
        print(f"{app_prefix} (USER) >>> {single_query}")
        await _query_one(single_query)
        return

    print(f"*** {title} ({url}) ***")
    print("Please note that answers are generated by LLMs and may be incorrect.")
    print()
    print("Type 'quit' to exit.")
    print("\n" + "-" * 40 + "\n")

    while True:
        q = input(f"{app_prefix} (USER) >>> ")
        if q.lower() == "quit":
            break
        await _query_one(q)

    print(f"\n{app_prefix} >>> Bye!")