xiaozhi-sdk 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

xiaozhi_sdk/__init__.py

@@ -1,155 +1,3 @@
-import asyncio
-import json
-import os
-import uuid
-import wave
-import websockets
-from collections import deque
-from typing import Dict
+__version__ = "0.2.0"
 
-from xiaozhi_sdk.config import INPUT_SERVER_AUDIO_SAMPLE_RATE, WSS_URL
-from xiaozhi_sdk.iot import OtaDevice
-from xiaozhi_sdk.mcp import McpTool
-from xiaozhi_sdk.opus import AudioOpus
-from xiaozhi_sdk.utils import read_audio_file, get_wav_info
-
-
-class XiaoZhiWebsocket(McpTool):
-
-    def __init__(self, message_handler_callback=None, url=None, audio_sample_rate=16000, audio_channels=1):
-        super().__init__()
-        self.url = url or WSS_URL
-        self.audio_sample_rate = audio_sample_rate
-        self.audio_channels = audio_channels
-        self.audio_opus = AudioOpus(audio_sample_rate, audio_channels)
-        self.client_id = str(uuid.uuid4())
-        self.mac_addr = None
-        self.message_handler_callback = message_handler_callback
-
-        self.hello_received = asyncio.Event()
-        self.session_id = ""
-        self.audio_queue = deque()
-        self.websocket = None
-        self.message_handler_task = None
-
-    async def send_hello(self, aec: bool):
-        hello_message = {
-            "type": "hello",
-            "version": 1,
-            "features": {"aec": aec, "mcp": True},
-            "transport": "websocket",
-            "audio_params": {
-                "format": "opus",
-                "sample_rate": INPUT_SERVER_AUDIO_SAMPLE_RATE,
-                "channels": 1,
-                "frame_duration": 60,
-            },
-        }
-        await self.websocket.send(json.dumps(hello_message))
-        await asyncio.wait_for(self.hello_received.wait(), timeout=10.0)
-
-    async def start_listen(self):
-        listen_message = {
-            "session_id": self.session_id,
-            "type": "listen",
-            "state": "start",
-            "mode": "realtime"
-        }
-        await self.websocket.send(json.dumps(listen_message))
-
-    async def set_mcp_tool_callback(self, tool_func: Dict[str, callable]):
-        self.tool_func = tool_func
-
-    async def activate_iot_device(self):
-        ota = OtaDevice(self.mac_addr, self.client_id)
-
-        data = await ota.activate_device()
-        if data.get("activation"):
-            await self.send_demo_get_code_audio()
-            challenge = data["activation"]["challenge"]
-            await asyncio.sleep(3)
-            for _ in range(10):
-                if await ota.check_activate(challenge):
-                    break
-                await asyncio.sleep(3)
-
-    async def init_connection(self, mac_addr: str, aec: bool = False):
-        self.mac_addr = mac_addr
-        headers = {
-            "Authorization": "Bearer test-token",
-            "Protocol-Version": "1",
-            "Device-Id": mac_addr,
-            "Client-Id": self.client_id,
-        }
-
-        self.websocket = await websockets.connect(uri=self.url, additional_headers=headers)
-        self.message_handler_task = asyncio.create_task(self.message_handler())
-        await self.send_hello(aec)
-        await self.start_listen()
-        asyncio.create_task(self.activate_iot_device())
-
-    async def send_demo_get_code_audio(self):
-        current_dir = os.path.dirname(os.path.abspath(__file__))
-        wav_path = os.path.join(current_dir, "../file/greet.wav")
-        framerate, nchannels = get_wav_info(wav_path)
-        audio_opus = AudioOpus(framerate, nchannels)
-
-        for pcm_data in read_audio_file(wav_path):
-            opus_data = await audio_opus.pcm_to_opus(pcm_data)
-            await self.websocket.send(opus_data)
-
-    async def send_silence_audio(self, duration_seconds: float = 1.2):
-        # 发送 静音数据
-        frames_count = int(duration_seconds * 1000 / 60)
-        pcm_frame = b"\x00\x00" * int(INPUT_SERVER_AUDIO_SAMPLE_RATE / 1000 * 60)
-
-        for _ in range(frames_count):
-            await self.send_audio(pcm_frame)
-
-    async def send_audio(self, pcm: bytes):
-        if not self.websocket:
-            return
-
-        state = self.websocket.state
-        if state == websockets.protocol.State.OPEN:
-            opus_data = await self.audio_opus.pcm_to_opus(pcm)
-            await self.websocket.send(opus_data)
-        elif state in [websockets.protocol.State.CLOSED, websockets.protocol.State.CLOSING]:
-            if self.message_handler_callback:
-                await self.message_handler_callback({"type": "websocket", "state": "close", "source": "sdk.send_audio"})
-            await asyncio.sleep(0.5)
-        else:
-            await asyncio.sleep(0.1)
-
-    async def message_handler(self):
-        try:
-            async for message in self.websocket:
-                if isinstance(message, bytes):
-                    pcm_array = await self.audio_opus.opus_to_pcm(message)
-                    self.audio_queue.extend(pcm_array)
-                else:
-                    data = json.loads(message)
-                    message_type = data["type"]
-
-                    if message_type == "hello":
-                        self.hello_received.set()
-                        self.session_id = data["session_id"]
-                    elif message_type == "mcp":
-                        await self.mcp(data)
-                    elif self.message_handler_callback:
-                        await self.message_handler_callback(data)
-        except websockets.ConnectionClosed:
-            if self.message_handler_callback:
-                await self.message_handler_callback(
-                    {"type": "websocket", "state": "close", "source": "sdk.message_handler"})
-
-    async def close(self):
-        if self.message_handler_task and not self.message_handler_task.done():
-            self.message_handler_task.cancel()
-            try:
-                await self.message_handler_task
-            except asyncio.CancelledError:
-                pass
-
-        if self.websocket:
-            await self.websocket.close()
+from xiaozhi_sdk.core import XiaoZhiWebsocket  # noqa

