src/cli.py

from __future__ import annotations

import asyncio
from typing import AsyncIterator

import httpx
import typer
from colorama import Fore, Style, init


API_URL = "http://localhost:8000"

app = typer.Typer(add_completion=False, help="Interact with the LLM backend API")


async def _get_sessions(user: str, server: str) -> list[str]:
    async with httpx.AsyncClient(base_url=server) as client:
        resp = await client.get(f"/sessions/{user}")
        resp.raise_for_status()
        data = resp.json()
        return data.get("sessions", [])


async def _stream_chat(
    user: str, session: str, prompt: str, server: str
) -> AsyncIterator[str]:
    async with httpx.AsyncClient(base_url=server, timeout=None) as client:
        async with client.stream(
            "POST",
            "/chat/stream",
            json={"user": user, "session": session, "prompt": prompt},
        ) as resp:
            resp.raise_for_status()
            async for line in resp.aiter_lines():
                if line:
                    yield line


async def _chat_loop(user: str, server: str) -> None:
    init(autoreset=True)
    sessions = await _get_sessions(user, server)
    session = "default"
    if sessions:
        typer.echo("Existing sessions:")
        for idx, name in enumerate(sessions, 1):
            typer.echo(f"  {idx}. {name}")
        choice = typer.prompt(
            "Select session number or enter new name", default=str(len(sessions))
        )
        if choice.isdigit() and 1 <= int(choice) <= len(sessions):
            session = sessions[int(choice) - 1]
        else:
            session = choice.strip() or session
    else:
        session = typer.prompt("Session name", default=session)

    typer.echo(
        f"Chatting as {Fore.GREEN}{user}{Style.RESET_ALL} in session '{session}'"
    )

    while True:
        try:
            msg = typer.prompt(f"{Fore.CYAN}You{Style.RESET_ALL}")
        except EOFError:
            break
        if msg.strip().lower() in {"exit", "quit"}:
            break
        async for part in _stream_chat(user, session, msg, server):
            typer.echo(f"{Fore.YELLOW}{part}{Style.RESET_ALL}")


@app.callback(invoke_without_command=True)
def main(
    user: str = typer.Option("default", "--user", "-u"),
    server: str = typer.Option(API_URL, "--server", "-s"),
) -> None:
    """Start an interactive chat session."""

    asyncio.run(_chat_loop(user, server))


if __name__ == "__main__":  # pragma: no cover - manual execution
    app()