scripts/chat_client.py

#!/usr/bin/env python3
"""Bee CLI Chat Client — Talk to Bee AGI via the local server.

Usage:
    python chat_client.py                    # Connect to localhost:8000
    python chat_client.py --host bee.local  # Custom host
"""

import argparse
import json
import sys
import time

import httpx
import websocket


def chat_rest(host: str, domain: str = "general"):
    """REST-based chat (non-streaming)."""
    url = f"http://{host}/v1/chat/completions"
    messages = []

    print(f"Bee AGI Chat (REST) — Domain: {domain}")
    print("Type '/quit' to exit, '/domain <name>' to switch")
    print("-" * 50)

    while True:
        user_input = input("\nYou: ").strip()
        if not user_input:
            continue
        if user_input == "/quit":
            break
        if user_input.startswith("/domain "):
            domain = user_input.split(maxsplit=1)[1]
            print(f"Switched to domain: {domain}")
            continue

        messages.append({"role": "user", "content": user_input})

        payload = {
            "model": "bee",
            "messages": messages,
            "max_tokens": 256,
            "temperature": 0.8,
            "stream": False,
            "domain": domain,
        }

        t0 = time.time()
        try:
            r = httpx.post(url, json=payload, timeout=120)
            r.raise_for_status()
            data = r.json()
            reply = data["choices"][0]["message"]["content"]
            elapsed = (time.time() - t0) * 1000

            print(f"\nBee ({elapsed:.0f}ms): {reply}")
            messages.append({"role": "assistant", "content": reply})

        except Exception as e:
            print(f"Error: {e}")


def chat_ws(host: str, domain: str = "general"):
    """WebSocket streaming chat."""
    ws_url = f"ws://{host}/v1/chat"
    messages = []

    print(f"Bee AGI Chat (WebSocket streaming) — Domain: {domain}")
    print("Type '/quit' to exit, '/domain <name>' to switch")
    print("-" * 50)

    ws = websocket.create_connection(ws_url)

    while True:
        user_input = input("\nYou: ").strip()
        if not user_input:
            continue
        if user_input == "/quit":
            break
        if user_input.startswith("/domain "):
            domain = user_input.split(maxsplit=1)[1]
            print(f"Switched to domain: {domain}")
            continue

        messages.append({"role": "user", "content": user_input})

        ws.send(json.dumps({
            "messages": messages,
            "max_tokens": 256,
            "temperature": 0.8,
            "domain": domain,
        }))

        print("\nBee: ", end="", flush=True)
        full_reply = []

        while True:
            try:
                msg = json.loads(ws.recv())
                if msg["type"] == "token":
                    print(msg["content"], end="", flush=True)
                    full_reply.append(msg["content"])
                elif msg["type"] == "done":
                    print()
                    messages.append({"role": "assistant", "content": "".join(full_reply)})
                    break
            except websocket.WebSocketConnectionClosedException:
                print("\n[Connection closed]")
                return
            except Exception as e:
                print(f"\n[Error: {e}]")
                break

    ws.close()


def main():
    parser = argparse.ArgumentParser(description="Bee CLI Chat Client")
    parser.add_argument("--host", default="localhost:8000", help="Bee server host:port")
    parser.add_argument("--ws", action="store_true", help="Use WebSocket streaming")
    parser.add_argument("--domain", default="general", help="Default domain adapter")
    args = parser.parse_args()

    # Check server health
    try:
        r = httpx.get(f"http://{args.host}/health", timeout=5)
        data = r.json()
        print(f"Bee server: {data}")
    except Exception as e:
        print(f"Cannot connect to Bee server at {args.host}: {e}")
        print("Start the server first: python -m bee.server")
        sys.exit(1)

    if args.ws:
        chat_ws(args.host, args.domain)
    else:
        chat_rest(args.host, args.domain)

    print("Goodbye.")


if __name__ == "__main__":
    main()