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

xiaozhi_sdk/core.py

@@ -0,0 +1,269 @@
+import asyncio
+import json
+import logging
+import os
+import re
+import uuid
+from collections import deque
+from typing import Any, Callable, Deque, Dict, Optional
+
+import websockets
+
+from xiaozhi_sdk.config import INPUT_SERVER_AUDIO_SAMPLE_RATE
+from xiaozhi_sdk.iot import OtaDevice
+from xiaozhi_sdk.mcp import McpTool
+from xiaozhi_sdk.utils import get_wav_info, read_audio_file, setup_opus
+
+setup_opus()
+from xiaozhi_sdk.opus import AudioOpus
+
+logger = logging.getLogger("xiaozhi_sdk")
+
+
+class XiaoZhiWebsocket(McpTool):
+
+    def __init__(
+        self,
+        message_handler_callback: Optional[Callable] = None,
+        url: Optional[str] = None,
+        ota_url: Optional[str] = None,
+        audio_sample_rate: int = 16000,
+        audio_channels: int = 1,
+        wake_word: str = "",
+    ):
+        super().__init__()
+        self.url = url
+        self.ota_url = ota_url
+        self.audio_channels = audio_channels
+        self.audio_opus = AudioOpus(audio_sample_rate, audio_channels)
+        self.wake_word = wake_word
+
+        # 客户端标识
+        self.client_id = str(uuid.uuid4())
+        self.mac_addr: Optional[str] = None
+        self.aec = False
+        self.websocket_token = ""
+
+        # 回调函数
+        self.message_handler_callback = message_handler_callback
+
+        # 连接状态
+        self.hello_received = asyncio.Event()
+        self.session_id = ""
+        self.websocket = None
+        self.message_handler_task: Optional[asyncio.Task] = None
+
+        # 输出音频
+        self.output_audio_queue: Deque[bytes] = deque()
+        self.is_playing: bool = False
+
+        # OTA设备
+        self.ota: Optional[OtaDevice] = None
+        self.iot_task: Optional[asyncio.Task] = None
+        self.wait_device_activated: bool = False
+
+        # mcp工具
+        self.mcp_tool_dict = {}
+
+    async def _send_hello(self, aec: bool) -> None:
+        """发送hello消息"""
+        hello_message = {
+            "type": "hello",
+            "version": 1,
+            "features": {"mcp": True, "aec": aec},
+            "transport": "websocket",
+            "audio_params": {
+                "format": "opus",
+                "sample_rate": 16000,
+                "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) -> None:
+        """开始监听"""
+        listen_message = {"session_id": self.session_id, "type": "listen", "state": "start", "mode": "realtime"}
+        await self.websocket.send(json.dumps(listen_message))
+
+    async def is_activate(self, ota_info):
+        """是否激活"""
+        if ota_info.get("activation"):
+            return False
+
+        return True
+
+    async def _activate_iot_device(self, license_key: str, ota_info: Dict[str, Any]) -> None:
+        """激活IoT设备"""
+        if not self.ota:
+            return
+
+        challenge = ota_info["activation"]["challenge"]
+        await asyncio.sleep(3)
+        self.wait_device_activated = True
+        for _ in range(10):
+            if await self.ota.check_activate(challenge, license_key):
+                self.wait_device_activated = False
+                break
+            await asyncio.sleep(3)
+
+    async def _send_demo_audio(self) -> None:
+        """发送演示音频"""
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        wav_path = os.path.join(current_dir, "../file/audio/greet.wav")
+        framerate, channels = get_wav_info(wav_path)
+        audio_opus = AudioOpus(framerate, channels)
+
+        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)
+        await self.send_silence_audio()
+
+    async def send_wake_word(self, wake_word: str) -> None:
+        """发送唤醒词"""
+        await self.websocket.send(
+            json.dumps({"session_id": self.session_id, "type": "listen", "state": "detect", "text": wake_word})
+        )
+
+    async def send_silence_audio(self, duration_seconds: float = 1.2) -> None:
+        """发送静音音频"""
+        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 _handle_websocket_message(self, message: Any) -> None:
+        """处理接受到的WebSocket消息"""
+
+        # audio data
+        if isinstance(message, bytes):
+            pcm_array = await self.audio_opus.opus_to_pcm(message)
+            self.output_audio_queue.extend(pcm_array)
+            return
+
+        # json message
+        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:
+            if data["type"] == "tts":
+                if data["state"] == "sentence_start":
+                    self.is_playing = True
+                    # self.output_audio_queue.clear()
+                else:
+                    self.is_playing = False
+
+            await self.message_handler_callback(data)
+
+    async def _message_handler(self) -> None:
+        """消息处理器"""
+        try:
+            async for message in self.websocket:
+                await self._handle_websocket_message(message)
+        except websockets.ConnectionClosed:
+            if self.message_handler_callback:
+                await self.message_handler_callback(
+                    {"type": "websocket", "state": "close", "source": "sdk.message_handler"}
+                )
+                logger.debug("[websocket] close")
+
+    async def set_mcp_tool(self, mcp_tool_list) -> None:
+        """设置MCP工具"""
+        for mcp_tool in mcp_tool_list:
+            self.mcp_tool_dict[mcp_tool["name"]] = mcp_tool
+
+    async def connect_websocket(self, websocket_token):
+        """连接websocket"""
+        headers = {
+            "Authorization": "Bearer {}".format(websocket_token),
+            "Protocol-Version": "1",
+            "Device-Id": self.mac_addr,
+            "Client-Id": self.client_id,
+        }
+        try:
+            self.websocket = await websockets.connect(uri=self.url, additional_headers=headers)
+        except websockets.exceptions.InvalidMessage as e:
+            logger.error("[websocket] 连接失败，请检查网络连接或设备状态。当前链接地址: %s, 错误信息：%s", self.url, e)
+            return
+        self.message_handler_task = asyncio.create_task(self._message_handler())
+
+        await self._send_hello(self.aec)
+        await self._start_listen()
+        logger.debug("[websocket] Connection successful")
+        await asyncio.sleep(0.5)
+
+    async def init_connection(
+        self, mac_addr: str, aec: bool = False, serial_number: str = "", license_key: str = ""
+    ) -> None:
+        """初始化连接"""
+        mac_pattern = r"^([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}$"
+        if not re.match(mac_pattern, mac_addr):
+            raise ValueError(f"无效的MAC地址格式: {mac_addr}。正确格式应为 XX:XX:XX:XX:XX:XX")
+
+        self.mac_addr = mac_addr.lower()
+        self.aec = aec
+
+        self.ota = OtaDevice(self.mac_addr, self.client_id, self.ota_url, serial_number)
+        ota_info = await self.ota.activate_device()
+        ws_url = ota_info.get("websocket", {}).get("url")
+        self.url = self.url or ws_url
+
+        if not self.url:
+            logger.warning("[websocket] 未找到websocket链接地址")
+            return
+
+        if "tenclass.net" not in self.url and "xiaozhi.me" not in self.url:
+            logger.warning("[websocket] 检测到非官方服务器，当前链接地址: %s", self.url)
+
+        self.websocket_token = ota_info["websocket"]["token"]
+        await self.connect_websocket(self.websocket_token)
+
+        if not await self.is_activate(ota_info):
+            self.iot_task = asyncio.create_task(self._activate_iot_device(license_key, ota_info))
+            logger.debug("[IOT] 设备未激活")
+
+        if self.wake_word:
+            await self.send_wake_word(self.wake_word)
+
+    async def send_audio(self, pcm: bytes) -> None:
+        """发送音频数据"""
+        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.wait_device_activated:
+                logger.debug("[websocket] Server actively disconnected, reconnecting...")
+                await self.connect_websocket(self.websocket_token)
+            elif self.message_handler_callback:
+                await self.message_handler_callback({"type": "websocket", "state": "close", "source": "sdk.send_audio"})
+                self.websocket = None
+                logger.debug("[websocket] Server actively disconnected")
+
+            await asyncio.sleep(0.5)
+        else:
+            await asyncio.sleep(0.1)
+
+    async def close(self) -> None:
+        """关闭连接"""
+        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.iot_task:
+            self.iot_task.cancel()
+
+        if self.websocket:
+            await self.websocket.close()

xiaozhi_sdk/iot.py

@@ -1,50 +1,84 @@
-import aiohttp
+import hashlib
+import hmac
 import json
+import logging
+from typing import Any, Dict, Optional
+
+import aiohttp
 
+from xiaozhi_sdk import __version__
 from xiaozhi_sdk.config import OTA_URL
 
-USER_AGENT = "XiaoXhi-SDK/1.0"
+# 常量定义
+BOARD_TYPE = "xiaozhi-sdk-box"
+USER_AGENT = "xiaozhi-sdk/{}".format(__version__)
+BOARD_NAME = "xiaozhi-sdk"
+
+logger = logging.getLogger("xiaozhi_sdk")
+
+
+class OtaDevice:
+    """
+    OTA设备管理类
 
+    用于处理设备的激活和挑战验证操作。
 
-class OtaDevice(object):
+    Attributes:
+        ota_url (str): OTA服务器URL
+        mac_addr (str): 设备MAC地址
+        client_id (str): 客户端ID
+        serial_number (str): 设备序列号
+    """
+
+    def __init__(self, mac_addr: str, client_id: str, ota_url: Optional[str] = None, serial_number: str = "") -> None:
+        self.ota_url = ota_url or OTA_URL
+        self.ota_url = self.ota_url.rstrip("/")
 
-    def __init__(self, mac_addr: str, client_id: str, serial_number: str = ""):
         self.mac_addr = mac_addr
         self.client_id = client_id
         self.serial_number = serial_number
 
-    async def activate_device(self):
-        header = {
+    def _get_base_headers(self) -> Dict[str, str]:
+        return {
             "user-agent": USER_AGENT,
             "Device-Id": self.mac_addr,
             "Client-Id": self.client_id,
             "Content-Type": "application/json",
-            "serial-number": self.serial_number,
         }
+
+    async def activate_device(self) -> Dict[str, Any]:
+        headers = self._get_base_headers()
+        headers["serial-number"] = self.serial_number
+
         payload = {
-            "application": {"version": "1.0.0"},
+            "application": {"version": __version__},
             "board": {
-                "type": "xiaozhi-sdk-box",
-                "name": "xiaozhi-sdk-main",
+                "type": BOARD_TYPE,
+                "name": BOARD_NAME,
             },
         }
+
         async with aiohttp.ClientSession() as session:
-            async with session.post(OTA_URL, headers=header, data=json.dumps(payload)) as response:
-                data = await response.json()
-                return data
+            async with session.post(self.ota_url + "/", headers=headers, data=json.dumps(payload)) as response:
+                if response.status != 200:
+                    err_text = await response.text()
+                    raise Exception(err_text)
+                response.raise_for_status()
+                return await response.json()
+
+    async def check_activate(self, challenge: str, license_key: str = "") -> bool:
+        url = f"{self.ota_url}/activate"
+
+        headers = self._get_base_headers()
+
+        hmac_instance = hmac.new(license_key.encode(), challenge.encode(), hashlib.sha256)
+        hmac_result = hmac_instance.hexdigest()
+
+        payload = {"serial_number": self.serial_number, "challenge": challenge, "hmac": hmac_result}
 
-    async def check_activate(self, challenge: str):
-        url = OTA_URL + "/activate"
-        header = {
-            "user-agent": USER_AGENT,
-            "Device-Id": self.mac_addr,
-            "Client-Id": self.client_id,
-            "Content-Type": "application/json",
-        }
-        payload = {
-            "serial_number": self.serial_number,
-            "challenge": challenge,
-        }
         async with aiohttp.ClientSession() as session:
-            async with session.post(url, headers=header, data=json.dumps(payload)) as response:
-                return response.status == 200
+            async with session.post(url, headers=headers, data=json.dumps(payload)) as response:
+                is_ok = response.status == 200
+                if not is_ok:
+                    logger.debug("[IOT] wait for activate device...")
+                return is_ok

xiaozhi_sdk/mcp.py

@@ -1,75 +1,171 @@
+import asyncio
+import copy
 import json
+import logging
+import time
+from typing import Any, Dict
 
+import numpy as np
 import requests
 
-from xiaozhi_sdk.config import VL_URL
-from xiaozhi_sdk.data import mcp_initialize_payload, mcp_tools_payload, mcp_tool_conf
+from xiaozhi_sdk.utils.mcp_tool import _get_random_music_info
+
+logger = logging.getLogger("xiaozhi_sdk")
+
+mcp_initialize_payload: Dict[str, Any] = {
+    "jsonrpc": "2.0",
+    "id": 1,
+    "result": {
+        "protocolVersion": "2024-11-05",
+        "capabilities": {"tools": {}},
+        "serverInfo": {"name": "", "version": "0.0.1"},
+    },
+}
+
+mcp_tools_payload: Dict[str, Any] = {
+    "jsonrpc": "2.0",
+    "id": 2,
+    "result": {"tools": []},
+}
 
 
 class McpTool(object):
 
     def __init__(self):
         self.session_id = ""
-        self.vl_token = ""
+        self.explain_url = ""
+        self.explain_token = ""
         self.websocket = None
-        self.tool_func = {}
+        self.mcp_tool_dict = {}
+        self.is_playing = False
+        self.message_handler_callback = None
 
     def get_mcp_json(self, payload: dict):
         return json.dumps({"session_id": self.session_id, "type": "mcp", "payload": payload})
 
     def _build_response(self, request_id: str, content: str, is_error: bool = False):
-        return self.get_mcp_json({
-            "jsonrpc": "2.0",
-            "id": request_id,
-            "result": {
-                "content": [{"type": "text", "text": content}],
-                "isError": is_error,
-            },
-        })
+        return self.get_mcp_json(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": {
+                    "content": [{"type": "text", "text": content}],
+                    "isError": is_error,
+                },
+            }
+        )
 
     async def analyze_image(self, img_byte: bytes, question: str = "这张图片里有什么？"):