xiaozhi_sdk/__main__.py

@@ -1,90 +1,11 @@
-import argparse
-import asyncio
-import re
-import time
-from collections import deque
+import logging
 
-import numpy as np
-import sounddevice as sd
-
-from xiaozhi_sdk import XiaoZhiWebsocket
-
-input_audio = deque()
-
-is_play_audio = False
-
-
-async def message_handler_callback(message):
-    print("message received:", message)
-
-
-async def assistant_audio_play(audio_queue):
-    global is_play_audio
-    # 创建一个持续播放的流
-    stream = sd.OutputStream(samplerate=16000, channels=1, dtype=np.int16)
-    stream.start()
-    last_time = None
-
-    while True:
-
-        if not audio_queue:
-            await asyncio.sleep(0.01)
-            if last_time and time.time() - last_time > 1:
-                is_play_audio = False
-            continue
-
-        is_play_audio = True
-        pcm_data = audio_queue.popleft()
-        stream.write(pcm_data)
-        last_time = time.time()
-
-
-class Client:
-    def __init__(self, mac_address):
-        self.mac_address = mac_address
-        self.xiaozhi = None
-        pass
-
-    async def start(self):
-        self.xiaozhi = XiaoZhiWebsocket(message_handler_callback)
-        await self.xiaozhi.init_connection(self.mac_address, aec=False)
-        asyncio.create_task(assistant_audio_play(self.xiaozhi.audio_queue))
-
-    def callback_func(self, indata, frames, time, status):
-        pcm = (indata.flatten() * 32767).astype(np.int16).tobytes()
-        input_audio.append(pcm)
-
-    async def process_audio(self):
-        while True:
-            if not input_audio:
-                await asyncio.sleep(0.02)
-                continue
-            pcm = input_audio.popleft()
-            if not is_play_audio:
-                await self.xiaozhi.send_audio(pcm)
-
-
-def mac_address(string):
-    """验证是否为有效的MAC地址"""
-    if re.fullmatch(r'([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}', string):
-        return string
-    else:
-        raise argparse.ArgumentTypeError(f"无效的MAC地址格式: '{string}'")
-
-
-async def main():
-    parser = argparse.ArgumentParser(description="这是一个小智SDK。")
-    parser.add_argument('device', type=mac_address, help='你的小智设备的MAC地址 (格式: XX:XX:XX:XX:XX:XX)')
-
-    args = parser.parse_args()
-    client = Client(args.device)
-    await client.start()
-    await asyncio.sleep(2)
-
-    with sd.InputStream(callback=client.callback_func, channels=1, samplerate=16000, blocksize=960):
-        print("Recording... Press Ctrl+C to stop.")
-        await client.process_audio()  # 持续处理音频
+from xiaozhi_sdk.cli import main
 
