client_demo.py

#!/usr/bin/env python3
"""
Simple CLI client for testing the Kaanta Tax Assistant API.

Example:
    python client_demo.py --question "Compute PAYE for 1500000 income" \
        --input gross_income=1500000
"""

from __future__ import annotations

import argparse
import json
from typing import Dict, Optional

import httpx


def _parse_inputs(raw_pairs: Optional[list[str]]) -> Optional[Dict[str, float]]:
    if not raw_pairs:
        return None

    parsed: Dict[str, float] = {}
    for item in raw_pairs:
        if "=" not in item:
            raise argparse.ArgumentTypeError(
                f"Calculator input '{item}' must be in key=value form."
            )
        key, value = item.split("=", 1)
        key = key.strip()
        if not key:
            raise argparse.ArgumentTypeError("Input keys cannot be empty.")
        try:
            parsed[key] = float(value)
        except ValueError as exc:
            raise argparse.ArgumentTypeError(
                f"Value for '{key}' must be numeric."
            ) from exc
    return parsed


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Send a test question to a running Kaanta Tax Assistant API."
    )
    parser.add_argument(
        "--base-url",
        default="https://eniiyanu-kaanta.hf.space",
        help="Base URL of the service (default: %(default)s).",
    )
    parser.add_argument(
        "--question",
        required=True,
        help="User question or task to send to the assistant.",
    )
    parser.add_argument(
        "--as-of",
        help="Optional YYYY-MM-DD date context for tax calculations.",
    )
    parser.add_argument(
        "--tax-type",
        default="PIT",
        help="Tax type for calculator runs (PIT, CIT, VAT).",
    )
    parser.add_argument(
        "--jurisdiction",
        default="state",
        help="Jurisdiction filter used by the rules engine.",
    )
    parser.add_argument(
        "--input",
        dest="inputs",
        action="append",
        metavar="key=value",
        help="Calculator input (repeatable). Example: --input gross_income=1500000",
    )
    parser.add_argument(
        "--rule-id",
        dest="rule_ids",
        action="append",
        help="Optional whitelist of rule IDs to evaluate (repeat flag for multiple).",
    )
    parser.add_argument(
        "--no-rag-quotes",
        action="store_true",
        help="Skip RAG enrichment when running the calculator.",
    )
    parser.add_argument(
        "--hf-token",
        help="Optional Hugging Face access token when querying a private Space.",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=60.0,
        help="HTTP timeout in seconds (default: %(default)s).",
    )
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    try:
        inputs = _parse_inputs(args.inputs)
    except argparse.ArgumentTypeError as exc:
        parser.error(str(exc))
        return

    payload = {
        "question": args.question,
        "as_of": args.as_of,
        "tax_type": args.tax_type.upper() if args.tax_type else None,
        "jurisdiction": args.jurisdiction,
        "inputs": inputs,
        "with_rag_quotes_on_calc": not args.no_rag_quotes,
        "rule_ids_whitelist": args.rule_ids,
    }

    # Remove fields that FastAPI would reject when left as None.
    payload = {k: v for k, v in payload.items() if v is not None}

    url = args.base_url.rstrip("/") + "/v1/query"
    headers = {}
    if args.hf_token:
        headers["Authorization"] = f"Bearer {args.hf_token}"

    def do_request(target: str) -> httpx.Response:
        return httpx.post(target, json=payload, headers=headers, timeout=args.timeout)

    tried_urls = [url]
    try:
        response = do_request(url)
        if response.status_code == 404 and "/proxy" not in url:
            proxy_url = args.base_url.rstrip("/") + "/proxy/v1/query"
            response = do_request(proxy_url)
            tried_urls.append(proxy_url)
        response.raise_for_status()
    except httpx.TimeoutException:
        parser.exit(1, f"Request timed out after {args.timeout} seconds\n")
    except httpx.HTTPStatusError as exc:
        locations = " -> ".join(tried_urls)
        parser.exit(
            1,
            f"Server returned HTTP {exc.response.status_code} for {locations}:\n"
            f"{exc.response.text}\n",
        )
    except httpx.RequestError as exc:
        parser.exit(1, f"Request failed: {exc}\n")

    print(json.dumps(response.json(), indent=2))


if __name__ == "__main__":
    main()