-        headers = {"Authorization": f"Bearer {self.vl_token}"}
+        headers = {"Authorization": f"Bearer {self.explain_token}"}
         files = {"file": ("camera.jpg", img_byte, "image/jpeg")}
         payload = {"question": question}
-
-        response = requests.post(VL_URL, files=files, data=payload, headers=headers)
-        return response.json()
+        init_time = time.time()
+        try:
+            response = requests.post(self.explain_url, files=files, data=payload, headers=headers, timeout=8)
+            res_json = response.json()
+        except Exception as e:
+            logger.error("[MCP] 图片解析 error: %s", e)
+            return "网络异常", True
+        if res_json.get("error"):
+            return res_json, True
+        logger.debug("[MCP] 图片解析耗时：%s", time.time() - init_time)
+        return res_json, False
+
+    async def play_custom_music(self, tool_func, arguments):
+        pcm_array, is_error = await tool_func(arguments)
+        while True:
+            if not self.is_playing:
+                break
+            await asyncio.sleep(0.1)
+        pcm_array = await self.audio_opus.change_sample_rate(np.array(pcm_array))
+        self.output_audio_queue.extend(pcm_array)
 
     async def mcp_tool_call(self, mcp_json: dict):
         tool_name = mcp_json["params"]["name"]
-        tool_func = self.tool_func[tool_name]
-        
+        mcp_tool = self.mcp_tool_dict[tool_name]
+        arguments = mcp_json["params"]["arguments"]
+        try:
+            if tool_name == "play_custom_music":
+                # v1 返回 url
+                music_info = await _get_random_music_info(arguments["id_list"])
+                if not music_info.get("url"):
+                    tool_res, is_error = {"message": "播放失败"}, True
+                else:
+                    tool_res, is_error = {"message": "正在为你播放: {}".format(arguments["music_name"])}, False
+                    data = {
+                        "type": "music",
+                        "state": "start",
+                        "url": music_info["url"],
+                        "text": arguments["music_name"],
+                        "source": "sdk.mcp_music_tool",
+                    }
+                    await self.message_handler_callback(data)
+
+                # v2 音频放到输出
+                # asyncio.create_task(self.play_custom_music(tool_func, arguments))
+
+            elif mcp_tool.get("is_async"):
+                tool_res, is_error = await mcp_tool["tool_func"](arguments)
+            else:
+                tool_res, is_error = mcp_tool["tool_func"](arguments)
+        except Exception as e:
+            logger.error("[MCP] tool_name: %s, error: %s", tool_name, e)
+            return self._build_response(mcp_json["id"], "工具调用失败", True)
+
+        if is_error:
+            logger.error("[MCP] tool_name: %s, error: %s", tool_name, tool_res)
+            return self._build_response(mcp_json["id"], "工具调用失败, {}".format(tool_res), True)
+
         if tool_name == "take_photo":