+logger = logging.getLogger("xiaozhi_sdk")
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    try:
+        main()
+    except KeyboardInterrupt:
+        logger.debug("Stopping...")

xiaozhi_sdk/cli.py

@@ -0,0 +1,231 @@
+import asyncio
+import io
+import logging
+import time
+from collections import deque
+from typing import Optional
+
+import click
+import colorlog
+import numpy as np
+import sounddevice as sd
+from PIL import ImageGrab
+
+from xiaozhi_sdk import XiaoZhiWebsocket
+from xiaozhi_sdk.config import INPUT_SERVER_AUDIO_SAMPLE_RATE
+
+# 定义自定义日志级别
+INFO1 = 21
+INFO2 = 22
+INFO3 = 23
+
+# 添加自定义日志级别到logging模块
+logging.addLevelName(INFO1, "INFO1")
+logging.addLevelName(INFO2, "INFO2")
+logging.addLevelName(INFO3, "INFO3")
+
+
+# 为logger添加自定义方法
+def info1(self, message, *args, **kwargs):
+    if self.isEnabledFor(INFO1):
+        self._log(INFO1, message, args, **kwargs)
+
+
+def info2(self, message, *args, **kwargs):
+    if self.isEnabledFor(INFO2):
+        self._log(INFO2, message, args, **kwargs)
+
+
+def info3(self, message, *args, **kwargs):
+    if self.isEnabledFor(INFO3):
+        self._log(INFO3, message, args, **kwargs)
+
+
+# 将自定义方法添加到Logger类
+logging.Logger.info1 = info1
+logging.Logger.info2 = info2
+logging.Logger.info3 = info3
+
+# 配置彩色logging
+handler = colorlog.StreamHandler()
+handler.setFormatter(
+    colorlog.ColoredFormatter(
+        "%(log_color)s%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        log_colors={
+            "DEBUG": "white",
+            "INFO": "white",
+            "INFO1": "green",
+            "INFO2": "cyan",
+            "INFO3": "blue",
+            "WARNING": "yellow",
+            "ERROR": "red",
+            "CRITICAL": "red,bg_white",
+        },
+    )
+)
+
+logger = logging.getLogger("xiaozhi_sdk")
+logger.addHandler(handler)
+logger.setLevel(logging.DEBUG)
+
+# 全局状态
+input_audio_buffer: deque[bytes] = deque()
+is_playing_audio = False
+is_end = False
+
+
+def get_image_byte(data):
+    logger.info("请截图需要识别的内容:")
+    time.sleep(1)
+    for _ in range(20):
+        im = ImageGrab.grabclipboard()
+        if not im:
+            time.sleep(0.3)
+            continue
+        if im.mode == "RGBA":
+            im = im.convert("RGB")
+
+        byte_io = io.BytesIO()
+        im.save(byte_io, format="JPEG", quality=30)
+        # im.save("./test.jpg", format='JPEG', quality=30)
+
+        img_bytes = byte_io.getvalue()
+        logger.info("截图成功")
+
+        # if platform.system() == "Darwin":
+        #     subprocess.run("pbcopy", input=b"")
+
+        return img_bytes, False
+
+    logger.error("截图失败, 请在10秒内完成截图")
+    return "截图失败", True
+
+
+async def handle_message(message):
+    """处理接收到的消息"""
+    global is_end
+    if message["type"] == "stt":  # 人类语音
+        logger.info1("human: %s", message["text"])
+    elif message["type"] == "tts" and message["state"] == "sentence_start":  # AI语音
+        logger.info2("AI: %s", message["text"])
+    elif message["type"] == "tts" and message["state"] == "stop":
+        pass
+        # logger.info2("播放结束")
+        # logger.info("聆听中...")
+    elif message["type"] == "llm":  # 表情
+        logger.info3("emotion: %s", message["text"])
+    else:  # 其他消息
+        pass
+        # logger.info("other: %s", message)
+
+    if message["type"] == "websocket" and message["state"] == "close":
+        is_end = True
+
+
+async def play_assistant_audio(audio_queue: deque[bytes], enable_audio):
+    """播放音频流"""
+    global is_playing_audio
+
+    stream = None
+    if enable_audio:
+        stream = sd.OutputStream(samplerate=INPUT_SERVER_AUDIO_SAMPLE_RATE, channels=1, dtype=np.int16)
+        stream.start()
+    last_audio_time = None
+
+    while True:
+        if is_end:
+            return
+
+        if not audio_queue:
+            await asyncio.sleep(0.01)
+            if last_audio_time and time.time() - last_audio_time > 1:
+                is_playing_audio = False
+            continue
+
+        is_playing_audio = True
+        pcm_data = audio_queue.popleft()
+        if stream:
+            stream.write(pcm_data)
+        last_audio_time = time.time()
+
+
+class XiaoZhiClient:
+    """小智客户端类"""
+
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        ota_url: Optional[str] = None,
+        wake_word: str = "",
+    ):
+        self.xiaozhi: Optional[XiaoZhiWebsocket] = None
+        self.url = url
+        self.ota_url = ota_url
+        self.mac_address = ""
+        self.wake_word = wake_word
+
+    async def start(self, mac_address: str, serial_number: str, license_key: str, enable_audio):
+        """启动客户端连接"""
+        self.mac_address = mac_address
+        self.xiaozhi = XiaoZhiWebsocket(handle_message, url=self.url, ota_url=self.ota_url, wake_word=self.wake_word)
+        from xiaozhi_sdk.utils.mcp_tool import take_photo
+
+        take_photo["tool_func"] = get_image_byte
+
+        await self.xiaozhi.set_mcp_tool([take_photo])
+        await self.xiaozhi.init_connection(
+            self.mac_address, aec=False, serial_number=serial_number, license_key=license_key
+        )
+
+        asyncio.create_task(play_assistant_audio(self.xiaozhi.output_audio_queue, enable_audio))
+
+    def audio_callback(self, indata, frames, time, status):
+        """音频输入回调函数"""
+        pcm_data = (indata.flatten() * 32767).astype(np.int16).tobytes()
+        input_audio_buffer.append(pcm_data)
+
+    async def process_audio_input(self):
+        """处理音频输入"""
+        while True:
+
+            if is_end:
+                return
+
+            if not input_audio_buffer:
+                await asyncio.sleep(0.02)
+                continue
+
+            pcm_data = input_audio_buffer.popleft()
+            if not is_playing_audio:
+                await self.xiaozhi.send_audio(pcm_data)
+
+
+async def run_client(
+    mac_address: str, url: str, ota_url: str, serial_number: str, license_key: str, enable_audio: bool, wake_word: str
+):
+    """运行客户端的异步函数"""
+    logger.debug("Recording... Press Ctrl+C to stop.")
+    client = XiaoZhiClient(url, ota_url, wake_word)
+    await client.start(mac_address, serial_number, license_key, enable_audio)
+
+    with sd.InputStream(callback=client.audio_callback, channels=1, samplerate=16000, blocksize=960):
+        await client.process_audio_input()
+
+
+@click.command()
+@click.argument("mac_address")
+@click.option("--url", help="服务端websocket地址")
+@click.option("--ota_url", help="OTA地址")
+@click.option("--serial_number", default="", help="设备的序列号")
+@click.option("--license_key", default="", help="设备的授权密钥")
+@click.option("--enable_audio", default=True, help="是否开启音频播放")
+@click.option("--wake_word", default="", help="唤醒词")
+def main(
+    mac_address: str, url: str, ota_url: str, serial_number: str, license_key: str, enable_audio: bool, wake_word: str
+):
+    """小智SDK客户端
+
+    MAC_ADDRESS: 设备的MAC地址 (格式: XX:XX:XX:XX:XX:XX)
+    """
+    asyncio.run(run_client(mac_address, url, ota_url, serial_number, license_key, enable_audio, wake_word))

xiaozhi_sdk/config.py

@@ -1,5 +1,3 @@
 INPUT_SERVER_AUDIO_SAMPLE_RATE = 16000
 
-WSS_URL = "wss://api.tenclass.net/xiaozhi/v1/"
-OTA_URL = "https://api.tenclass.net/xiaozhi/ota/"
-VL_URL = "http://api.xiaozhi.me/mcp/vision/explain"
+OTA_URL = "https://api.tenclass.net/xiaozhi/ota"