File size: 4,229 Bytes
db82745 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | #!/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()
|