-            res = await self.analyze_image(tool_func(None), mcp_json["params"]["arguments"]["question"])
-        else:
-            res = tool_func(mcp_json["params"]["arguments"])
-        
-        content = json.dumps(res, ensure_ascii=False)
-        return self._build_response(mcp_json["id"], content)
+            tool_res, is_error = await self.analyze_image(tool_res, mcp_json["params"]["arguments"]["question"])
+
+        content = json.dumps(tool_res, ensure_ascii=False)
+        return self._build_response(mcp_json["id"], content, is_error)
 
     async def mcp(self, data: dict):
         payload = data["payload"]
         method = payload["method"]
 
         if method == "initialize":
-            self.vl_token = payload["params"]["capabilities"]["vision"]["token"]
+            self.explain_url = payload["params"]["capabilities"]["vision"]["url"]
+            # self.explain_url = "http://82.157.143.133:8000/vision/explain"
+            self.explain_token = payload["params"]["capabilities"]["vision"]["token"]
+
             mcp_initialize_payload["id"] = payload["id"]
             await self.websocket.send(self.get_mcp_json(mcp_initialize_payload))
 
+        elif method == "notifications/initialized":
+            # print("\nMCP 工具初始化")
+            pass
+
+        elif method == "notifications/cancelled":
+            logger.error("[MCP] 工具加载失败")
+
         elif method == "tools/list":
             mcp_tools_payload["id"] = payload["id"]
-            for name, func in self.tool_func.items():
-                if func:
-                    mcp_tool_conf[name]["name"] = name
-                    mcp_tools_payload["result"]["tools"].append(mcp_tool_conf[name])
+            tool_name_list = []
+            mcp_tool_dict = copy.deepcopy(self.mcp_tool_dict)
+            for _, mcp_tool in mcp_tool_dict.items():
+                tool_name_list.append(mcp_tool["name"])
+                tool_func = mcp_tool.pop("tool_func", None)
+                if not tool_func:
+                    logger.error("[MCP] Tool %s has no tool_func", mcp_tool["name"])
+                    return
+                mcp_tool.pop("is_async", None)
+                mcp_tools_payload["result"]["tools"].append(mcp_tool)
 
             await self.websocket.send(self.get_mcp_json(mcp_tools_payload))
+            logger.debug("[MCP] 加载成功，当前可用工具列表为：%s", tool_name_list)
 
         elif method == "tools/call":
-            print("tools/call", payload)
             tool_name = payload["params"]["name"]
-            if not self.tool_func.get(tool_name):
-                raise Exception("Tool not found")
+
+            if not self.mcp_tool_dict.get(tool_name):
+                logger.warning("[MCP] Tool not found: %s", tool_name)
+                return
 
             mcp_res = await self.mcp_tool_call(payload)
             await self.websocket.send(mcp_res)
+            logger.debug("[MCP] Tool %s called", tool_name)
+        else:
+            logger.warning("[MCP] unknown method %s: %s", method, payload)

xiaozhi_sdk/opus.py

@@ -1,13 +1,11 @@
-import os
-from xiaozhi_sdk import INPUT_SERVER_AUDIO_SAMPLE_RATE
+import math
 
-# 设置 opus 库路径
-os.environ["DYLD_LIBRARY_PATH"] = "/opt/homebrew/lib:" + os.environ.get("DYLD_LIBRARY_PATH", "")
-os.environ["LIBRARY_PATH"] = "/opt/homebrew/lib:" + os.environ.get("LIBRARY_PATH", "")
 import av
 import numpy as np
 import opuslib
 
+from xiaozhi_sdk.config import INPUT_SERVER_AUDIO_SAMPLE_RATE
+
 
 class AudioOpus:
 
@@ -33,11 +31,16 @@ class AudioOpus:
         pcm_bytes = pcm_array.tobytes()
         return self.opus_encoder.encode(pcm_bytes, 960)
 
-    async def change_sample_rate(self, pcm_array):
+    @staticmethod
+    def to_n_960(samples) -> np.ndarray:
+        n = math.ceil(samples.shape[0] / 960)
+        arr_padded = np.pad(samples, (0, 960 * n - samples.shape[0]), mode="constant", constant_values=0)
+        return arr_padded.reshape(n, 960)
+
+    async def change_sample_rate(self, pcm_array) -> np.ndarray:
         if self.sample_rate == INPUT_SERVER_AUDIO_SAMPLE_RATE:
-            return pcm_array.reshape(1, 960)
+            return self.to_n_960(pcm_array)
 
-        c = int(self.sample_rate / INPUT_SERVER_AUDIO_SAMPLE_RATE)
         frame = av.AudioFrame.from_ndarray(np.array(pcm_array).reshape(1, -1), format="s16", layout="mono")
         frame.sample_rate = INPUT_SERVER_AUDIO_SAMPLE_RATE  # Assuming input is 16kHz
         resampled_frames = self.resampler.resample(frame)
@@ -49,10 +52,9 @@ class AudioOpus:
         )
         new_frame.sample_rate = self.sample_rate
         new_samples = new_frame.to_ndarray().flatten()
-        arr_padded = np.pad(new_samples, (0, 960 * c - new_samples.shape[0]), mode="constant", constant_values=0)
-        return arr_padded.reshape(c, 960)
+        return self.to_n_960(new_samples)
 
-    async def opus_to_pcm(self, opus):
+    async def opus_to_pcm(self, opus) -> np.ndarray:
         pcm_data = self.opus_decoder.decode(opus, 960)
         pcm_array = np.frombuffer(pcm_data, dtype=np.int16)
         samples = await self.change_sample_rate(pcm_array)